文本生成系统

相关源文件

• llama/generation.py
• llama/tokenizer.py

文本生成系统是 Llama 2 代码库的核心组件，负责使用预训练的 Transformer 模型生成文本。本文档解释了系统如何处理提示、管理分词以及为通用文本补全和基于聊天的应用生成连贯的文本响应。有关 Transformer 架构实现的信息，请参阅模型架构。

系统概览

文本生成系统提供了一个使用 Llama 2 模型的高级接口，处理模型加载、分词处理以及使用各种采样参数生成文本的复杂性。它充当用户提示和模型输出之间的桥梁。

来源: llama/generation.py51-395 llama/tokenizer.py14-68

核心组件

Llama 类

Llama 类作为文本生成的主要接口，封装了一个预训练的 Transformer 模型和分词器。它提供了文本补全和聊天补全的方法。

来源: llama/generation.py51 llama/generation.py125-127 llama/tokenizer.py14-68

构建 Llama 实例

build 静态方法使用指定的模型权重和分词器初始化 Llama 实例

1. 初始化分布式处理和模型并行
2. 设置 CUDA 设备
3. 加载模型检查点
4. 初始化 Transformer 模型和分词器

来源: llama/generation.py52-123

文本生成过程

Generate 方法

generate 方法是处理给定分词提示的文本生成的核心函数

来源: llama/generation.py129-231

生成过程的关键组件是：

1. 分词准备：将输入提示转换为词元张量并准备生成
2. 增量生成：使用模型的正向传播逐个生成词元
3. 采样：应用温度和 top-p 采样来控制随机性
4. EOS 检测：检测序列结束以适当地终止生成

Top-p (核心) 采样

系统使用 top-p（核心）采样来控制文本生成中的随机性。

来源: llama/generation.py398-421

公共 API

文本续写

text_completion 方法为通用文本提示生成补全

1. 对输入提示进行分词
2. 调用 generate 方法
3. 将生成的词元解码回文本
4. 以字典列表的形式返回补全

Input: "The best way to learn is to"
Output: "The best way to learn is to teach others what you've learned. When you explain a concept to someone else, you solidify your own understanding and identify gaps in your knowledge."

来源: llama/generation.py233-282

聊天补全

chat_completion 方法处理基于对话的交互

1. 使用特殊词元（例如 [INST]、[/INST]、<）格式化对话>)
2. 验证对话结构（用户/助手消息交替）
3. 分词并生成响应
4. 包含安全检查以防止提示注入

来源: llama/generation.py284-395

分词

Tokenizer 类提供了文本和词元 ID 之间转换的方法

方法	目的	关键参数
编码 (encode)	将文本转换为 token ID	s (输入字符串), bos (添加起始词元), eos (添加结束词元)
解码 (decode)	将词元 ID 转换为文本	t (词元 ID 列表)

特殊词元包括

• 序列起始 (BOS)：当 bos=True 时添加到序列的开头
• 序列结束 (EOS)：当 eos=True 时添加到序列的末尾
• 填充 (PAD)：用于批量处理

来源: llama/tokenizer.py14-68

安全措施

系统实施安全措施以防止提示注入

1. 特殊标签检测：检查用户输入中是否包含特殊标签，如 [INST]、[/INST]、<<SYS>>、<</SYS>>
2. 错误处理：对不安全的请求返回预定义的错误消息
3. 对话结构验证：确保对话遵循预期的用户/助手交替模式

SPECIAL_TAGS = ["[INST]", "[/INST]", "<<SYS>>", "<</SYS>>"]
UNSAFE_ERROR = "Error: special tags are not allowed as part of the prompt."

来源: llama/generation.py43-48 llama/generation.py321-322 llama/generation.py374-395

配置参数

文本生成系统可以使用各种参数进行配置，以控制生成过程。

参数	描述	默认值
温度 (temperature)	控制随机性（值越高随机性越大）	0.6
top_p	核心采样概率阈值	0.9
max_gen_len	生成序列的最大长度	None（使用模型的 max_seq_len - 1）
logprobs	是否返回词元对数概率	否
echo	是否在输出中包含提示	False（仅限文本补全）

来源: llama/generation.py129-231 llama/generation.py233-282 llama/generation.py284-395