prepare_data.py

"""
数据预处理脚本
将OpenWebText数据集转换为token ids并保存为二进制文件
"""
import os
import argparse
import numpy as np
from tqdm import tqdm
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor
from transformers import GPT2Tokenizer


def tokenize_file(args):
    """对单个文件进行tokenize"""
    file_path, tokenizer_name = args
    tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_name)

    try:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            text = f.read()
        if text.strip():
            tokens = tokenizer.encode(text)
            return tokens
    except Exception as e:
        print(f"Error processing {file_path}: {e}")
    return []


def process_openwebtext(
    data_dir: str,
    output_dir: str,
    tokenizer_name: str = "gpt2",
    val_ratio: float = 0.005,
    test_ratio: float = 0.005,
    num_workers: int = 8
):
    """
    处理OpenWebText数据集

    Args:
        data_dir: 解压后的OpenWebText目录
        output_dir: 输出目录
        tokenizer_name: tokenizer名称
        val_ratio: 验证集比例
        test_ratio: 测试集比例（用于计算PPL）
        num_workers: 并行处理的worker数量
    """
    os.makedirs(output_dir, exist_ok=True)

    # 收集所有txt文件
    print("Collecting files...")
    all_files = []

    # OpenWebText解压后可能有多层目录结构
    for root, dirs, files in os.walk(data_dir):
        for fname in files:
            if fname.endswith('.txt'):
                all_files.append(os.path.join(root, fname))

    print(f"Found {len(all_files)} text files")

    if len(all_files) == 0:
        print("No .txt files found. Checking for .xz files...")
        xz_files = list(Path(data_dir).rglob("*.xz"))
        if xz_files:
            print(f"Found {len(xz_files)} .xz files. Please decompress them first:")
            print("  cd {} && find . -name '*.xz' -exec xz -d {{}} \\;".format(data_dir))
        return

    # 划分训练集、验证集和测试集
    np.random.seed(42)
    np.random.shuffle(all_files)

    test_size = int(len(all_files) * test_ratio)
    val_size = int(len(all_files) * val_ratio)

    test_files = all_files[:test_size]
    val_files = all_files[test_size:test_size + val_size]
    train_files = all_files[test_size + val_size:]

    print(f"Train files: {len(train_files)}, Val files: {len(val_files)}, Test files: {len(test_files)}")

    # 处理函数
    def process_split(files, output_file, split_name):
        print(f"\nProcessing {split_name} split...")

        all_tokens = []
        args_list = [(f, tokenizer_name) for f in files]

        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            for tokens in tqdm(
                executor.map(tokenize_file, args_list),
                total=len(files),
                desc=f"Tokenizing {split_name}"
            ):
                if tokens:
                    all_tokens.extend(tokens)

        # 保存为二进制文件
        print(f"Total tokens in {split_name}: {len(all_tokens):,}")
        arr = np.array(all_tokens, dtype=np.uint16)
        arr.tofile(output_file)
        print(f"Saved to {output_file}")

        return len(all_tokens)

    # 处理训练集、验证集和测试集
    train_tokens = process_split(
        train_files,
        os.path.join(output_dir, "train.bin"),
        "train"
    )
    val_tokens = process_split(
        val_files,
        os.path.join(output_dir, "val.bin"),
        "val"
    )
    test_tokens = process_split(
        test_files,
        os.path.join(output_dir, "test.bin"),
        "test"
    )

    # 保存元信息
    meta = {
        "tokenizer": tokenizer_name,
        "train_tokens": train_tokens,
        "val_tokens": val_tokens,
        "test_tokens": test_tokens,
        "train_files": len(train_files),
        "val_files": len(val_files),
        "test_files": len(test_files),
    }

    import json
    with open(os.path.join(output_dir, "meta.json"), 'w') as f:
        json.dump(meta, f, indent=2)

    print("\nDone!")
    print(f"Train tokens: {train_tokens:,}")
    print(f"Val tokens: {val_tokens:,}")
    print(f"Test tokens: {test_tokens:,}")


def process_huggingface_openwebtext(
    output_dir: str,
    tokenizer_name: str = "gpt2",
    val_ratio: float = 0.005,
    test_ratio: float = 0.005,
    num_proc: int = 8
):
    """
    使用Hugging Face datasets库处理OpenWebText
    这种方式更简单，会自动下载数据
    """
    from datasets import load_dataset

    os.makedirs(output_dir, exist_ok=True)

    print("Loading OpenWebText from Hugging Face...")
    dataset = load_dataset("openwebtext")

    # 加载tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_name)
    # 我们会自行按block_size切块，这里不需要长度警告
    tokenizer.model_max_length = int(1e30)

    def tokenize_function(examples):
        return {"tokens": [tokenizer.encode(text) for text in examples["text"]]}

    print("Tokenizing...")
    tokenized = dataset.map(
        tokenize_function,
        batched=True,
        num_proc=num_proc,
        remove_columns=["text"],
        desc="Tokenizing"
    )

    # 划分并保存：单次顺序遍历原始数据集，按随机标签分流到train/val/test
    # 避免train_test_split产生的索引映射导致随机I/O
    n = len(tokenized["train"])
    print(f"Total documents: {n:,}")

    # 生成每个文档的split分配：0=train, 1=val, 2=test
    rng = np.random.RandomState(42)
    assignments = np.zeros(n, dtype=np.uint8)
    perm = rng.permutation(n)
    test_size = int(n * test_ratio)
    val_size = int(n * val_ratio)
    assignments[perm[:test_size]] = 2
    assignments[perm[test_size:test_size + val_size]] = 1

    split_arrays = {"train": [], "val": [], "test": []}
    split_map = {0: "train", 1: "val", 2: "test"}

    full_data = tokenized["train"]
    batch_size = 50000
    for i in tqdm(range(0, n, batch_size), desc="Collecting tokens"):
        end = min(i + batch_size, n)
        batch_tokens = full_data[i:end]["tokens"]
        batch_assignments = assignments[i:end]
        for j, tokens in enumerate(batch_tokens):
            split_name = split_map[batch_assignments[j]]
            split_arrays[split_name].append(np.array(tokens, dtype=np.uint16))

    for split_name, arrays in split_arrays.items():
        arr = np.concatenate(arrays)
        output_file = os.path.join(output_dir, f"{split_name}.bin")
        arr.tofile(output_file)
        print(f"Saved {split_name}: {len(arr):,} tokens to {output_file}")

    print("Done!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Preprocess OpenWebText for GPT-2 training")
    parser.add_argument(
        "--source",
        type=str,
        choices=["local", "huggingface"],
        default="huggingface",
        help="Data source: 'local' for downloaded files, 'huggingface' for automatic download"
    )
    parser.add_argument(
        "--data_dir",
        type=str,
        default="./data/openwebtext",
        help="Directory containing OpenWebText txt files (for local source)"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="./data",
        help="Output directory for processed data"
    )
    parser.add_argument(
        "--tokenizer",
        type=str,
        default="gpt2",
        help="Tokenizer to use"
    )
    parser.add_argument(
        "--val_ratio",
        type=float,
        default=0.01,
        help="Validation set ratio"
    )
    parser.add_argument(
        "--test_ratio",
        type=float,
        default=0.005,
        help="Test set ratio (for PPL calculation)"
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=8,
        help="Number of parallel workers"
    )

    args = parser.parse_args()

    if args.source == "huggingface":
        process_huggingface_openwebtext(
            output_dir=args.output_dir,
            tokenizer_name=args.tokenizer,
            val_ratio=args.val_ratio,
            test_ratio=args.test_ratio,
            num_proc=args.num_workers
        )
    else:
        process_openwebtext(
            data_dir=args.data_dir,
            output_dir=args.output_dir,
            tokenizer_name=args.tokenizer,
            val_ratio=args.val_ratio,
            test_ratio=args.test_ratio,
            num_workers=args.num_workers
        )