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() 做了三件事:

  1. 读取原始文本的所有行。
  2. 用正则表达式把非英文字母替换成空格。
  3. 去掉首尾空格,并统一转成小写。

这是一种教学用清洗方式。它让后面的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。

词表要解决三个问题:

  1. 编号问题:每个token必须有唯一id。
  2. 未知词问题:测试时可能出现训练集中没有的token,需要用 <unk> 兜底。
  3. 频率问题:高频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最初是一种压缩算法,用在分词器里时,可以理解为“不断合并最常见的相邻符号对”:

  1. 先把文本变成基础符号序列。byte-level BPE通常从UTF-8字节开始,因此天然能覆盖中文、英文、emoji和各种符号。
  2. 统计所有相邻pair的出现频率,例如 ("t", "h")、("i", "n")。
  3. 选择出现频率最高的pair,合并成一个新token。
  4. 重复合并,直到达到目标词表大小。
  5. 编码时按学到的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:基于HuggingFace tokenizers,可以训练也可以推理,接口比较通用。
  • 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做的事情不只是“读文本”。它完成了完整的预处理链路:

  1. 从Parquet row group读取一批原始字符串。
  2. 多线程批量编码,并在每篇文档前加BOS。
  3. 把不同长度的文档装进固定长度训练行。
  4. 用Best-Fit策略尽量选择能放入剩余空间的最长文档。
  5. 如果没有文档能完整放下,就裁剪一篇文档填满剩余空间。
  6. 生成 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

模型不是一次学会“整篇文章”,而是在每个位置都学习条件概率:

P(xt+1∣x1,x2,...,xt)P(x_{t+1} \mid x_1, x_2, ..., x_t)

因此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 封尺素

正在取尺素...