Peeking under the covers to see next token probabilities

[glaforge]

Hi all,

I've read this interesting article that shows how to use Keras and the Gemma model to peek under the hood and seen the next token probabilities.

I'm curious if 1) it's possible to do something like that when running Gemma with Jlama, and 2) if yes, how can you drive Jlama to do that?

Guillaume

[tjake]

Hi,

Yeah you can get the raw probabilities in the sample method.

Jlama/jlama-core/src/main/java/com/github/tjake/jlama/model/AbstractModel.java

Line 213 in 8d3b5ba

public int sample(AbstractTensor output, float temperature, float uniformSample, AbstractTensor logits) {

Currently it does argmax to pick the top value based on the temperature but it should be simple to make this method pluggable (PRs welcome!)