NLP基础:文本预处理、词元化、词表构建与BPE分词器
本文介绍NLP中的文本预处理基础步骤,包括数据集读取、文本清洗、词元化、词表构建,并结合nanochat项目说明BPE分词器训练、对话渲染和预训练数据加载流程。
- 人工智能
- 数据科学
NLP基础:文本预处理、词元化、词表构建与BPE分词器
学习语言模型时,很容易把“文本预处理”理解成一段无聊的清洗代码:读文件、正则替换、split、统计词频、转成数字。这样理解没有错,但不够完整。
对GPT这类模型来说,文本预处理不是训练前的杂活,而是在回答一个更根本的问题:
如何把人类写出来的开放文本,稳定、可逆、高效地变成模型可以学习的整数序列?
神经网络不能直接吃字符串。它看到的是张量,是数字。于是我们必须先设计一套规则,把文本映射成token id,再通过Embedding层把token id变成连续向量。这个映射规则看似只是工程实现,实际上会影响模型的上下文长度、训练效率、泛化能力、未知词处理、多语言支持,甚至会影响模型在代码、数字、中文、emoji等文本上的表现。
这篇文章分成两层:
- 前半部分用D2L的“时光机器”数据集讲清楚文本预处理的基本概念:清洗、词元化、词表、索引映射。
- 后半部分结合nanochat源码,看现代GPT风格Tokenizer为什么会使用byte-level BPE、特殊token、对话渲染和固定长度数据打包。
为什么需要文本预处理
原始文本有几个特点:
- 它是字符串,不是模型可以直接计算的数字。
- 它长度不固定,有的句子很短,有的文档很长。
- 它包含大小写、标点、空格、换行、数字、代码、中文、英文、emoji等复杂符号。
- 它的词是开放集合,新词、拼写变体、专有名词永远会出现。
而语言模型需要的是一种稳定输入:
文本 -> token序列 -> token id序列 -> 张量
比如:
"the time machine"
-> ["the", "time", "machine"]
-> [1, 19, 50]
这一步要同时满足几个目标。
第一,要能把任意文本转成数字。否则模型根本无法训练。
第二,词表必须有限。模型的Embedding矩阵大小是:
vocab_size x n_embd
如果词表无限大,Embedding层也就无限大,训练不可能进行。
第三,序列不能太长。同样一句话,如果按字符切分,token数量会很多;如果按常见子词切分,token数量会少很多。上下文窗口是有限的,token越省,模型在一次forward里能看到的内容越多。
第四,信息最好不要丢。对GPT类生成模型来说,Tokenizer最好满足:
decode(encode(text)) == text
也就是说,文本转成token id之后,再解码回来应该和原文一致。如果分词阶段直接丢了大小写、标点或空格,那么模型就永远学不到这些信息。
几个基本概念
在正式看代码前,先把几个词说清楚。
语料库(corpus):模型用来学习的文本集合。可以是一本文学作品,也可以是网页、书籍、论文、代码、对话数据的混合。
词元(token):文本被切分后的基本单位。token不一定是“单词”,也可以是字符、子词、字节片段、标点、空格,甚至是特殊控制符。
词表(vocabulary):所有允许出现的token集合。词表中的每个token都会被分配一个整数id。
索引映射(token to id):把token映射成整数的字典。例如:
"the" -> 1
"time" -> 19
"machine" -> 50
反向映射(id to token):把整数id还原回token,方便模型生成后解码成人类可读文本。
这些概念串起来就是:
原始文本
-> 词元化 tokenize
-> 词表 vocabulary
-> token id
-> Embedding
-> Transformer
词元粒度如何选择
最朴素的想法是按单词切分:
"the time machine" -> ["the", "time", "machine"]
这很符合人类直觉,但训练GPT时会遇到问题。
| 切分方式 | 优点 | 问题 |
|---|---|---|
| 按字符 | 词表极小,几乎没有未知词 | 序列太长,模型要自己从字符拼出单词 |
| 按单词 | 语义直观,序列较短 | 词表爆炸,新词和拼写变体难处理 |
| 按子词 / BPE | 兼顾词表大小和序列长度 | Tokenizer训练和实现更复杂 |
| 按字节级BPE | 可覆盖任意Unicode文本,几乎无未知词 | token边界不一定符合人类直觉 |
早期教学代码常用word或char,因为它们容易理解。现代LLM更常使用byte-level BPE或类似的子词算法,因为真实语料太复杂:中文没有天然空格,代码里有大量符号,数字和URL格式变化多,英文里还有词形变化、大小写和专有名词。
所以文本预处理的核心不是“怎么split字符串”,而是在取舍:
词表大小、序列长度、可逆性、多语言覆盖、训练效率
下面先用一个小数据集理解最基础的word/char词元化和词表构建,再过渡到GPT使用的BPE分词器。
教学示例:从时光机器数据集开始
读取数据集
先从一个小语料开始。D2L常用《The Time Machine》做语言模型入门数据集,因为它足够小,方便我们观察每一步结果。
这里的 read_time_machine() 做了三件事:
- 读取原始文本的所有行。
- 用正则表达式把非英文字母替换成空格。
- 去掉首尾空格,并统一转成小写。
这是一种教学用清洗方式。它让后面的word/char统计更简单,但也会丢掉标点、大小写和数字。对于现代GPT预训练,这种清洗通常太激进,因为模型需要学习标点、格式、代码、大小写、空格等细节。也就是说:教学场景可以简化文本,生成式大模型训练则更强调保留原文信息。
import collections
import re
from d2l import torch as d2l
d2l.DATA_HUB['time_machine'] = (d2l.DATA_URL + 'timemachine.txt','090b5e7e70c295757f55df93cb0a180b9691891a')
def read_time_machine():
with open(d2l.download('time_machine'),'r') as f:
lines = f.readlines()
return [re.sub('[^A-Za-z]+',' ',line).strip().lower() for line in lines]
lines = read_time_machine()
print(f'文本总行数:{len(lines)}')
print(lines[0])
print(lines[10])
文本总行数:3221
the time machine by h g wells
twinkled and his usually pale face was flushed and animated the
词元化:把字符串切成可计数的单位
词元化的目标,是把一整段字符串切成一个个离散单位。后面构建词表、统计词频、转索引,都是基于这些单位完成的。
下面这个 tokenize() 支持两种模式:
word:按空格切成单词。char:按字符切分。
这两个模式正好能体现词元粒度的差异。
按word切分时,序列更短,但词表更大。例如 traveller、travelling、travelled 可能是不同token。
按char切分时,词表很小,但序列会变长。例如 machine 会变成 m a c h i n e 七个token,模型需要更多步才能看到完整单词。
def tokenize(lines,token='word'):
if token == 'word':
return [line.split() for line in lines]
elif token == 'char':
return [list(line) for line in lines]
else:
print("错误: 未知的词元类型: " + token)
tokens = tokenize(lines)
for i in range(11):
print(tokens[i])
['the', 'time', 'machine', 'by', 'h', 'g', 'wells']
[]
[]
[]
[]
['i']
[]
[]
['the', 'time', 'traveller', 'for', 'so', 'it', 'will', 'be', 'convenient', 'to', 'speak', 'of', 'him']
['was', 'expounding', 'a', 'recondite', 'matter', 'to', 'us', 'his', 'grey', 'eyes', 'shone', 'and']
['twinkled', 'and', 'his', 'usually', 'pale', 'face', 'was', 'flushed', 'and', 'animated', 'the']
词表:把离散符号变成整数
词元化之后,我们得到的是字符串列表,但模型仍然不能直接使用字符串。下一步要建立词表,把每个token分配一个整数id。
词表要解决三个问题:
- 编号问题:每个token必须有唯一id。
- 未知词问题:测试时可能出现训练集中没有的token,需要用
<unk>兜底。 - 频率问题:高频token通常更重要,可以优先保留;低频token可以过滤或交给子词机制处理。
实际是一个字典,用来将字符串类型的词元映射到从0开始的数字索引中。我们对训练集的所有文档合并对他们的唯一词元进行统计,得到的结果为语料库。根据每个词元出现的频率为其分配索引。
class Vocab:
def __init__(self, tokens=None, min_freq=0, reserved_tokens=None):
if tokens is None:
tokens = []
if reserved_tokens is None:
reserved_tokens = []
# 按照出现频率排序
counter = count_corpus(tokens)
"""
* counter.items() 是字典的一个方法。它返回一个包含字典所有项((键, 值) 对,即 (key, value) tuples)的视图对象。例如,如果 counter 是 {'the': 10, 'a': 5},counter.items() 会提供像 ('the', 10) 和 ('a', 5) 这样的元素。
* sorted() 在对元素进行比较之前,会先将每个元素传递给 key 指定的函数,然后根据这些函数返回的结果进行排序。
* lambda匿名函数见下
"""
self._token_freqs = sorted(counter.items(), key=lambda x: x[1], reverse=True)
# 未知词元的索引为0
self.idx_to_token = ['<unk>'] + reserved_tokens
self.token_to_idx = {token: idx for idx, token in enumerate(self.idx_to_token)}
for token, freq in self._token_freqs:
if freq < min_freq:
break
if token not in self.token_to_idx:
self.idx_to_token.append(token)
self.token_to_idx[token] = len(self.idx_to_token) - 1
# 定义指定方法`__len__`. 这个方法允许在这个类对象上使用内置函数 `len()`。 (e.g., `len(my_vocab)`).
def __len__(self):
return len(self.idx_to_token)
# 定义特殊方法`__getitem__`。这允许使用方括号表示法进行查找(例如,`my_vocab['hello']`或`my_voab[['hello','world']])。
def __getitem__(self, tokens):
if not isinstance(tokens, (list, tuple)):
return self.token_to_idx.get(tokens, self.unk)
return [self.__getitem__(token) for token in tokens]
# 进行反向查找index--->token
def to_tokens(self, indices):
if not isinstance(indices, (list, tuple)):
return self.idx_to_token[indices]
return [self.idx_to_token[index] for index in indices]
# `@property` 是一个装饰器,允许我们向调用属性的方式一样来调用函数。E.g., `my_vocab.unk` instead of `my_vocab.unk()`.
@property
def unk(self):
return 0
@property
def token_freqs(self):
return self._token_freqs
def count_corpus(tokens):
"""统计词元出现的频率"""
if len(tokens) == 0 or isinstance(tokens[0], list):
# 将词元展开成一个列表
tokens = [token for line in tokens for token in line]
return collections.Counter(tokens)
lambda 匿名函数 (Anonymous Function): lambda x: x[1] 定义了一个简单的匿名函数(没有名字的函数)。
lambda: 关键字,表示这是一个 lambda 函数。x: 参数名,代表输入给这个 lambda 函数的单个元素。在这个上下文中,因为sorted正在处理counter.items()的结果,所以x会是一个(键, 值)对的元组,例如('the', 10)。: x[1]: 冒号后面是函数的返回值。x[1]表示取元组x中索引为 1 的元素,也就是那个键值对中的值 (value)——即词元的频率。key=lambda x: x[1]的整体作用: 告诉sorted函数,在排序(词元, 频率)对时,应该根据频率(元组的第二个元素x[1])来排序,而不是根据词元本身(元组的第一个元素x[0])。
这里按频率排序,不只是为了输出好看。语言数据有非常明显的长尾分布:少数高频词出现极多,大量低频词只出现一两次。高频词如果被分配到稳定的id,模型在训练早期就能频繁看到它们,对基础语言规律的学习更稳定。
但这种传统 Vocab 方式也有明显缺陷:只要一个词没在训练集中出现过,它就只能变成 <unk>。例如训练集中没有 tokenization,测试时这个词就丢成未知词。GPT类模型通常不能接受这种信息损失,所以后面才会转向BPE这类子词分词方法。
vocab = Vocab(tokens)
print(list(vocab.token_to_idx.items())[:10])
[('<unk>', 0), ('the', 1), ('i', 2), ('and', 3), ('of', 4), ('a', 5), ('to', 6), ('was', 7), ('in', 8), ('that', 9)]
for i in [0,10]:
print('文本: ', tokens[i])
print('索引: ', vocab[tokens[i]])
文本: ['the', 'time', 'machine', 'by', 'h', 'g', 'wells']
索引: [1, 19, 50, 40, 2183, 2184, 400]
文本: ['twinkled', 'and', 'his', 'usually', 'pale', 'face', 'was', 'flushed', 'and', 'animated', 'the']
索引: [2186, 3, 25, 1044, 362, 113, 7, 1421, 3, 1045, 1]
def load_corpus_time_machine(max_tokens=-1):
"""返回时光机器数据集的词元索引列表和词表"""
lines = read_time_machine()
tokens = tokenize(lines, 'char')
vocab = Vocab(tokens)
# 因为时光机器数据集中的每个文本行不一定是一个句子或段落,所以将所有的文本行展平到一个列表中
corpus = [vocab[token] for line in tokens for token in line]
if max_tokens > 0:
corpus = corpus[:max_tokens]
return corpus, vocab
corpus, vocab = load_corpus_time_machine()
len(corpus), len(vocab)
(170580, 28)
这里特意使用了 char 模式,所以最终词表只有28个左右。字符级语言模型很好教学,因为几乎不会出现未知词;但它不是训练大模型的理想方案。
原因是序列长度会被放大。比如一句话按单词可能只有10个token,按字符可能变成50个token。Transformer的Attention计算复杂度和序列长度强相关,序列越长,计算和显存压力越大。因此真实GPT会使用更折中的子词Tokenizer。
基于RustBPE和tiktoken实现高效对话模型分词器(Tokenizer)
前面手写的 Vocab 已经能完成“词元 -> 整数id”的映射,但它更适合教学。真正训练GPT类语言模型时,直接按单词切分会遇到词表爆炸和未知词问题;直接按字符切分虽然没有未知词,但序列会变得很长,模型需要花大量上下文长度去拼出常见单词。
BPE(Byte Pair Encoding)正好介于二者之间:常见片段会被合并成较长token,罕见词又能拆回更小的字节或子词片段。因此现代LLM通常不会使用传统word vocab,而是使用byte/subword级别的Tokenizer。
下面结合 /Volumes/PSSD/工具/nanochat/nanochat 这个项目,看看一个工程化分词器是如何接入训练流程的。
为什么现代LLM选择BPE
BPE解决的是word和char之间的矛盾。
如果按单词切分:
unbelievable -> ["unbelievable"]
unbelievably -> ["unbelievably"]
两个词会占用两个不同词表项。训练集中没见过的新词还会变成 <unk>。
如果按字符切分:
unbelievable -> ["u", "n", "b", "e", "l", "i", "e", "v", "a", "b", "l", "e"]
虽然不会有未知词,但序列太长,模型要自己反复学习常见字符组合。
BPE的想法是:常见组合就合并成一个token,罕见组合就拆成更小单位。
unbelievable -> ["un", "believ", "able"]
unbelievably -> ["un", "believ", "ably"]
这只是示意,不同Tokenizer的实际切分结果可能不同。关键在于:BPE不是用人工语言学规则规定词根词缀,而是从语料统计中自动学习哪些片段值得合并。
因此BPE有几个优势:
- 高频词或高频片段会变短,节省上下文长度。
- 低频词可以拆成子词,不必变成
<unk>。 - byte-level BPE可以从UTF-8字节出发,覆盖中文、英文、代码、emoji等复杂文本。
- 词表大小可以人为控制,例如32K、50K、100K。
所以训练LLM时,Tokenizer本质上是在做一种压缩:把文本压成更短、更适合模型学习的token序列,同时尽量保持可逆。
BPE的核心思想
BPE最初是一种压缩算法,用在分词器里时,可以理解为“不断合并最常见的相邻符号对”:
- 先把文本变成基础符号序列。byte-level BPE通常从UTF-8字节开始,因此天然能覆盖中文、英文、emoji和各种符号。
- 统计所有相邻pair的出现频率,例如
("t", "h")、("i", "n")。 - 选择出现频率最高的pair,合并成一个新token。
- 重复合并,直到达到目标词表大小。
- 编码时按学到的merge规则把文本切成token id;解码时再把token id还原为字节/字符串。
举个直观例子。假设语料中经常出现:
l o w
l o w e r
l o w e s t
一开始每个字符都是基础符号。如果相邻pair ("l", "o") 最常见,就先合并成 lo;如果 ("lo", "w") 又很常见,就继续合并成 low。这样高频片段逐渐变成更长token。
BPE训练完成后,词表里既有基础字节,也有合并出来的常见片段。编码新文本时,就按照这些merge规则尽可能合并;遇到罕见词,也可以退回到更小的字节片段。
一个极简版的BPE训练过程可以写成这样:
def get_stats(ids):
counts = collections.Counter()
for pair in zip(ids, ids[1:]):
counts[pair] += 1
return counts
def merge(ids, pair, new_id):
new_ids = []
i = 0
while i < len(ids):
if i < len(ids) - 1 and (ids[i], ids[i + 1]) == pair:
new_ids.append(new_id)
i += 2
else:
new_ids.append(ids[i])
i += 1
return new_ids
def train_bpe(text, vocab_size):
assert vocab_size >= 256
ids = list(text.encode("utf-8"))
merges = {}
vocab = {idx: bytes([idx]) for idx in range(256)}
for i in range(vocab_size - 256):
stats = get_stats(ids)
pair = max(stats, key=stats.get)
new_id = 256 + i
ids = merge(ids, pair, new_id)
merges[pair] = new_id
vocab[new_id] = vocab[pair[0]] + vocab[pair[1]]
return merges, vocab
这个版本只展示算法骨架。nanochat没有在Python里手写完整BPE,而是用 rustbpe 负责训练,用 tiktoken 负责高效推理编码。
这也是工程实践和教学代码的区别:教学代码追求看懂算法,真实训练追求速度、稳定性、可逆性和大规模数据吞吐。
nanochat的Tokenizer结构
nanochat的核心分词器在 nanochat/tokenizer.py 中。文件里保留了两套实现:
HuggingFaceTokenizer:基于HuggingFacetokenizers,可以训练也可以推理,接口比较通用。RustBPETokenizer:使用rustbpe训练merge规则,再构造tiktoken.Encoding做快速encode/decode,这是项目实际默认使用的版本。
项目还定义了一组特殊token:
SPECIAL_TOKENS = [
"<|bos|>",
"<|user_start|>",
"<|user_end|>",
"<|assistant_start|>",
"<|assistant_end|>",
"<|python_start|>",
"<|python_end|>",
"<|output_start|>",
"<|output_end|>",
]
这些特殊token不来自普通文本统计,而是由训练脚本显式加入词表。它们负责告诉模型“文档开始了”“用户消息开始/结束了”“助手消息开始/结束了”“工具调用开始/结束了”等结构信息。
其中最重要的是 <|bos|>。在预训练数据加载阶段,nanochat会让每个文档都以BOS开头。这样模型在一个训练样本内可以明确知道新文档边界,不会把两个无关文档误以为是连续上下文。
特殊token的作用不是表示普通词义,而是表示数据结构。可以把它们理解成“给模型看的标记符”。
例如聊天数据如果只拼成普通文本:
你好
你好,我是助手
模型很难知道哪一句是用户说的,哪一句是助手应该学习输出的。加入特殊token后,结构就清楚了:
<|user_start|>你好<|user_end|>
<|assistant_start|>你好,我是助手<|assistant_end|>
这就是为什么对话模型不仅需要Tokenizer,还需要chat template或conversation rendering。模型不是天生知道“用户”和“助手”的角色,这些角色边界必须在训练数据中显式表达。
GPT风格的预切分正则
BPE不是直接在整篇文本上无约束地合并。nanochat先用一个接近GPT-4风格的正则表达式把文本预切分成较合理的片段:
SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""
这个正则大致处理几类文本:
- 英文缩写后缀,如
's、're、've。 - 字母序列,即普通单词片段。
- 数字序列。nanochat使用
\p{N}{1,2},也就是最多两个数字一组。项目注释里说明,32K词表下{1,2}比{1,3}更适合小模型,因为三位数字组合会消耗更多词表空间。 - 标点、换行和空白。空格是否贴在后续单词上,会直接影响token边界。
这样做的目的不是完成最终分词,而是给BPE一个更稳定的搜索空间。BPE再在这些片段内部学习常见字节组合。
如果完全不做预切分,BPE可能学到一些跨类别的奇怪合并,比如把单词末尾和下一个标点、换行、空格以不稳定方式绑在一起。预切分相当于先给文本划出比较合理的边界,再让BPE在边界内部学习合并。
这一步体现了Tokenizer设计中的一个重要原则:不是所有统计上高频的pair都值得合并。我们希望合并规则既能压缩文本,又不要破坏太多自然边界。
训练分词器:从Parquet文本到tokenizer.pkl
nanochat的预训练数据不是一个大txt,而是一组Parquet分片。nanochat/dataset.py 负责下载和读取这些分片:
BASE_URL = "https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle/resolve/main"
MAX_SHARD = 6542
DATA_DIR = os.path.join(base_dir, "base_data_climbmix")
读取时,训练集使用除最后一个分片外的所有Parquet,验证集固定使用最后一个分片:
def parquets_iter_batched(split, start=0, step=1):
parquet_paths = list_parquet_files()
parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]
for filepath in parquet_paths:
pf = pq.ParquetFile(filepath)
for rg_idx in range(start, pf.num_row_groups, step):
rg = pf.read_row_group(rg_idx)
texts = rg.column("text").to_pylist()
yield texts
scripts/tok_train.py 会把这些batch展开成文本迭代器,并限制训练材料规模:
def text_iterator():
nchars = 0
for batch in parquets_iter_batched(split="train"):
for doc in batch:
doc_text = doc
if len(doc_text) > args.doc_cap:
doc_text = doc_text[:args.doc_cap]
nchars += len(doc_text)
yield doc_text
if nchars > args.max_chars:
return
默认参数是:
--max-chars 2000000000:最多用20亿字符训练分词器。--doc-cap 10000:单篇文档最多截到10000字符,避免超长文档主导merge统计。--vocab-size 32768:词表大小为32K。
这里的 doc_cap 很重要。Tokenizer训练本质上是在统计pair频率,如果少数超长文档占比过大,它们的写作风格、重复模板或噪声可能会主导merge规则。限制单篇文档长度,是为了让训练Tokenizer时看到的语料分布更均衡。
vocab_size 也是一个取舍。词表太小,文本会被切得很碎,序列变长;词表太大,Embedding和输出层参数变多,低频token也更难学好。32K是很多小到中等规模模型常用的折中点。
训练入口非常直接:
tokenizer = RustBPETokenizer.train_from_iterator(text_iter, args.vocab_size)
tokenizer.save(tokenizer_dir)
在 RustBPETokenizer.train_from_iterator 里,流程可以拆成两步:
tokenizer = rustbpe.Tokenizer()
vocab_size_no_special = vocab_size - len(SPECIAL_TOKENS)
tokenizer.train_from_iterator(text_iterator, vocab_size_no_special, pattern=SPLIT_PATTERN)
mergeable_ranks = {bytes(k): v for k, v in tokenizer.get_mergeable_ranks()}
tokens_offset = len(mergeable_ranks)
special_tokens = {name: tokens_offset + i for i, name in enumerate(SPECIAL_TOKENS)}
enc = tiktoken.Encoding(
name="rustbpe",
pat_str=tokenizer.get_pattern(),
mergeable_ranks=mergeable_ranks,
special_tokens=special_tokens,
)
这里有几个细节:
- 词表大小要先减去特殊token数量,因为特殊token后面单独插入。
- 普通token来自BPE merge规则,特殊token直接映射到词表尾部。
- 保存到磁盘的是
tokenizer.pkl,后续训练和推理会通过get_tokenizer()加载它。
训练后还有一个很关键的自检:
encoded = tokenizer.encode(test_text)
decoded = tokenizer.decode(encoded)
assert decoded == test_text
分词器必须满足 decode(encode(text)) == text。如果这个性质不成立,说明文本经过分词后会丢信息,语言模型训练就不可靠。
这和前面的D2L清洗代码形成对比。教学示例里我们把标点和大小写删掉,是为了降低入门难度;但GPT的Tokenizer应该尽量保留原文,因为生成模型最终要还原自然文本、代码格式、换行、标点和各种符号。
encode、decode和批量编码
nanochat把常用接口封装成统一形式:
def encode(self, text, prepend=None, append=None, num_threads=8):
if isinstance(text, str):
ids = self.enc.encode_ordinary(text)
if prepend is not None:
ids.insert(0, prepend_id)
if append is not None:
ids.append(append_id)
elif isinstance(text, list):
ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)
...
return ids
def decode(self, ids):
return self.enc.decode(ids)
几个设计点值得注意:
- 普通文本使用
encode_ordinary,不会意外把文本里的特殊token字符串当成控制符。 prepend和append可以传特殊token字符串,也可以直接传token id。- 传入
list[str]时走encode_ordinary_batch,可以多线程批量编码,这是训练数据加载时的主要路径。
例如:
tokenizer = get_tokenizer()
bos = tokenizer.get_bos_token_id()
ids = tokenizer.encode("The chemical formula of water is", prepend=bos)
text = tokenizer.decode(ids)
此时 ids 是模型真正看到的整数序列,decode 则把模型生成的token id还原为人类可读文本。
对话数据如何渲染成token
预训练阶段只需要普通文档文本,但SFT/RL阶段要处理多轮对话。nanochat在 RustBPETokenizer.render_conversation 中把结构化conversation渲染成两组数组:
ids: list[int]:完整token序列。mask: list[int]:同长度监督掩码,1表示该位置参与助手回复训练,0表示不参与loss。
简化后的逻辑如下:
ids, mask = [], []
add_tokens(bos, 0)
for message in messages:
if message["role"] == "user":
add_tokens(user_start, 0)
add_tokens(tokenizer.encode(content), 0)
add_tokens(user_end, 0)
elif message["role"] == "assistant":
add_tokens(assistant_start, 0)
add_tokens(tokenizer.encode(content), 1)
add_tokens(assistant_end, 1)
也就是说,用户输入、分隔符、工具输出等内容主要用于提供上下文,不作为模型要学习预测的答案;助手回复正文和 <|assistant_end|> 才是SFT阶段重点监督的部分。
如果conversation第一条是system消息,nanochat会把system内容合并进后面的user消息。这是一个工程折中:保留system指令语义,同时保持后续消息严格按 user -> assistant -> user -> assistant 交替,方便做一致性检查。
训练数据加载:从文档到input/target
分词器训练好后,真正训练基础模型时,scripts/base_train.py 会加载:
tokenizer = get_tokenizer()
token_bytes = get_token_bytes(device=device)
vocab_size = tokenizer.get_vocab_size()
然后用 nanochat/dataloader.py 构造训练和验证加载器:
train_loader = tokenizing_distributed_data_loader_with_state_bos_bestfit(
tokenizer,
args.device_batch_size,
args.max_seq_len,
split="train",
device=device,
)
这个DataLoader做的事情不只是“读文本”。它完成了完整的预处理链路:
- 从Parquet row group读取一批原始字符串。
- 多线程批量编码,并在每篇文档前加BOS。
- 把不同长度的文档装进固定长度训练行。
- 用Best-Fit策略尽量选择能放入剩余空间的最长文档。
- 如果没有文档能完整放下,就裁剪一篇文档填满剩余空间。
- 生成
inputs和targets,用于自回归next-token prediction。
这里的“固定长度训练行”不是随便定的。GPU训练需要把一批样本堆成形状规整的张量,例如:
inputs.shape = (B, T)
targets.shape = (B, T)
但真实文档长度是不规则的。如果每篇文档单独训练,会造成大量padding和低效计算;如果简单拼接,又可能把完全无关的文档接在一起。nanochat的做法是把文档尽量打包到固定长度行里,并用BOS标记文档边界,在效率和语义边界之间做折中。
关键代码是:
row_capacity = T + 1
token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads)
cpu_inputs.copy_(row_buffer[:, :-1])
cpu_targets.copy_(row_buffer[:, 1:])
yield inputs, targets, state_dict
为什么是 T + 1?因为模型输入长度是 T,目标也是 T,但目标相当于输入向左错一位:
row_buffer: [BOS, t1, t2, t3, ..., tT]
inputs: [BOS, t1, t2, ..., t(T-1)]
targets: [t1, t2, t3, ..., tT]
这样每个位置都在训练“根据前文预测下一个token”。
这就是GPT预训练最核心的监督信号:不需要人工标注问题答案,只要有大量文本,就能把每个token的下一个token当作标签。
输入: The chemical formula of water is
目标: chemical formula of water is H
模型不是一次学会“整篇文章”,而是在每个位置都学习条件概率:
因此DataLoader构造 inputs / targets 的方式,直接决定了模型每一步到底在预测什么。
nanochat这个加载器还有两个实用特性:
- DDP分片:不同rank读取不同row group,避免多卡重复消费相同数据。
- 断点恢复:
state_dict = {"pq_idx": pq_idx, "rg_idx": rg_idx, "epoch": epoch},训练中断后可以近似从上次位置继续。
为什么要缓存token_bytes
scripts/tok_train.py 训练完分词器后,还会生成 token_bytes.pt:
token_strings = [tokenizer.decode([token_id]) for token_id in range(vocab_size)]
token_bytes = []
for token_id in range(vocab_size):
token_str = token_strings[token_id]
if token_str in special_set:
token_bytes.append(0)
else:
token_bytes.append(len(token_str.encode("utf-8")))
它记录每个token对应多少UTF-8字节。特殊token不属于普通文本内容,所以字节数记为0。
这个张量用于评估 bits per byte(bpb)。普通loss会受词表大小影响:词表大时token更长,平均token loss和小词表不完全可比。bpb则按字节归一化:
bpb = total_nats / (math.log(2) * total_bytes)
因此当你调整BPE词表大小、预切分规则或训练语料时,bpb比“平均每token loss”更适合横向比较。
换句话说,token级loss回答的是“每预测一个token平均错多少”,但不同Tokenizer的token长度不同,不能直接公平比较。bpb回答的是“每压缩一个字节平均需要多少bit的信息量”,更接近跨Tokenizer的统一尺度。
分词器压缩率评估
scripts/tok_eval.py 会拿当前分词器和GPT-2、GPT-4风格分词器对比压缩率。核心指标是:
encoded = tokenizer.encode(text)
encoded_bytes = text.encode("utf-8")
ratio = len(encoded_bytes) / len(encoded)
ratio 表示平均每个token承载多少字节。ratio越高,通常说明同样文本会被编码成更短的token序列,模型能在固定上下文长度里看到更多内容。但这不是唯一目标:过大的词表会增加embedding和lm head开销,也可能让低频token学得不好。因此nanochat默认选用32K词表,是压缩率、训练速度和模型容量之间的折中。
整体流程总结
把传统文本预处理和nanochat工程代码连起来,可以得到完整的数据路径:
原始数据集
-> 清洗/重打包为Parquet分片
-> 训练BPE分词器 tokenizer.pkl
-> 生成 token_bytes.pt
-> DataLoader读取Parquet文本
-> tokenizer批量编码并添加BOS
-> 文档打包为固定长度token行
-> 构造 inputs / targets
-> GPT执行next-token prediction训练
如果在本地复现nanochat的分词器训练,大致命令是:
python -m nanochat.dataset -n 170
python -m scripts.tok_train --vocab-size 32768
python -m scripts.tok_eval
前半篇的 Vocab 让我们理解了“字符串词元如何映射成整数”;后半篇的nanochat源码则展示了工业化训练里还要解决的问题:可逆编码、特殊token、批量编码、文档边界、分布式读取、固定长度打包,以及与模型loss可比性相关的评估指标。文本预处理不是模型训练前的一步小清洗,而是直接决定模型看见什么、如何看见、以及训练信号是否干净的一整套数据管线。
尺素
0 封尺素
正在取尺素...