DataWhale Task02：从零预练习一个tiny-llama 20923

DataWhale Task02：从零预练习一个tiny-llama 20923

原文link：https://github.com/KMnO4-zx/tiny-llm
开源内容：https://github.com/datawhalechina/tiny-universe
   写在前面！！：因为没有编译环境，这期主要讲代码的理解
  起首，安装python脚本，安装依赖库:
pip install -r requirements.txt
练习步调：

• 练习Tokenizer： python train_vocab.py --download True --vocab_size 4096
  
• 数据预处置处罚：python preprocess.py
  
• 练习模型：python train.py
  
• 使用模型天生文本：python sample.py --prompt "One day, Lily met a Shoggoth"
  

具体步调：

1.练习Tokenizer：

在自然语言处置处罚（NLP）中，Tokenizer（分词器）是将输入文本分割成更小单元（通常称为token）的工具或方法。Token可以是词、子词、字符、乃至是标点符号，具体取决于所使用的分词方法。Tokenizer 的主要作用是将原始文本转换成恰当模型处置处罚的情势。
以下是常见的两种分词方法：

• 词级别分词（Word-level tokenization）：直接按词进行分割，通常基于空格或标点符号。比方：
  
  
  
  • 输入：“I love programming.”
    
  • 输出：["I", "love", "programming", "."]
    
  
  
• 子词级别分词（Subword-level tokenization）：将单词进一步拆解成更小的子单元，常见于一些基于子词的模型（如BPE或WordPiece）。这对处置处罚未登录词（out-of-vocabulary words）尤其有效。
  
  
  
  • 输入：“programming”
    
  • 输出：["pro", "gram", "ming"]
    
  
  

Tokenizer 通常是语言模型的前处置处罚步调，用于将自然语言文本转换为可以输入模型的数值表现。
下载数据集并练习:
python train_vocab.py --download True --vocab_size 4096
tokenizer.py

1. def download_file(url: str, fname: str, chunk_size=1024):
2.     """发送HTTP GET请求以流式方式获取文件"""
3.     ···
5. def download():
6.     """执行 download_file 下载数据集"""
7.     ···
9. def train_vocab(vocab_size: int=32000, num_shards: int=20):
10.     """
11.     vocab_size: int, 词汇表的大小，决定分词器的词汇量。
12.     num_shards: int, 用于加快词汇表训练的效率，指定要处理的分片数量。
13.     """
14.     # 确保词汇表大小为正数
15.     assert vocab_size > 0, "Vocab size must be positive"
17.     # SentencePiece 模型的前缀路径，将用于保存分词器
18.     prefix = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}")
20.     # 1) 将多个分片中的文本导出为单个文本文件 tiny.txt
21.     tiny_file = os.path.join(DATA_CACHE_DIR, "tiny.txt")
22.     data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
23.     shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
25.     # 创建 tiny.txt 文件并写入指定数量的分片中的文本
26.     print(f"Writing temporary file {tiny_file} with {num_shards} shards...")
27.     with open(tiny_file, "w", encoding="utf-8") as of:
28.         # 遍历前 num_shards 个分片
29.         for shard in tqdm(shard_filenames[:num_shards]):
30.             with open(shard, "r") as f:
31.                 data = json.load(f)  # 读取分片中的JSON数据
32.             # 遍历每个例子，将其中的故事文本写入 tiny.txt 文件
33.             for example in data:
34.                 text = example["story"]
35.                 text = text.strip()  # 去除文本首尾的空白字符
36.                 of.write(text + "\n")  # 每个文本写入一行
38.     # 输出生成的 tiny.txt 文件的大小
39.     print(f"Size is: {os.path.getsize(tiny_file) / 1024 / 1024:.2f} MB")
41.     # 2) 使用 SentencePiece 训练分词器
42.     print("Will now train the vocab...")
43.     spm.SentencePieceTrainer.train(
44.         input=tiny_file,         # 输入文件为之前生成的 tiny.txt
45.         model_prefix=prefix,     # 模型前缀路径
46.         model_type="bpe",        # 使用 Byte-Pair Encoding (BPE) 训练分词器
47.         vocab_size=vocab_size,   # 词汇表大小
48.         self_test_sample_size=0, # 自测样本大小设置为 0
49.         input_format="text",     # 输入文件格式为纯文本
50.         character_coverage=1.0,  # 覆盖所有字符（包括非常见字符）
51.         num_threads=os.cpu_count(),  # 使用 CPU 的线程数
52.         split_digits=True,       # 拆分数字
53.         allow_whitespace_only_pieces=True,  # 允许仅由空格组成的词元
54.         byte_fallback=True,      # 启用字节级回退
55.         unk_surface=r" \342\201\207 ",  # UNK token 表示未知字符的方式
56.         normalization_rule_name="identity"  # 使用“identity”归一化规则
57.     )
59.     # 3) 可选的清理操作，询问用户是否删除临时文件 tiny.txt
60.     dec = input(f"Delete the temporary file {tiny_file}? [y/N] ")
61.     if dec.lower() == "y":
62.         os.remove(tiny_file)  # 删除临时文件
63.         print(f"Deleted {tiny_file}")
65.     # 输出模型保存的路径
66.     print(f"Trained tokenizer is in {prefix}.model")
67.     print("Done.")

复制代码

代码理解：
主要使用了 SentencePiece 库来天生基于 BPE（Byte-Pair Encoding）的方法。练习过程由多个步调构成，涉及从数据集中读取文本、天生临时文件、练习词汇表并最终生存模型。
train_vocab概述：这个函数的目标是：

• 从多个 JSON 文件中提取文本数据。
  
• 将这些数据写入一个临时文件（tiny.txt）。
  
• 使用 SentencePiece 对提取的文本进行分词器的练习。
  
• 最终根据练习设置天生一个词汇表巨细为 vocab_size 的分词器模型。
  

vocab_size: 定义练习出的分词器的词汇表巨细（即模型中包含多少个词元）。
num_shards: 指定从数据集中读取的分片数，用于加速处置处罚多个数据分片
代码详解：
assert vocab_size > 0, "Vocab size must be positive"
确保词汇表巨细为整数，用来确保 vocab_size 是大于 0 的有效值，如果不满意条件则抛出错误
prefix = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}")
设置模型的路径前缀，将词汇表练习的模型路径设置为 DATA_CACHE_DIR 目录下的 tok{vocab_size}。练习后的模型将生存为 {prefix}.model 文件。
tiny_file = os.path.join(DATA_CACHE_DIR, "tiny.txt") data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data") shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
将数据集的多个 JSON 文件归并成一个 tiny.txt 文件，

• data_dir 是数据集所在的目录，代码会读取该目录下的所有 JSON 文件。
  
• shard_filenames 是读取到的 JSON 文件路径列表，按文件名排序。
  

1. with open(tiny_file, "w", encoding="utf-8") as of:
2.     for shard in tqdm(shard_filenames[:num_shards]):
3.         with open(shard, "r") as f:
4.             data = json.load(f)
5.         for example in data:
6.             text = example["story"]
7.             text = text.strip()
8.             of.write(text + "\n")

复制代码

打开 tiny.txt 作为写入的目标文件。
遍历前 num_shards 个 JSON 文件，每个文件包含许多 JSON 格式的“故事”。
对每个 JSON 文件中的每个“故事”进行处置处罚，将此中的 story 字段文本提取出来并写入 tiny.txt。
print(f"Size is: {os.path.getsize(tiny_file) / 1024 / 1024:.2f} MB")
输入天生文件（tiny.txt）的巨细
spm.SentencePieceTrainer.train( input=tiny_file, model_prefix=prefix, model_type="bpe", vocab_size=vocab_size, ... )
使用 SentencePiece 练习分词器
input=tiny_file: 使用之前天生的 tiny.txt 文件作为练习输入。
model_prefix=prefix: 模型文件将以 prefix 为前缀生存。
model_type="bpe": 采用 BPE（Byte-Pair Encoding）算法练习分词器。
vocab_size: 指定词汇表的巨细。
其他参数用于控制分词器练习时的一些细节，比方字符覆盖率、线程数、是否拆分数字等。
可选的清理步调，练习完成后，询问用户是否删除临时文件 tiny.txt，如果用户输入 y 则删除该文件。
print(f"Trained tokenizer is in {prefix}.model")
输出模型生存路径，模型文件会生存为 {prefix}.model，这行代码输出模型生存的路径。
展示怎样使用该类来处置处罚 TinyStory 数据集中的故事文本

1. class Tokenizer:
2.     def __init__(self, tokenizer_model=None):
3.         """
4.         初始化分词器。加载预训练的SentencePiece模型，并设置一些特殊的token ID。
5.    class Tokenizer:
6.     def __init__(self, tokenizer_model=None):
7.         """
8.         初始化分词器。加载预训练的SentencePiece模型，并设置一些特殊的token ID。
10.         参数:
11.         tokenizer_model: str, 可选，分词器模型的路径，如果不指定则使用默认路径 TOKENIZER_MODEL。
12.         """
13.         # 如果提供了分词器模型路径，使用该路径；否则使用默认模型路径
14.         model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL
15.         # 确保模型文件存在
16.         assert os.path.isfile(model_path), model_path
18.         # 加载 SentencePiece 模型
19.         self.sp_model = SentencePieceProcessor(model_file=model_path)
20.         self.model_path = model_path
22.         # 获取分词器的特殊token和词汇表大小
23.         self.n_words: int = self.sp_model.vocab_size()  # 词汇表大小
24.         self.bos_id: int = self.sp_model.bos_id()       # 句子开头 (BOS) 的ID
25.         self.eos_id: int = self.sp_model.eos_id()       # 句子结尾 (EOS) 的ID
26.         self.pad_id: int = self.sp_model.pad_id()       # 填充 (PAD) 的ID
28.         # 验证分词器词汇表大小是否正确
29.         assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
33.     def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
34.         """
35.         将字符串编码为词元ID列表。可以选择是否添加句子开头 (BOS) 和句子结尾 (EOS) 标记。
37.         参数:
38.         s: str, 要编码的字符串。
39.         bos: bool, 是否在编码的词元列表前添加 BOS 标记。
40.         eos: bool, 是否在编码的词元列表末尾添加 EOS 标记。
42.         返回:
43.         List[int]: 编码后的词元ID列表。
44.         """
45.         # 确保输入是字符串类型
46.         assert type(s) is str
47.         # 使用SentencePiece将字符串编码为词元ID
48.         t = self.sp_model.encode(s)
49.         # 如果需要BOS标记，将其添加到词元列表开头
50.         if bos:
51.             t = [self.bos_id] + t
52.         # 如果需要EOS标记，将其添加到词元列表末尾
53.         if eos:
54.             t = t + [self.eos_id]
55.         return t
57.     def decode(self, t: List[int]) -> str:
58.         """
59.         将词元ID列表解码为字符串。
61.         参数:
62.         t: List[int], 词元ID列表。
64.         返回:
65.         str: 解码后的字符串。s
66.         """
67.         return self.sp_model.decode(t)     
68.    
69.    
70.    
71.    

复制代码

此中：
1.__init__方法

tokenizer_model: 这是一个可选参数，用于指定分词器模型的路径。如果未提供，则使用默认路径 TOKENIZER_MODEL
model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL
assert os.path.isfile(model_path), model_path
起首，检查是否提供了自定义的分词器模型路径。如果没有，就使用默认路径 TOKENIZER_MODEL。
使用 assert 语句确保指定路径的模型文件存在。如果文件不存在，抛出错误并输出文件路径。
self.sp_model = SentencePieceProcessor(model_file=model_path)
self.model_path = model_path
起首，检查是否提供了自定义的分词器模型路径。如果没有，就使用默认路径 TOKENIZER_MODEL。
使用 assert 语句确保指定路径的模型文件存在。如果文件不存在，抛出错误并输出文件路径。
self.sp_model = SentencePieceProcessor(model_file=model_path)
self.model_path = model_path
SentencePieceProcessor 是 SentencePiece 的焦点处置处罚器，负责加载模型。model_file=model_path 使得处置处罚器加载指定路径下的分词器模型。
self.n_words: int = self.sp_model.vocab_size()
self.bos_id: int = self.sp_model.bos_id()
self.eos_id: int = self.sp_model.eos_id()
self.pad_id: int = self.sp_model.pad_id()

• n_words: 模型的词汇表巨细（即分词器中包含多少个词元）。
  
• bos_id: 句子开头的特殊标记 BOS 的 ID。
  
• eos_id: 句子末了的特殊标记 EOS 的 ID。
  
• pad_id: 填充标记 PAD 的 ID。
  

1. assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()

复制代码

通过断言确保 vocab_size() 和 get_piece_size() 的结果划一，验证分词器的词汇表巨细是正确的。
2.encode方法

   encode 方法用于将输入字符串编码为词元ID列表，同时可以选择是否在结果中添加 BOS 和 EOS 标记。
  s: 要编码的字符串。
bos: 是否在词元列表开头添加 BOS 标记。
eos: 是否在词元列表末了添加 EOS 标记。
assert type(s) is str
t = self.sp_model.encode(s)
起首，检查 s 是否为字符串类型。
使用 SentencePiece 的 encode 方法将字符串编码为词元ID列表。
if bos:
t = [self.bos_id] + t
if eos:
t = t + [self.eos_id]

• 如果 bos 参数为真，则在词元列表开头插入 BOS 标记。
  
• 如果 eos 参数为真，则在词元列表末了添加 EOS 标记。
  

末了，返回处置处罚后的词元ID列表。
3.decode方法

   decode 方法用于将词元ID列表解码回原始字符串。
  t: 词元ID的列表。
return self.sp_model.decode(t)
直接使用 SentencePiece 的 decode 方法将 ID 列表解码回对应的文本字符串。
2.数据预处置处罚

具体代码的分析，具体代码见github。
将文本数据转换为模型可以大概理解的数字序列
该代码实现了分词数据的预处置处罚、加载和批处置处罚，适用于大规模语言模型的练习任务。它主要分为以下几部分：
1. 分片处置处罚函数 process_shard 和预处置处罚函数 pretokenize

这两个函数负责将文本数据进行分词处置处罚并生存为二进制文件。只管函数体未展示出来，但从函数签名可以推测：

• process_shard 负责处置处罚单个数据分片，将此中的文本数据分词并存储为 .bin 文件。
  
• pretokenize 负责对所有分片进行批量处置处罚，调用 process_shard 对每个分片进行分词预处置处罚。
  

2. PretokDataset 类

该类继承自 torch.utils.data.IterableDataset，用于从磁盘加载预处置处罚好的分词数据并将其提供给 PyTorch 模型练习。
焦点功能：

• 初始化数据集：通过 __init__ 方法，指定数据集的分别方式（练习集或测试集）、最大序列长度、词汇表巨细以及词汇表来源（如 Llama2 或自定义词汇表）。
  
• 加载分片文件：在 __iter__ 方法中：
  
  
  
  • 根据 vocab_source 决定要加载的分片文件路径。
    
  • 对于练习集，加载所有分片；对于测试集，加载第一个分片。
    
  
  
• 批量天生：通过 memmap 方式读取二进制文件中的数据，确保大文件可以从磁盘加载而不完全占用内存。每个分片文件被分割为多个批次，max_seq_len 决定了每个批次的长度。
  
  
  
  • 模型输入 x 是当前批次的前 max_seq_len 个词元。
    
  • 模型输出 y 是对应下一个词元，用于构建语言模型的监视学习任务。
    
  
  

其他功能：

• 随机性：通过 worker_info 和 rank 实现数据的并行加载和分布式练习。在数据加载过程中，基于 worker_id 和 rank 创建随机数种子，确保差异进程、线程之间的数据处置处罚是唯一的，不会重复。
  
• 数据天生：__iter__ 方法天生 (x, y) 对，供 PyTorch 模型练习。
  

3. Task 类

Task 类封装了数据集的批处置处罚流程，并为外部调用提供了一个静态方法 iter_batches，用于迭代天生批次数据。
焦点功能：

• 批量加载数据：使用 torch.utils.data.DataLoader 创建批次迭代器，将 PretokDataset 的输出打包为固定巨细的批次。
  
• 数据迁移到装备：在每个批次中，将数据移动到指定的计算装备（如 GPU），并使用 non_blocking=True 以加速数据拷贝的速度。
  

其他功能：

• 支持并行加载：可以通过 num_workers 参数指定数据加载器的并行线程数，以加速数据加载过程。
  

3: 练习模型

这个 generate 方法实现了一个基于语言模型的文本天生过程，徐徐天生新 token 并将其附加到现有序列 idx 中。该方法可以通过多次前向流传和采样来天生新序列，使用了温度和 top_k 战略来控制天生过程。
关键流程解释：

• 初始输入 idx：
  
  
  
  • idx 是外形为 (bz, seq_len) 的长整型张量，表现输入序列。
    
  • bz 是 batch size，seq_len 是序列的长度。
    
  
  
• 序列长度控制：
  
  
  
  • 当输入序列长度超过模型的最大序列长度 self.args.max_seq_len 时，序列会被截断，只保存末了的 max_seq_len 个 token（即上下文）。
    
  • 这包管了模型的输入不会超过它的上下文窗口巨细。
    
  
  
• 前向流传：
  
  
  
  • 对于当前输入序列 idx_cond，模型进行前向流传，输出猜测 logits。logits[:, -1, :] 表现我们只取末了一个 token 的输出，因为这是当前天生 token 的猜测概率分布。
    
  
  
• 天生新 token：
  
  
  
  • 确定性采样（当 temperature == 0.0 时）：直接选择概率最高的下一个 token（torch.topk(logits, k=1)）。
    
  • 随机采样（当 temperature > 0 时）：先将 logits 按 temperature 缩放，然后进行采样。temperature 值越高，天生的文本越随机；值越低，天生的文本越确定。
    
    
    
    • top_k 采样：如果指定了 top_k，只从 top k 个最大概的 token 中进行采样，这可以减少天生的随机性，避免模型天生一些不合理的 token。
      
    • Softmax 和 Multinomial：F.softmax 将 logits 转换为概率分布，torch.multinomial 根据该分布采样出下一个 token。
      
    
    
  
  
• 更新序列：
  
  
  
  • 将新天生的 token（idx_next）附加到现有序列 idx 背面，并继承迭代直到天生所需的 token 数量。
    
  
  
• 返回天生的序列：
  
  
  
  • 最终返回更新后的 idx，它包含初始输入序列和新天生的 token。
    
  
  

参数：

• idx: 输入的 token 序列，外形为 (batch_size, seq_len)。
  
• max_new_tokens: 要天生的 token 数量。
  
• temperature: 控制采样随机性的参数，值越低天生越确定，越高天生越随机。
  
• top_k: 限制从 top k 个 token 中进行采样，值较低时可提高天生质量。
  

代码示例的逻辑流：

• 每次天生时，输入序列 idx 截取为最大答应的长度。
  
• 通过模型获取序列末了一个位置的 logits。
  
• 根据温度和 top_k 参数选择下一个 token。
  
• 将新天生的 token 加入到序列，重复天生直到到达 max_new_tokens。
  

这个方法主要用于简单的文本天生任务，恰当在推理模式下运行。
4: 使用模型天生文本

在模型练习完成后，会在output目录下天生一个ckpt.pt文件，这个文件就是我们练习好的模型。我们可以使用以下下令天生文本。
python sample.py --prompt "One day, Lily met a Shoggoth"
sample.py代码详解：

1.     class TextGenerator:
2.     def __init__(self,
3.                  checkpoint='output/ckpt.pt',  # 模型检查点路径
4.                  tokenizer_model_path='tok4096.model',  # 分词器模型路径
5.                  seed=1337,  # 随机种子，确保可重复性
6.                  device=None,  # 设备，优先使用 CUDA，如果没有可用的 CUDA，则使用 CPU
7.                  dtype="float32"):  # 数据类型，默认为 float32，可以选择 float16 或 bfloat16
8.         """
9.         初始化 TextGenerator 类，加载模型、设置设备和分词器等。
10.         """
12.         #模型加载配置
14. ​        self.checkpoint = checkpoint  # 保存的模型检查点路径
15. ​        self.tokenizer_model_path = tokenizer_model_path  # 分词器模型文件路径
16. ​        self.seed = seed  # 随机数种子，用于生成的可重复性
17. ​        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')  # 根据硬件条件选择设备
18. ​        self.dtype = dtype  # 模型的浮点数类型
19. ​        self.device_type = 'cuda' if 'cuda' in self.device else 'cpu'  # 判断当前设备是否为 CUDA
20.     # 设置随机种子，确保生成的可重复性
21.     torch.manual_seed(seed)  # 设置 CPU 随机种子
22.     torch.cuda.manual_seed(seed)  # 设置 CUDA 随机种子
23.     torch.backends.cuda.matmul.allow_tf32 = True  # 允许 CUDA 使用 TF32 精度进行矩阵乘法运算
24.     torch.backends.cudnn.allow_tf32 = True  # 允许 cuDNN 使用 TF32 精度加速
25.    
26.     # 根据 dtype 选择适当的自动混合精度上下文
27.     ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype]
28.     self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(device_type=self.device_type, dtype=ptdtype)
29.    
30.     # 加载模型检查点文件
31.     checkpoint_dict = torch.load(self.checkpoint, map_location=self.device)  # 加载模型参数
32.     gptconf = ModelArgs(**checkpoint_dict['model_args'])  # 初始化模型参数
33.     self.model = Transformer(gptconf)  # 实例化 Transformer 模型
34.     state_dict = checkpoint_dict['model']  # 获取模型状态字典
35.    
36.     # 去除状态字典中的不必要前缀
37.     unwanted_prefix = '_orig_mod.'  # 这个前缀在保存时可能被添加，现在要去除它
38.     for k, v in list(state_dict.items()):
39.         if k.startswith(unwanted_prefix):
40.             state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)  # 去除不必要的前缀
41.    
42.     # 加载模型参数到模型中
43.     self.model.load_state_dict(state_dict, strict=False)
44.     # 计算模型参数量
45.     num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
46.     print(f"Model has {num_params} parameters.")
47.     # 设置模型为评估模式（evaluation mode），防止训练模式下的 dropout 等操作影响结果
48.     self.model.eval()
49.     # 将模型放置到正确的设备上（GPU 或 CPU）
50.     self.model.to(self.device)
51.     # 初始化分词器
52.     self.tokenizer = Tokenizer(tokenizer_model=self.tokenizer_model_path)  # 根据指定的路径加载分词器
54. def sample(self,
55.            start="Hello!",  # 生成文本的起始提示词，可以是任意字符串
56.            num_samples=3,  # 生成样本的数量，默认生成 3 个样本
57.            max_new_tokens=256,  # 每个样本生成的最大 token 数，默认最多生成 256 个 token
58.            temperature=1.0,  # 控制生成的随机性，1.0 为标准，值越大越随机
59.            top_k=300):  # 保留概率最高的 top_k 个 token，限制生成时的选择范围
60.     """
61.     根据给定的起始文本生成样本。
62.    
63.     :param start: 生成文本的起始提示词
64.     :param num_samples: 要生成的文本样本数
65.     :param max_new_tokens: 每个样本生成的最大 token 数
66.     :param temperature: 控制生成的随机性，值越小生成越确定，值越大生成越随机
67.     :param top_k: 限制生成时选择的 token 范围
68.     :return: 生成的文本样本列表
69.     """
70.     # 如果 start 是以 'FILE:' 开头，表示从文件中读取起始文本
71.     if start.startswith('FILE:'):
72.         with open(start[5:], 'r', encoding='utf-8') as f:
73.             start = f.read()  # 读取文件内容作为起始文本
74.    
75.     # 将起始文本编码为 token id 序列
76.     start_ids = self.tokenizer.encode(start, bos=True, eos=False)  # bos=True 表示加上句首标记，eos=False 表示不加句尾标记
77.     x = (torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...])  # 将编码后的 token id 转为 PyTorch 张量
78.    
79.     generated_texts = []  # 用于保存生成的文本样本
80.     with torch.no_grad():  # 禁用梯度计算，提升效率
81.         with self.ctx:  # 进入自动混合精度的上下文（如果是 GPU 并使用 float16 时）
82.             for k in range(num_samples):  # 循环生成指定数量的样本
83.                 y = self.model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)  # 生成文本
84.                 generated_texts.append(self.tokenizer.decode(y[0].tolist()))  # 解码生成的 token 序列为可读文本
85.    
86.     return generated_texts  # 返回生成的文本样本

复制代码

实现了一个 TextGenerator 类，用于加载预练习语言模型，并根据输入提示天生文本样本。下面是对各个部分的具体解释：
1. 初始化 (__init__ 方法)

__init__ 方法用于加载模型、设置装备（如 CPU 或 CUDA）、配置分词器和初始化随机种子，以确保天生的文本可重复。
参数：

• checkpoint: 模型的检查点文件路径，包含模型的权重和配置。
  
• tokenizer_model_path: 分词器模型的路径，负责将文本转化为 token 序列。
  
• seed: 随机种子，确保每次天生的文本可重复。
  
• device: 选择使用的计算装备，优先选择 CUDA，如果没有可用的 CUDA，则使用 CPU。
  
• dtype: 数据类型，默认为 float32，也可以选择 float16 或 bfloat16 来加速计算。
  

主要步调：

• 设置随机种子：通过 torch.manual_seed 和 torch.cuda.manual_seed 确保 CPU 和 CUDA 装备上的随机性是可控的。
  
• 装备配置：如果可用，则使用 CUDA，否则使用 CPU；设置主动混淆精度（torch.amp.autocast）来提高天生效率。
  
• 加载模型检查点：从指定的 checkpoint 路径加载预练习模型，去除状态字典中的不必要前缀（如 _orig_mod.），然后将模型权重加载到模型中。
  
• 评估模式：设置模型为评估模式（self.model.eval()），禁用掉练习过程中使用的 dropout 等机制，以确保天生结果的稳定性。
  
• 分词器初始化：加载分词器，用于将文本转化为 token，并将天生的 token 序列解码回可读文本。
  

2. 天生文本 (sample 方法)

sample 方法负责根据给定的起始文本天生指定数量的样本，返回的是天生的文本列表。
参数：

• start: 天生文本的起始提示词，可以是恣意字符串，也可以是文件路径（以 FILE: 开头）。
  
• num_samples: 要天生的文本样本数量，默认天生 3 个样本。
  
• max_new_tokens: 每个样本天生的最大 token 数量，默认值为 256。
  
• temperature: 控制天生文本的随机性，值越高越随机；值越低天生的文本越接近模型的确定性猜测。
  
• top_k: 限制天生时的选择范围，只保存概率最高的 top k 个 token，如允许以避免天生低概率但大概不合理的 token。
  

主要步调：

• 文件输入支持：如果 start 参数以 FILE: 开头，表现从指定的文件中读取起始文本作为天生提示词。
  
• 分词器编码：将起始文本编码为 token id 序列，添加句首标记（bos=True），天生的序列作为模型的初始输入。
  
• 天生样本：在禁用梯度计算（torch.no_grad()）的上下文中，使用模型的 generate 方法天生指定数量的文本样本。
  
• 解码：将天生的 token 序列通过分词器解码成可读的文本，并生存到列表中。
  

3. 模型天生逻辑

在 sample 方法中，调用 self.model.generate 实行天生逻辑。该方法是之前定义的，用于根据给定的输入序列徐徐天生新 token，并支持多种采样战略（如 temperature 和 top_k 采样）。天生的 token 序列通过分词器解码回文本后返回。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。