构建超7GB大文本的transformers数据集时规避内存错误并保证输出一致

你之前修改的代码存在两个核心偏差，导致输出和原逻辑不一致：

• 错误1：将文本行拆分为词列表后传入tokenizer
  原逻辑的样本粒度是1整行文本对应1个编码样本，tokenizer收到一个文本字符串作为单个输入时，会对整行完成加特殊token、截断处理，输出1条input_ids。你手动line.split()拆成单词列表后，tokenizer会把每个单词当成独立的输入样本，自然会输出每个单词单独编码的短序列，完全偏离原样本粒度。
• 错误2：循环中覆盖样本列表而非累积结果
  原逻辑是把所有行的编码结果统一存入self.examples，你在循环里每次执行self.examples = ...赋值，会直接覆盖上一次循环存的内容，最终列表里只剩最后一行的编码结果，和你看到的重复输出最后一行的现象完全对应。
• 额外对齐问题：逐行读取时没有还原原逻辑的行格式
  原逻辑用f.read().splitlines()读取文件时，会自动去掉每行末尾的换行符\n；直接逐行遍历文件句柄时，读到的line会带末尾的换行符，如果不处理会导致输入文本和原逻辑有差异，编码结果出现偏差。

要在流式逐行读大文件、不触发内存错误的前提下，得到和原逻辑完全一致的编码结果，只需要按以下规则调整代码：

1. 保持和原逻辑完全一致的行过滤规则、行格式处理（去掉行尾换行符）
2. 不拆分原文本行，保持整行作为单个样本传入tokenizer
3. 逐行/分批累积编码结果到样本列表，不做覆盖
4. 控制每批传入tokenizer的文本行数，把内存占用控制在极低水平

最小改动1:1对齐版代码 #

这个版本的输出和原LineByLineTextDataset的编码结果没有任何差异：

# 初始化空样本列表，逐行累积结果
self.examples = []
with open(file_path, encoding="utf-8") as f:
    for raw_line in f:
        # 去掉行尾换行符，完全对齐原逻辑splitlines()的处理
        line = raw_line.rstrip("\n")
        # 和原逻辑完全一致的空行/纯空白行过滤
        if len(line) > 0 and not line.isspace():
            # 直接传整行文本，不做split拆分，和原逻辑的单样本输入完全一致
            encoding = tokenizer(
                line,
                add_special_tokens=True,
                truncation=True,
                max_length=block_size
            )
            # 追加当前行的编码结果，不覆盖已有样本
            self.examples.append({
                "input_ids": torch.tensor(encoding["input_ids"], dtype=torch.long)
            })

大文件优化版代码 #

逐行编码的效率较低，可以通过攒小批量文本的方式提升速度，同时依然保持极低内存占用，编码结果和原逻辑100%一致：

self.examples = []
# 调整这个批次大小控制内存占用，一般设为500-2000即可
tokenize_batch_size = 1000
current_line_batch = []

with open(file_path, encoding="utf-8") as f:
    for raw_line in f:
        line = raw_line.rstrip("\n")
        if len(line) > 0 and not line.isspace():
            current_line_batch.append(line)
        
        # 攒够一批就批量编码，避免一次性读入全文件
        if len(current_line_batch) >= tokenize_batch_size:
            batch_encoding = tokenizer(
                current_line_batch,
                add_special_tokens=True,
                truncation=True,
                max_length=block_size
            )
            # 批量追加编码结果
            for input_ids in batch_encoding["input_ids"]:
                self.examples.append({
                    "input_ids": torch.tensor(input_ids, dtype=torch.long)
                })
            # 清空当前批，释放内存
            current_line_batch = []
    
    # 处理文件末尾剩余的不足一批的行
    if current_line_batch:
        batch_encoding = tokenizer(
            current_line_batch,
            add_special_tokens=True,
            truncation=True,
            max_length=block_size
        )
        for input_ids in batch_encoding["input_ids"]:
            self.examples.append({
                "input_ids": torch.tensor(input_ids, dtype=torch.long)
            })

以上两种方案，传入tokenizer的文本内容、参数设置、样本粒度都和原源码完全一致，不会出现编码结果偏差。内存占用仅和当前攒的文本批次大小相关，和原文件7GB的大小无关，完全不会触发内存错误。

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