gm.text.Greedy

class gemma.gm.text.Greedy

Bases: gemma.gm.text._sampling.SamplingMethod

Greedy sampling.

get_next_tokens(

logits: kauldron.ktyping.array_type_meta.Float['*B V'],

rng: kauldron.ktyping.array_type_meta.UInt32['2'] | kauldron.ktyping.array_type_meta.Fry[''] | kauldron.ktyping.array_type_meta.KdPRNGKey,

) → kauldron.ktyping.array_type_meta.Int['*B']

Returns the next tokens to generate.

Parameters:

• logits – Logits, as returned by the model (i.e. before softmax).

• rng – A random key.

Returns:

The next tokens to generate.