BERT预训练数据集构建：从数据清洗到特征工程

1. 项目概述：BERT预训练数据集的实现价值

在自然语言处理领域，预训练语言模型已经成为标配技术方案。2018年问世的BERT（Bidirectional Encoder Representations from Transformers）通过Transformer架构和掩码语言建模（MLM）任务，实现了上下文相关的词向量表示。但真正决定模型效果的，往往是训练数据的质量与处理方式。

我曾参与过多个工业级NLP项目的BERT微调工作，发现许多团队把90%的精力放在模型调参上，却忽视了数据工程这个基础环节。实际上，在相同模型架构下，优质的数据预处理能使下游任务性能提升5-15%。《动手学深度学习》第69节聚焦的正是这个常被轻视的关键环节——从原始文本到模型可消化数据的完整转换流程。

这个实现过程包含三个核心价值点：

1. 数据规范化：将维基百科、图书语料等异构文本转化为统一的预训练格式
2. 特征工程：构建适合MLM和NSP（下一句预测）任务的特征结构
3. 性能优化：处理超长文本时的内存效率方案

2. 核心需求解析

2.1 预训练数据的特殊性

与常见的分类/标注数据集不同，BERT预训练数据需要满足两个独特需求：

• 动态掩码要求：每个epoch需要重新生成随机掩码位置，避免模型记忆固定模式
• 文档级连贯性：NSP任务需要保持原始文档的段落关系，不能简单随机打乱句子

以维基百科数据为例，原始XML格式包含大量标记和注释信息。我们需要提取纯净文本的同时，保留章节、段落等结构信息。这涉及到：

python复制<page>
  <title>深度学习</title>
  <revision>
    <text>== 历史 ==\n1958年...\n== 架构 ==\n神经网络由...\n</text>
  </revision>
</page>

2.2 内存效率挑战

BERT预训练通常需要处理TB级文本。直接加载全部数据到内存显然不现实。我们的实现方案需要：

1. 流式处理：按文档逐批加载
2. 缓存机制：保存预处理中间结果
3. 并行化：利用多进程加速特征生成

3. 数据预处理全流程实现

3.1 原始文本清洗

首先安装必要的依赖库：

bash复制pip install wikiextractor html2text tqdm

使用改进版的WikiExtractor进行内容提取，保留段落分隔信息：

python复制from wikiextractor import clean, extract_pages

def process_wiki_dump(input_file):
    for page in extract_pages(input_file):
        if not page.is_redirect():
            text = clean(page.text)
            yield {
                'title': page.title,
                'paragraphs': [p for p in text.split('\n\n') if p.strip()]
            }

关键技巧：保留原始段落结构对NSP任务至关重要，清洗时需禁用常规的句子打乱操作

3.2 文档分块策略

BERT的输入长度限制（通常512 tokens）要求我们对长文档进行智能分块：

1. 基础分块：按段落自然分割
2. 重叠分块：对技术文档采用滑动窗口（stride=128）
3. 特殊处理：代码片段等非自然语言内容单独标记

实现代码示例：

python复制def chunk_document(paragraphs, max_len=512, stride=128):
    chunks = []
    current_chunk = []
    current_len = 0
    
    for para in paragraphs:
        tokens = tokenizer.tokenize(para)
        if current_len + len(tokens) > max_len:
            if current_chunk:
                chunks.append(' '.join(current_chunk))
                current_chunk = current_chunk[-stride:]
                current_len = len(current_chunk)
        current_chunk.extend(tokens)
        current_len += len(tokens)
    
    return chunks

4. 特征工程实现细节

4.1 动态掩码生成

不同于静态预处理方案，我们实现动态掩码以增强模型鲁棒性：

python复制def create_masked_lm_predictions(tokens, mask_prob=0.15):
    cand_indices = [i for i, token in enumerate(tokens) 
                   if token not in ['[CLS]', '[SEP]']]
    random.shuffle(cand_indices)
    
    output_tokens = list(tokens)
    masked_lm_positions = []
    masked_lm_labels = []
    
    for idx in cand_indices[:int(len(cand_indices)*mask_prob)]:
        masked_token = None
        # 80%概率替换为[MASK]
        if random.random() < 0.8:
            masked_token = "[MASK]"
        else:
            # 10%概率保留原词
            if random.random() < 0.5:
                masked_token = tokens[idx]
            # 10%概率替换为随机词
            else:
                masked_token = random.choice(vocab_list)
        
        output_tokens[idx] = masked_token
        masked_lm_positions.append(idx)
        masked_lm_labels.append(tokens[idx])
    
    return output_tokens, masked_lm_positions, masked_lm_labels

4.2 下一句预测(NSP)样本生成

正样本：连续的两个文档片段
负样本：随机组合的不相关片段

python复制def create_next_sentence_example(chunks, index):
    current_chunk = chunks[index]
    next_index = index + 1
    
    # 50%概率生成负样本
    if random.random() < 0.5 or next_index >= len(chunks):
        random_index = random.randint(0, len(chunks)-1)
        while random_index == index:
            random_index = random.randint(0, len(chunks)-1)
        return (current_chunk, chunks[random_index], 0)  # 0表示负样本
    else:
        return (current_chunk, chunks[next_index], 1)  # 1表示正样本

5. 性能优化实战技巧

5.1 内存映射技术

对于超大规模数据集，使用内存映射文件避免OOM：

python复制import mmap

class MemoryMappedDataset:
    def __init__(self, file_path):
        self.file = open(file_path, 'r+b')
        self.mm = mmap.mmap(self.file.fileno(), 0)
        
    def __getitem__(self, index):
        # 实现索引逻辑
        pass

5.2 多进程预处理

利用Python的multiprocessing加速：

python复制from multiprocessing import Pool

def parallel_preprocess(texts, workers=8):
    with Pool(workers) as p:
        results = list(tqdm(p.imap(process_function, texts), total=len(texts)))
    return results

实测数据：在32核服务器上处理100GB维基百科数据，多进程方案比单线程快17倍

6. 常见问题与解决方案

6.1 中文特殊处理

中文需要不同的分词策略：

python复制# 使用jieba分词替代默认的空白分割
import jieba

def chinese_tokenizer(text):
    return [word for word in jieba.cut(text) if word.strip()]

6.2 长文档处理技巧

对于法律、医疗等领域的超长文档：

1. 关键句提取（基于TF-IDF）
2. 层次化分块（先按章节，再按段落）
3. 增加特殊分隔符标记不同层级

6.3 数据平衡策略

当使用多源数据（维基+新闻+论坛）时：

1. 按领域分层采样
2. 动态调整采样频率
3. 添加领域标识符token

7. 完整实现示例

以下是从原始文本到训练数据的端到端流程：

python复制# 1. 加载原始数据
documents = load_wiki_dump("enwiki-latest-pages-articles.xml.bz2")

# 2. 分块处理
chunks = []
for doc in documents:
    chunks.extend(chunk_document(doc['paragraphs']))

# 3. 生成训练样本
train_examples = []
for i in range(len(chunks)):
    tokens = tokenizer.tokenize(chunks[i])
    masked_tokens, positions, labels = create_masked_lm_predictions(tokens)
    next_sent = create_next_sentence_example(chunks, i)
    
    train_examples.append({
        'input_ids': tokenizer.convert_tokens_to_ids(masked_tokens),
        'masked_lm_positions': positions,
        'masked_lm_labels': tokenizer.convert_tokens_to_ids(labels),
        'next_sentence_label': next_sent[2]
    })

# 4. 保存TFRecord格式
writer = tf.io.TFRecordWriter("bert_pretrain_data.tfrecord")
for example in train_examples:
    tf_example = create_tf_example(example)
    writer.write(tf_example.SerializeToString())
writer.close()

在实际项目中，我发现三个关键点最容易出问题：

1. 文本编码不一致（特别是混合多语言数据时）
2. 动态掩码的随机种子管理
3. 内存泄漏问题（尤其在多进程环境下）

建议在正式训练前，先用小规模数据（如1%）跑通全流程验证数据质量。我曾经因为跳过这个步骤，导致浪费了3天的GPU计算资源——最终发现是数据预处理环节的一个小bug导致所有标签错位。