如何在大规模数据集上训练分词器？T5波斯语分词器训练卡顿问题求解

问题解决方法 #

1. 卡住原因定位 #

你当前代码卡住并非是train_from_iterator不支持800万条规模的数据集，问题出在自定义的batch_iterator用了dataset[i: i + batch_length]的分片索引逻辑，这类操作会触发datasets库把对应范围的全量数据一次性加载到内存，800万条数据的内存占用超过Colab Pro的内存上限，才会出现卡死的情况。

2. 最优方案：流式加载全量数据直接训练（无需拆分合并） #

SentencePiece的训练逻辑本身支持流式输入全量数据，不需要拆分训练再合并，你只需要修改数据集加载方式和迭代逻辑，避免全量数据入内存即可，最终得到的分词器是基于全量800万条数据训练的，效果远好于拆分后合并的方案：

import datasets

from t5_tokenizer_model import SentencePieceUnigramTokenizer


vocab_size = 32_000
input_sentence_size = None # 不需要限制大小，直接用全量数据

# 加载数据集时开启流式模式，不会一次性把全量数据加载到内存
dataset = datasets.load_dataset(
    "oscar", 
    name="unshuffled_deduplicated_fa", 
    split="train",
    streaming=True
)

tokenizer = SentencePieceUnigramTokenizer(unk_token="<unk>", eos_token="</s>", pad_token="<pad>")

# 重写迭代器，逐行遍历流式数据集，无内存溢出风险
def batch_iterator(batch_length=100):
    batch = []
    for idx, item in enumerate(dataset):
        if input_sentence_size is not None and idx >= input_sentence_size:
            break
        batch.append(item["text"])
        if len(batch) == batch_length:
            yield batch
            batch = []
    # 返回最后不足100条的剩余数据
    if batch:
        yield batch


# 直接用全量数据训练，不会卡住
tokenizer.train_from_iterator(
    iterator=batch_iterator(),
    vocab_size=vocab_size,
    show_progress=True,
)

# 保存分词器
tokenizer.save("/content/drive/MyDrive/Pouramini/tokenizer.json")

3. 分块训练再合并的方案（仅特殊场景使用，效果不如全量训练） #

如果你的场景确实必须拆分训练，可以按照以下步骤操作：

• 把全量数据集随机拆分为多个10万条的子集
• 分别在每个子集上训练分词器，得到对应的词汇表
• 合并所有词汇表，统计每个token的出现频率、Unigram算法得分
• 按照得分从高到低排序，取前32000个token作为最终词汇表，初始化新的分词器加载该词汇表即可

注意：这种拆分训练的方式会损失全局词汇统计的准确性，分词效果会弱于全量流式训练的结果。

内容的提问来源于stack exchange，提问作者Ahmad