前言
上一篇我们已经对海量的OpenWebText数据完成了清洗去重，并选择将其存为高效的Arrow格式，这个步骤完成后，我们已经为训练自己的大语言模型（LLM）打下了坚实的基础。那么问题来了，存储为Arrow格式的数据，接下来，如何将这些原始文本数据转化为模型能够“消化”的格式呢？
别担心，这篇教程将手把手完成以下关键步骤：
- BPE分词器：基于我们的Arrow数据集，训练一个专属的BPE（Byte-Pair Encoding）分词器。我们会探讨不同方案，并给出推荐。
- 构建数据管道：学习如何使用PyTorch的
Dataset和DataLoader，结合训练好的BPE词表，高效地将Arrow数据构造成模型训练所需的batch。 - 攻克长文本难题：OpenWebText中不乏长篇大论。对于decoder-only这类模型，粗暴截断可不是好主意。我们将调研并实现一种更合适的长文本处理方案，同时兼顾大数据集的处理效率。
- 完整代码实战：提供所有步骤的完整、可运行Python代码。
1. BPE分词器
在我们开始训练LLM之前，首先需要一个“字典”，也就是分词器(Tokenizer)。它的任务是将原始文本切分成模型能够理解的最小单元——词元(Token)。对于大规模语料，字节对编码 (Byte-Pair Encoding, BPE) 是一种非常流行且高效的分词算法。特别是其变种 字节级BPE (Byte-Level BPE)，由于其能覆盖所有字节，从而避免未知字符（UNK问题），在GPT系列模型中得到了广泛应用 。
1.1 为何需要定制分词器
虽然有很多预训练好的分词器，但它们可能并不完全适合我们精心清洗过的OpenWebText数据集。原因有二：
- 词表不匹配（Vocabulary Mismatch）：预训练分词器的词表是基于通用语料构建的，可能无法很好地覆盖特定数据集中的高频词汇或领域术语。
- 数据分布差异（Data Distribution Differences）：我们的数据清洗和去重过程可能改变了文本的统计特性，使得通用分词器的切分方式不再最优。
定制一个分词器，可以让它更适应我们的数据，从而可能带来更好的压缩率（用更少的token表示相同信息）、产生更具语义意义的token单元，并最终提升模型的性能 。
1.2 数据集加载
由于OpenWebText数据集非常庞大，一次性将所有文本加载到内存中训练分词器是不现实的。我们需要一个能够分批从Arrow文件中迭代读取文本数据的迭代器。
def get_text_iterator_for_tokenizer(dataset_path, text_column_name="text", batch_size=1000):
try:
dataset = load_from_disk(dataset_path)
print(f"成功从 '{dataset_path}' 加载数据集，包含 {len(dataset)} 条记录。")
except Exception as e:
print(f"加载数据集 '{dataset_path}' 失败: {e}")
if not isinstance(dataset, HFDataset):
raise TypeError("load_from_disk 未返回预期的 Hugging Face Dataset 对象。")
if text_column_name not in dataset.column_names:
raise ValueError(f"指定的文本列 '{text_column_name}' 不在数据集中。可用列: {dataset.column_names}")
print(f"开始从数据集的 '{text_column_name}' 列生成文本迭代器...")
for i in range(0, len(dataset), batch_size):
yield dataset[i : i + batch_size][text_column_name]
1.3 BPE训练
方案一：Hugging Face ByteLevelBPETokenizer 高层封装 (推荐初学者)
这是最直接且推荐初学者上手的方法。Hugging Face tokenizers 库中的 ByteLevelBPETokenizer 类已经为我们封装好了底层的BPE模型、字节级预分词器 (pre-tokenizer) 和解码器 (decoder)。
关键参数详解：
vocab_size(整数): 词表大小，这是一个需要仔细调整的超参数。它包括了最初的256个字节以及通过BPE合并规则学习到的新词元。常见的大小有几万到十几万，例如GPT-2使用的是50257。选择合适的词表大小是一个权衡：太小可能导致文本被切分得过细，序列变长；太大则会增加模型Embedding层参数量，且可能包含很多低频无用的词元。对于主要为英文的OpenWebText，45k-65k是一个不错的起点。但是考虑后续还要添加中文支持，所以这里我先设置为40k，后续看训练效果再进行调整。min_frequency(整数): 一个词元对(pair)需要出现的最小次数才能被合并成一个新的词元。special_tokens(列表): 特殊词元列表，例如 ["<s>", "<pad>", "</s>", "<unk>", "<mask>"]。它们的定义顺序非常重要，会直接决定它们在词表中的ID（通常从0开始）。虽然字节级BPE理论上因为覆盖了所有字节而不需要 <unk> 来表示未知字符，但将其作为通用特殊词元加入仍是良好实践，因为很多模型架构或下游任务可能会用到它 。add_prefix_space(布尔值，ByteLevelBPETokenizer构造函数中默认为False): 这个参数控制是否在每个单词的开头添加一个前导空格（通常是Ġ符号的内部表示）再进行分词。对于类似GPT-2的分词风格，通常推荐设置为True。ByteLevelBPETokenizer内部会使用pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space)。当设置为True时，可以更好地区分粘连词和词首的词，例如 "the" 和 "inthe" 中的 "the"。unicode_normalizer(字符串, 可选): 例如填入"nfkc"或"nfc"。虽然字节级BPE直接操作字节，但在预处理阶段进行Unicode归一化是一个好习惯。网页文本中充满了各种Unicode字符，同一个视觉上的字符可能有多种字节表示方式（比如带音调的字母的不同组合形式）。归一化能将它们统一成标准形式，这样BPE算法在学习合并规则时会更高效，生成的词表质量也更高，避免了因编码差异而产生的冗余词元。不过因为我们在数据清洗阶段就已经做过unicode规范化了，所以这里我们可以跳过这个步骤。trim_offsets(布尔值, 默认为 False): 是否在分词结果的偏移量(offset)中去除由add_prefix_space添加的前导空格。通常与add_prefix_space结合考虑，如果add_prefix_space为True，将trim_offsets也设为True可以使得解码后的文本更自然，不保留额外的Ġ标记。
参考代码如下：
from datasets import load_from_disk, load_dataset
from tokenizers import ByteLevelBPETokenizer
import os
# 假设 ARROW_DATASET_DIR 是你保存 Arrow 数据集的路径
ARROW_DATASET_DIR = "./openwebtext_arrow_dataset/train"
BPE_TOKENIZER_SAVE_DIR = "./bpe_tokenizer_from_arrow"
VOCAB_SIZE = 25000 # BPE词表大小示例
MIN_FREQUENCY = 2 # 最小词频示例
def get_openwebtext_dataset(arrow_data_path, split="train"):
return load_dataset("arrow", data_files={split: arrow_data_path + "/*.arrow"})[split]
def batch_iterator(dataset, text_column="text", batch_size=1000):
for i in range(0, len(dataset), batch_size):
yield dataset[i : i + batch_size][text_column]
def train_bpe_from_arrow_dataset(arrow_dataset_path: str, save_dir: str):
if not os.path.exists(arrow_dataset_path):
print(f"错误: Arrow数据集路径不存在: {arrow_dataset_path}")
return
print(f"从Arrow数据集加载数据: {arrow_dataset_path}")
dataset = get_openwebtext_dataset(arrow_dataset_path)
print(f"数据集加载完成，样本数: {len(dataset)}")
# 初始化并训练 BPE 分词器
tokenizer = ByteLevelBPETokenizer(add_prefix_space=False)
print("开始训练BPE分词器...")
tokenizer.train_from_iterator(
batch_iterator(dataset),
vocab_size=VOCAB_SIZE,
min_frequency=MIN_FREQUENCY,
special_tokens=[
"<s>",
"<pad>",
"</s>",
"<unk>",
"<mask>",
],
)
print("BPE分词器训练完成.")
os.makedirs(save_dir, exist_ok=True)
tokenizer.save(os.path.join(save_dir, "tokenizer.json"))
print(f"BPE分词器模型已保存到: {save_dir}/tokenizer.json")
# 调用BPE训练
if __name__ == "__main__":
train_bpe_from_arrow_dataset(ARROW_DATASET_DIR, BPE_TOKENIZER_SAVE_DIR)
优点:
- 实现非常简单，代码量少。
- 集成了字节级BPE的许多最佳实践，开箱即用。
缺点:
- 对于BPE模型本身的一些更底层的参数（例如dropout,continuing_subword_prefix,end_of_word_suffix），虽然可以在ByteLevelBPETokenizer的构造函数
__init__中设置（它们会传递给内部的models.BPE），但它们不是train_from_iterator方法的直接参数，控制起来不如方案二灵活。
方案二：从底层组件构建和训练BPE(为希望深入理解的进阶读者提供)
如果想更深入地理解BPE分词的每一个环节，或者需要更精细地控制某些参数，那么可以采用这种从底层组件构建的方式。这涉及到组合使用tokenizers.Tokenizer类以及models.BPE, pre_tokenizers.ByteLevel, decoders.ByteLevel,和trainers.BpeTrainer等模块。
这种方式允许你在初始化models.BPE时就设置一些高级参数：
- dropout (浮点数, 可选): BPE dropout的概率，这是一种正则化技术。在训练过程中，它会以一定的概率随机跳过某些最高频的词元对合并步骤。这样做可以使得模型学习到更多样化的子词表示，有时能提升模型的鲁棒性。
- continuing_subword_prefix (字符串, 可选): 用于标记非词首子词的前缀，例如BERT中常用的 "##"。 end_of_word_suffix (字符串, 可选): 用于标记词尾子词的后缀。 在 trainers.BpeTrainer 中，一个关键参数是：
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(): 必须显式指定初始字母表为所有256个字节。这是实现“字节级”BPE的核心。如果缺失这一步，BPE会基于文本中的字符进行训练，而不是字节 。
代码示例如下，因为方案选择上我们这次会偏向于使用方案一：ByteLevelBPETokenizer方式进行训练，所以下面的代码仅供参考，我也也没有实际跑过。
from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, normalizers
# 假设 openwebtext_dataset 和 batch_iterator 已定义
# vocab_size, min_frequency, special_tokens 也已定义 (同方案一)
# 1. 初始化 BPE 模型 (Model)
# 这里可以设置 dropout, continuing_subword_prefix 等参数
bpe_model = models.BPE(
# vocab=None, merges=None, # 从头训练
dropout=0.1, # 示例：启用BPE dropout
continuing_subword_prefix="##", # 示例：设置非词首子词前缀
# end_of_word_suffix="$", # 示例：设置词尾后缀
unk_token="<unk>" # BPE模型本身也可以定义<unk>
)
tokenizer_low_level = Tokenizer(bpe_model)
# 2. 设置范化器 (Normalizer) - 可选但推荐
# 范化器在预分词之前作用于原始文本字符串
tokenizer_low_level.normalizer = normalizers.Sequence()
# 3. 设置预分词器 (PreTokenizer)
# 对于字节级BPE，这里必须使用 ByteLevel
tokenizer_low_level.pre_tokenizer = pre_tokenizers.ByteLevel(
add_prefix_space=True, # 与方案一的推荐设置保持一致
trim_offsets=True # 移除偏移量中的前导空格
)
# 4. 设置解码器 (Decoder)
# 确保解码过程与预分词过程对应
tokenizer_low_level.decoder = decoders.ByteLevel()
# 5. 设置训练器 (Trainer)
bpe_trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
min_frequency=min_frequency,
show_progress=True,
special_tokens=special_tokens,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), #!! 关键：确保是字节级BPE的初始字母表
# continuing_subword_prefix 和 end_of_word_suffix 也可以在Trainer中设置，
# 但如果在Model中已设置，这里通常不需要重复。Trainer中的设置会覆盖Model中的。
)
# 6. 训练分词器
print("Starting BPE tokenizer training (Low-Level)...")
tokenizer_low_level.train_from_iterator(
batch_iterator(openwebtext_dataset, text_column="text", batch_size=100), # 同样使用迭代器
trainer=bpe_trainer,
length=len(openwebtext_dataset) # 提供样本总数
)
print("BPE tokenizer training completed.")
# 7. 保存分词器
# 这种方式通常直接保存为一个 tokenizer.json 文件，包含完整配置
tokenizer_save_path_low_level = "./my_bpe_tokenizer_low_level"
os.makedirs(tokenizer_save_path_low_level, exist_ok=True)
tokenizer_low_level.save(os.path.join(tokenizer_save_path_low_level, "tokenizer.json"))
print(f"Low-level tokenizer saved to {tokenizer_save_path_low_level}/tokenizer.json")
优点 (Pros):
- 对分词流程的每一个组件（范化、预分词、模型、解码、训练）都有完全的、精细的控制权。
- 可以灵活设置更多BPE模型本身的参数，如dropout。
缺点 (Cons):
- 代码相对复杂，需要对 tokenizers 库的各个组件有更深入的理解。
- 配置项较多，更容易出错。
1.4 小结与推荐
为了更清晰地对比这两种方案，请看下表：
| 特性 (Feature) | 方案一 (ByteLevelBPETokenizer) | 方案二 (底层组件构建) |
|---|---|---|
| 实现复杂度 | 低 | 高 |
| 控制力 | 中 | 高 |
| 推荐度 | ★★★★★ (初学者首选) | ★★★ (适合进阶用户或特殊需求) |
| 关键参数 | vocab_size, min_frequency, special_tokens, add_prefix_space, unicode_normalizer, trim_offsets | 方案一所有参数 + dropout, continuing_subword_prefix (在 models.BPE 中), initial_alphabet (在 BpeTrainer 中) |
建议： 对于刚开始接触大模型和分词器训练的同学，强烈推荐从方案一 (ByteLevelBPETokenizer)入手。它在易用性和功能性之间取得了极佳的平衡，能让你快速得到一个高质量的字节级BPE分词器。当你对整个流程非常熟悉，并且有特殊定制需求时，再考虑方案二也不迟。
关于词表大小 (vocab_size) 的快速提示：
选择一个合适的 vocab_size 非常重要，它直接影响模型的性能和效率。
- 太小：常见的词汇可能会被拆分成过多的子词，导致输入序列变长，增加计算负担，也可能损失一些词的整体语义。
- 太大：会增加模型输入层（Embedding层）的参数数量，可能导致模型体积增大，训练变慢。同时，过大的词表可能包含很多非常低频、甚至只出现过一两次的“词元”，这些词元的Embedding可能学习不充分。
- 权衡：需要在压缩率（即平均每个词被编码成的token数量）和词表覆盖度之间找到平衡。
- 实践参考：GPT-2的词表大小是50257，LLaMA是32000，Gemma是256000。对于主要是英文的OpenWebText，45k到65k通常是一个合理的起始范围。记住，
ByteLevelBPETokenizer的初始词表已经包含了256个字节，所以你设置的 vocab_size 是指最终（字节 + 通过BPE合并学习到的新词元）的总大小。 保存你的“字典”——至关重要的一步！
训练好分词器后，务必妥善保存。
- 使用
tokenizer.save_model("output_dir")会在指定目录下保存vocab.json(词表)和merges.txt(合并规则) 两个文件。 - 更推荐的方式是使用
tokenizer.save("output_dir/tokenizer.json")。这个方法会保存一个单独的tokenizer.json文件，它里面包含了分词器的所有配置信息：模型类型、范化规则、预分词规则、后处理规则、词表、合并规则等等。这个完整的tokenizer.json文件可以非常方便地被transformers.PreTrainedTokenizerFast类加载，实现与Hugging Face Transformers生态的无缝对接。这对于后续的模型训练和推理至关重要，能确保分词行为的一致性。
2. 长文本处理，告别简单截断
清洗后的OpenWebText数据集，每一行文本可能都非常长。在训练decoder-only类型的模型（如GPT系列）时，模型通常有一个最大的序列长度限制(例如 512, 1024, 2048 tokens)。如果简单地将超长的文本行直接截断，会丢失大量有价值的上下文信息，这对于依赖长程依赖关系的语言模型来说是十分不利的。
那么，如何更优雅地处理这些“大长篇”呢？
我可以通过滑动窗口分块来处理“长篇”的文章和内容。这是一种常见的策略：将一篇长文档切分成多个固定长度的文本块 (chunks)，并且允许这些文本块之间有一定的重叠 (overlap)。
工作流程：
- 取一篇长文档。
- 设定目标块长度
chunk_size(通常等于模型的max_seq_length)和重叠长度overlap_size。 - 从文档开头开始，取一个长度为
chunk_size的文本块。 - 下一个文本块从上一个块的
chunk_size - overlap_size位置开始，同样取chunk_size长度。 - 重复此过程，直到覆盖整个文档。
长文档: [-------------------------------------------------------------]
块 1: [===============] (长度 chunk_size)
块 2: [===============] (与块1有 overlap_size 的重叠)
块 3: [===============] (与块2有 overlap_size 的重叠)
...
优点：
- 最大化数据利用：一篇长文档可以生成多个训练样本。
- 上下文保持：重叠部分有助于模型在一定程度上学习到块与块之间的连续性信息，缓解了硬截断带来的上下文割裂问题。
- 实现简单：逻辑清晰，易于在数据处理流程中实现。
缺点：
- 数据冗余：重叠部分的数据在多个样本中重复出现。
- 上下文有限：尽管有重叠，但每个块的上下文仍然局限于
chunk_size。对于需要超长程依赖的任务，可能仍有局限。 - 数据集膨胀：生成的样本数量会显著增加，需要更多的存储和训练时间。
建议：滑动窗口分块是一个非常实用且广泛采用的策略。overlap_size的选择需要权衡，常见的取值可以是chunk_size的10%-25%。例如，如果chunk_size是512，overlap_size`可以设为50或128。
我们将在下一节的 Dataset 实现中融入这个策略。
3. 定制PyTorch Dataset
有了处理好的Arrow数据和训练好的 BPE 分词器，接下来就是构建数据加载管道，将数据高效地送入模型进行训练。考虑到OpenWebText数据集非常庞大，我们必须采用流式加载 (streaming) 的方式，避免将整个数据集一次性读入内存。同时，要平衡内存使用和数据加载到显存进行训练的效率。
PyTorch 提供了torch.utils.data.IterableDataset类，非常适合这种流式处理大数据的场景。我们将创建一个自定义的IterableDataset，它能够：
- 从 Arrow 文件中逐条或逐批读取原始文本。
- 对每条文本应用之前讨论的滑动窗口分块策略。
- 对每个文本块使用我们训练好的BPE分词器进行分词。
- 产出模型训练所需的input_ids和attention_mask。
Hugging Face datasets库也提供了将磁盘上的Arrow数据集转换为IterableDataset的便捷方法to_iterable_dataset()，或者在加载时直接使用streaming=True参数。这种方式可以大大简化实现，因为它底层已经处理了高效的文件读取和迭代。
参考代码如下：
class StreamingOpenWebTextDataset(IterableDataset):
def __init__(self, hf_dataset: Dataset, tokenizer: PreTrainedTokenizerFast, max_seq_length: int,
text_column: str = 'text', overlap_size: int = 128,
precompute_total_samples: bool = False,
num_proc_for_counting: int = None,
total_sample: int = None):
super().__init__()
self.hf_dataset = hf_dataset
# 使用PreTrainedTokenizerFast加载tokenizer
self.tokenizer = tokenizer
# 确保tokenizer有pad_token
if self.tokenizer.pad_token is None:
raise ValueError("tokenizer pad token is none")
self.max_seq_length = max_seq_length
self.text_column = text_column
if overlap_size >= max_seq_length:
raise ValueError("overlap_size 必须小于 max_seq_length。")
self.overlap_size = overlap_size
self.stride = self.max_seq_length - self.overlap_size
if self.stride <= 0:
# stride必须为正，否则滑动窗口无法前进
raise ValueError(f"Stride ({self.stride}) 必须为正。请检查 max_seq_length ({self.max_seq_length}) 和 overlap_size ({self.overlap_size})。")
if total_sample is None:
self._total_samples = None # 用于存储计算出的总样本数
else:
self._total_samples = total_sample
if precompute_total_samples:
print("Precomputing total number of samples. This may take a while for large datasets...")
if num_proc_for_counting is None:
num_proc_for_counting = os.cpu_count()
self._compute_total_samples(num_proc_for_counting=num_proc_for_counting)
print(f"Total number of samples precomputed: {self._total_samples}")
def _tokenize_and_chunk_document_stream(self, document_text: str):
"""
辅助函数：对单个文档进行分词和分块，并生成块。
"""
if not document_text or not isinstance(document_text, str):
return
encoded_document = self.tokenizer(
document_text,
add_special_tokens=False,
return_attention_mask=False,
truncation=False,
padding=False
)
doc_token_ids = encoded_document['input_ids']
if not doc_token_ids:
return
max_covered_exclusive_end_idx_in_doc = 0
for i in range(0, len(doc_token_ids), self.stride):
chunk_token_ids = doc_token_ids[i : i + self.max_seq_length]
if not chunk_token_ids:
continue
current_chunk_actual_exclusive_end_idx = i + len(chunk_token_ids)
if current_chunk_actual_exclusive_end_idx <= max_covered_exclusive_end_idx_in_doc:
continue
yield {"input_ids": chunk_token_ids}
max_covered_exclusive_end_idx_in_doc = current_chunk_actual_exclusive_end_idx
def _count_chunks_in_document_text(self, document_text: str) -> int:
"""
计算单个文档文本可以产生的 chunk 数量。
这个方法不 yield，而是直接返回计数。
"""
count = 0
if not document_text or not isinstance(document_text, str):
return 0
# 注意：这里的 tokenizer 调用与 _tokenize_and_chunk_document_stream 中的一致
encoded_document = self.tokenizer(
document_text,
add_special_tokens=False,
return_attention_mask=False,
truncation=False,
padding=False
)
doc_token_ids = encoded_document['input_ids']
if not doc_token_ids:
return 0
max_covered_exclusive_end_idx_in_doc = 0
for i in range(0, len(doc_token_ids), self.stride):
chunk_token_ids = doc_token_ids[i : i + self.max_seq_length]
if not chunk_token_ids:
continue
current_chunk_actual_exclusive_end_idx = i + len(chunk_token_ids)
if current_chunk_actual_exclusive_end_idx <= max_covered_exclusive_end_idx_in_doc:
continue
count += 1
max_covered_exclusive_end_idx_in_doc = current_chunk_actual_exclusive_end_idx
return count
def _map_function_to_count_chunks(self, example: dict) -> dict:
"""
这是传递给 hf_dataset.map() 的函数。
它从 example 中提取文本，并使用 _count_chunks_in_document_text 计算 chunk 数。
"""
document_text = example[self.text_column]
num_chunks = self._count_chunks_in_document_text(document_text)
return {"num_chunks_for_doc": num_chunks} # 返回包含新列的字典
def _map_function_to_count_chunks_batched(self, examples: dict) -> dict:
"""
这是传递给 hf_dataset.map(batched=True) 的函数。
它从 examples 中提取一批文本，并计算每个文本的 chunk 数。
"""
document_texts = examples[self.text_column]
num_chunks_list = []
for text in document_texts:
num_chunks_list.append(self._count_chunks_in_document_text(text))
return {"num_chunks_for_doc": num_chunks_list}
def _compute_total_samples(self, num_proc_for_counting: int = None, use_batched_map: bool = True, map_batch_size: int = 1000):
"""
使用 dataset.map() 并行计算总样本数。
"""
if self._total_samples is not None:
return self._total_samples
if num_proc_for_counting is None:
num_proc_for_counting = os.cpu_count()
if num_proc_for_counting is None: # os.cpu_count() might return None
num_proc_for_counting = 1
print(f"Number of processors for counting not specified, defaulting to {num_proc_for_counting}.")
if num_proc_for_counting > os.cpu_count():
print(f"Warning: num_proc_for_counting ({num_proc_for_counting}) is greater than available CPUs ({os.cpu_count()}). Setting to {os.cpu_count()}.")
num_proc_for_counting = os.cpu_count()
print(f"Calculating total samples using dataset.map() with num_proc={num_proc_for_counting}...")
# `datasets.map()` 需要一个函数，该函数接收一个 example (dict) 并返回一个 dict。
# `self` (包含 tokenizer, stride 等) 会被 pickle 并传递给子进程。
# PreTrainedTokenizerFast 通常是可 pickle 的。
# 选择 map 函数 (batched or not)
map_fn = self._map_function_to_count_chunks_batched if use_batched_map else self._map_function_to_count_chunks
# 执行 map 操作
# 这个操作会返回一个新的 Dataset 对象（或者如果原始数据集是内存映射的，可能会就地修改）
# 它会包含一个名为 "num_chunks_for_doc" 的新列。
dataset_with_counts = self.hf_dataset.map(
map_fn,
batched=use_batched_map,
batch_size=map_batch_size if use_batched_map else None,
num_proc=num_proc_for_counting,
desc="Counting chunks per document",
remove_columns=[col for col in self.hf_dataset.column_names if col != self.text_column and col != "num_chunks_for_doc"] # 移除不必要的列以节省内存
)
# 从新列中计算总和
# dataset_with_counts["num_chunks_for_doc"] 会是一个包含每个文档chunk数的列表或Arrow Array
total_count = sum(dataset_with_counts["num_chunks_for_doc"])
self._total_samples = total_count
return self._total_samples
def __len__(self):
"""
返回数据集中的样本总数。
如果未预计算，则会触发计算或抛出异常。
DataLoader需要这个方法来知道总共有多少数据，以便正确显示进度。
"""
if self._total_samples is None:
raise RuntimeError("Total number of samples has not been precomputed. "
"Initialize dataset with precompute_total_samples=True "
"or call _compute_total_samples() manually.")
return self._total_samples
def __iter__(self):
worker_info = torch.utils.data.get_worker_info()
if worker_info is None:
doc_indices_for_this_worker = range(len(self.hf_dataset))
else:
num_docs = len(self.hf_dataset)
num_workers = worker_info.num_workers
worker_id = worker_info.id
doc_indices_for_this_worker = range(worker_id, num_docs, num_workers)
for doc_idx in doc_indices_for_this_worker:
try:
document_text = self.hf_dataset[doc_idx][self.text_column]
except IndexError:
print(f"Warning: Worker {worker_info.id if worker_info else 'main'} tried to access out-of-bounds doc_idx {doc_idx}.")
continue
for sample in self._tokenize_and_chunk_document_stream(document_text):
yield sample
自定义 collate_fn：DataLoader会从Dataset中取出样本并组合成批次 (batch)。对于文本数据，由于每个样本（即文本块）的长度可能不同（特别是最后一个块，或者如果max_seq_length只是一个上限），我们需要一个collate_fn来将同一批次内的样本填充 (pad) 到该批次中最长样本的长度。
collate_f参考代码如下：
def custom_collate_fn(batch, tokenizer, max_seq_length):
"""
自定义的collate_fn，使用tokenizer的pad功能
"""
# 提取input_ids列表
input_ids_list = [item['input_ids'] for item in batch]
# 使用tokenizer的padding功能
padded_batch = tokenizer.pad(
{'input_ids': input_ids_list},
padding='longest', # 填充到batch中最长序列的长度
max_length=max_seq_length, # 但不超过最大长度
return_tensors='pt', # 返回PyTorch张量
return_attention_mask=True
)
return {
"input_ids": padded_batch['input_ids'],
"attention_mask": padded_batch['attention_mask']
}
4. 总结
- BPE分词器训练：学习了如何使用Hugging Face tokenizers从Arrow数据集高效训练BPE分词器。
- 长文本处理：掌握了滑动窗口分块策略，优雅地处理长文本，最大限度保留上下文信息。
- 流式数据加载：构建了基于
IterableDataset的流式Dataset，实现了内存友好且高效的数据投喂。
这些都是大模型训练中数据准备环节的核心技术点。虽然看起来步骤繁多，但理解了背后的原理并动手实践后，我们会发现它们并没有那么神秘。
数据准备是LLM项目中非常耗时但又极其重要的一环。希望这篇攻略能为我们后续的实际打下坚实的基础。这些环节准备完毕，接下来，我们就可以更有信心地去探索模型选择、训练技巧、模型评估等更广阔的天地啦！
关注我的公众号不走丢
附录
GitHub链接：github.com/JimmysAIPG/…
