gm.text.TopPSampling

class gemma.gm.text.TopPSampling(*, p: float = 0.9, temperature: float = 1.0)

Bases: gemma.gm.text._sampling.SamplingMethod

Top-p (Nucleus) Sampling.

p: float = 0.9

temperature: float = 1.0

get_next_tokens(

logits: kauldron.ktyping.array_type_meta.Float['... 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['...']

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.