gm.text.Greedy

class gemma.gm.text.Greedy

基类: gemma.gm.text._sampling.SamplingMethod

贪婪采样。

get_next_tokens(

logits: jaxtyping.Float[Array, '*B V'] | jaxtyping.Float[ndarray, '*B V'],

rng: jaxtyping.UInt32[Array, '2'] | jaxtyping.UInt32[ndarray, '2'] | jax.Array,

) → jaxtyping.Int[Array, '*B'] | jaxtyping.Int[ndarray, '*B']

返回要生成的下一个 token。

参数：

• logits – Logits，由模型返回（即 softmax 之前）。

• rng – 一个随机密钥。

返回：

要生成的下一个 token。