generate

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]

从模型中生成以提示为条件的 tokens，并返回生成的 logits。

参数:

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

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

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

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

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

• top_k (Optional[int]) – 如果指定，我们将采样限制为仅在 top_k 概率内的 token ID，默认为 None。

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

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

• custom_generate_next_token (Optional[Callable]) – 如果指定，我们将使用 custom_generate_next_token function。这通常仅在您想要为性能原因指定 torch.compile 版本的 generate next token 时才有用。如果为 None，我们使用默认的 generate_next_token()。默认为 None。

注意

此函数仅在 decoder-only 模型上进行了测试。

示例

>>> 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!

返回值:

包含两个张量的元组

• 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]。

返回类型:

Tuple[torch.Tensor, torch.Tensor]