【大模型】- 数据管道：从文本到训练

from datasets import load_dataset

dataset = load_dataset("text", data_files="data.txt")
# 或使用预制数据集
dataset = load_dataset("wikitext", "wikitext-2-raw-v1")

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")

def tokenize_function(examples):
    return tokenizer(examples["text"])

def create_windows(examples, block_size=128):
    # 连接所有文本
    concatenated = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated[list(examples.keys())[0]])
    
    # 按 block_size 分割
    result = {
        k: [t[i:i + block_size] 
            for i in range(0, total_length, block_size)]
        for k, t in concatenated.items()
    }
    return result

def group_texts(examples):
    # 每个块：前 n-1 个 token 作为输入，后 n-1 个 token 作为目标
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    
    # 调整到 block_size 的倍数
    total_length = (total_length // block_size) * block_size
    
    result = {}
    for k, t in concatenated_examples.items():
        result[k] = [
            t[i:i + block_size]  # 输入
            for i in range(0, total_length, block_size)
        ]
    
    # 目标是输入右移一位
    result["labels"] = [
        t[1:i + block_size + 1]
        for i in range(0, total_length, block_size)
    ]
    
    return result

import torch
from torch.utils.data import DataLoader

# 应用分词和分组
tokenized = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"]
)

grouped = tokenized.map(
    group_texts,
    batched=True
)

# 转换为 PyTorch 格式
grouped.set_format(type="torch", columns=["input_ids", "attention_mask"])

# 创建 DataLoader
dataloader = DataLoader(grouped["train"], batch_size=32, shuffle=True)

from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    input_ids = [item["input_ids"] for item in batch]
    labels = [item["labels"] for item in batch]
    
    # 填充到批次内最大长度
    input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=0)
    labels_padded = pad_sequence(labels, batch_first=True, padding_value=-100)
    
    # 创建注意力掩码（忽略填充值）
    attention_mask = (input_ids_padded != 0).long()
    
    return {
        "input_ids": input_ids_padded,
        "attention_mask": attention_mask,
        "labels": labels_padded
    }

dataloader = DataLoader(
    grouped["train"],
    batch_size=32,
    collate_fn=collate_fn
)

# 使用内存映射（Hugging Face 默认）
dataset = load_dataset("text", data_files="large.txt", streaming=True)

# 或使用 WebDataset 格式
import webdataset as wds

dataset = (
    wds.WebDataset("data-{000..099}.tar")
    .decode("pil")
    .to_tuple("input_ids", "labels")
    .batched(32)
)

def analyze_dataset(dataset):
    total_tokens = sum(len(x["input_ids"]) for x in dataset)
    avg_length = total_tokens / len(dataset)
    vocab_size = len(set(token for x in dataset for token in x["input_ids"]))
    
    print(f"总样本数: {len(dataset):,}")
    print(f"总 token 数: {total_tokens:,}")
    print(f"平均序列长度: {avg_length:.1f}")
    print(f"词汇表大小: {vocab_size:,}")
    print(f"估计大小: {total_tokens * 2 / 1e9:.2f} GB (float32)")

# 多进程加载
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,  # 多进程
    pin_memory=True,  # GPU 加速
    persistent_workers=True
)

# 预取
prefetch_factor=2  # 每个工作进程预取 2 个批次