生成

torchtune.generation.generate(model: TransformerDecoder, prompt: Tensor, *, max_generated_tokens: int, pad_id: int = 0, temperature: float = 1.0, top_k: Optional[int] = None, stop_tokens: Optional[List[int]] = None, rng: Optional[Generator] = None, custom_generate_next_token: Optional[Callable] = None) → Tuple[Tensor, Tensor]

根据提示生成模型的令牌，并返回生成的logits。

Parameters:

• model (TransformerDecoder) – 用于生成的模型

• 提示 (torch.Tensor) – 包含与给定提示相关联的标记ID的张量，形状为[seq_length]或[bsz x seq_length]。

• max_generated_tokens (int) – 要生成的令牌数量

• pad_id (int) – 用于填充的token ID，默认为0。

• 温度 (float) – 用于缩放预测的logits的值，默认值为1.0。

• top_k (可选[int]) – 如果指定，我们将采样修剪为仅包含在top_k概率内的token id， 默认为None。

• stop_tokens (Optional[List[int]]) – 如果指定，当生成任何这些标记时，生成将停止，默认为 None。

• rng (可选[torch.Generator]) – 随机数生成器，默认为 None。

• custom_generate_next_token (可选[Callable]) – 如果指定，我们将使用 custom_generate_next_token 函数。这通常只在您希望出于性能原因指定 torch.compile 版本的生成下一个令牌时有用。如果为 None，我们使用默认的 generate_next_token()。 默认值为 None。

注意

此函数仅在使用仅解码器模型时进行了测试。

示例

>>> model = torchtune.models.llama3.llama3_8b()
>>> tokenizer = torchtune.models.llama3.llama3_tokenizer()
>>> prompt = tokenizer.encode("Hi my name is")
>>> rng.manual_seed(42)
>>> output, logits = generate(model, torch.tensor(prompt), max_generated_tokens=100, pad_id=0)
>>> print(tokenizer.decode(output[0].tolist()))
Hi my name is Jeremy and I'm a friendly language model assistant!

Returns:

两个张量的元组：

• tokens (torch.Tensor): 生成的tokens的张量，

  形状为 [bsz x seq_len + num_generated_tokens]，其中 num_generated_tokens 可能小于 max_generated_tokens，如果提供了 stop_tokens。

• logits (torch.Tensor): 与生成的tokens相关的logits张量，

  形状为 [bsz x seq_len + num_generated_tokens x vocab_size]。

Return type:

元组[torch.Tensor, torch.Tensor]