data_preprocess.py

"""
数据预处理脚本 - Alpaca-style 指令微调数据集处理

功能：
1. 下载开源指令跟随数据集 (yahma/alpaca-cleaned)
2. 数据清洗：去重、过滤空输出/过短样本
3. 按 Alpaca prompt template 拼接指令
4. 划分训练集/测试集
5. 保存处理后的 JSON 文件供训练使用
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
from pathlib import Path
from typing import Any

import yaml

# 国内网络环境优先使用 HuggingFace 镜像
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

from datasets import load_dataset


# Alpaca 风格指令模板（与 Stanford Alpaca 一致）
ALPACA_PROMPT_TEMPLATE = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Input:
{input}

### Response:
{output}"""

ALPACA_PROMPT_NO_INPUT = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Response:
{output}"""


def load_config(config_path: str) -> dict[str, Any]:
    """加载 YAML 训练配置文件。"""
    with open(config_path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def build_prompt(instruction: str, input_text: str, output: str) -> str:
    """
    将单条样本拼接为 Alpaca-style 完整训练文本。
    有 input 时使用带 Input 字段的模板，否则使用无 Input 模板。
    """
    instruction = (instruction or "").strip()
    input_text = (input_text or "").strip()
    output = (output or "").strip()

    if input_text:
        return ALPACA_PROMPT_TEMPLATE.format(
            instruction=instruction,
            input=input_text,
            output=output,
        )
    return ALPACA_PROMPT_NO_INPUT.format(
        instruction=instruction,
        output=output,
    )


def build_inference_prompt(instruction: str, input_text: str) -> str:
    """构建推理时使用的 prompt（不含 Response 部分）。"""
    instruction = (instruction or "").strip()
    input_text = (input_text or "").strip()

    if input_text:
        return (
            "Below is an instruction that describes a task, paired with an input that provides further context. "
            "Write a response that appropriately completes the request.\n\n"
            f"### Instruction:\n{instruction}\n\n"
            f"### Input:\n{input_text}\n\n"
            "### Response:\n"
        )
    return (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        f"### Instruction:\n{instruction}\n\n"
        "### Response:\n"
    )


def sample_fingerprint(instruction: str, input_text: str, output: str) -> str:
    """生成样本指纹，用于去重。"""
    raw = f"{instruction}|||{input_text}|||{output}".strip().lower()
    return hashlib.md5(raw.encode("utf-8")).hexdigest()


def clean_dataset(records: list[dict[str, str]]) -> list[dict[str, str]]:
    """
    数据清洗：
    - 去除 instruction/output 为空的样本
    - 过滤过短输出（< 5 字符，通常为低质量）
    - 按内容指纹去重
    """
    cleaned: list[dict[str, str]] = []
    seen: set[str] = set()

    for row in records:
        instruction = (row.get("instruction") or "").strip()
        input_text = (row.get("input") or "").strip()
        output = (row.get("output") or "").strip()

        if not instruction or not output:
            continue
        if len(output) < 5:
            continue

        fp = sample_fingerprint(instruction, input_text, output)
        if fp in seen:
            continue
        seen.add(fp)

        cleaned.append(
            {
                "instruction": instruction,
                "input": input_text,
                "output": output,
            }
        )

    return cleaned


def split_dataset(
    records: list[dict[str, str]],
    train_ratio: float,
    seed: int,
) -> tuple[list[dict[str, str]], list[dict[str, str]]]:
    """按比例随机划分训练集与测试集。"""
    rng = random.Random(seed)
    shuffled = records.copy()
    rng.shuffle(shuffled)

    split_idx = int(len(shuffled) * train_ratio)
    train_set = shuffled[:split_idx]
    test_set = shuffled[split_idx:]
    return train_set, test_set


def format_for_training(records: list[dict[str, str]]) -> list[dict[str, str]]:
    """为训练阶段生成 text 字段。"""
    formatted = []
    for row in records:
        text = build_prompt(row["instruction"], row["input"], row["output"])
        formatted.append(
            {
                **row,
                "text": text,
                "prompt": build_inference_prompt(row["instruction"], row["input"]),
            }
        )
    return formatted


def save_json(data: list[dict[str, Any]], path: Path) -> None:
    """保存 JSON 文件。"""
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def main() -> None:
    parser = argparse.ArgumentParser(description="Alpaca-style 指令数据集预处理")
    parser.add_argument(
        "--config",
        type=str,
        default="configs/train_config.yaml",
        help="配置文件路径",
    )
    args = parser.parse_args()

    cfg = load_config(args.config)
    data_cfg = cfg["data"]
    seed = cfg["project"]["seed"]

    print(f"[1/5] 加载数据集: {data_cfg['dataset_name']}")
    dataset = load_dataset(data_cfg["dataset_name"], split="train")
    max_samples = int(data_cfg["max_samples"])
    if max_samples > 0:
        dataset = dataset.select(range(min(max_samples, len(dataset))))

    raw_records = [
        {
            "instruction": row.get("instruction", ""),
            "input": row.get("input", ""),
            "output": row.get("output", ""),
        }
        for row in dataset
    ]
    print(f"  原始样本数: {len(raw_records)}")

    print("[2/5] 数据清洗（去重、过滤低质量样本）")
    cleaned = clean_dataset(raw_records)
    print(f"  清洗后样本数: {len(cleaned)}")

    print("[3/5] 划分训练集/测试集")
    train_records, test_records = split_dataset(
        cleaned,
        train_ratio=float(data_cfg["train_ratio"]),
        seed=seed,
    )
    print(f"  训练集: {len(train_records)} | 测试集: {len(test_records)}")

    print("[4/5] 拼接 Alpaca prompt template")
    train_formatted = format_for_training(train_records)
    test_formatted = format_for_training(test_records)

    output_dir = Path(data_cfg["output_dir"])
    print(f"[5/5] 保存到 {output_dir}")
    save_json(train_formatted, output_dir / "train.json")
    save_json(test_formatted, output_dir / "test.json")
    save_json(
        {
            "dataset_name": data_cfg["dataset_name"],
            "max_samples": max_samples,
            "train_size": len(train_formatted),
            "test_size": len(test_formatted),
            "prompt_template": "Alpaca-style",
        },
        output_dir / "dataset_info.json",
    )

    print("数据预处理完成。")


if __name__ == "__main__":
    main()