Keras 3 API 文档 / KerasNLP / 采样器 / TopKSampler

TopKSampler

[来源]

TopKSampler

keras_nlp.samplers.TopKSampler(k=5, seed=None, **kwargs)

Top-K 采样器类。

此采样器实现了 top-k 搜索算法。简而言之,top-k 算法从概率最高的 K 个 token 中随机选择一个 token,其选择机会由概率决定。

参数

  • k: int,top-k 的 k 值。
  • seed: int。随机种子。默认为 None

调用参数

{{call_args}}

示例

causal_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en")

# Pass by name to compile.
causal_lm.compile(sampler="top_k")
causal_lm.generate(["Keras is a"])

# Pass by object to compile.
sampler = keras_nlp.samplers.TopKSampler(k=5, temperature=0.7)
causal_lm.compile(sampler=sampler)
causal_lm.generate(["Keras is a"])