Direction9/database/download_direction9_datasets.py

#!/usr/bin/env python
"""
方向 9 实验数据集下载脚本

根据「方向 9——跨模型谱收敛实验计划」下载所需数据集：
- MATH: 数学推理
- GSM8K: 多步算术
- flores-200: 低资源语言
- BIG-Bench Hard: 能力涌现探针
- HumanEval: 代码生成
- AlpacaEval: 指令跟随

用法:
    python database/download_direction9_datasets.py --all
    python database/download_direction9_datasets.py --dataset math --sample-num 500
"""

import os
import json
import argparse
from pathlib import Path
from typing import List, Dict, Optional

# 设置 HuggingFace 镜像
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from datasets import load_dataset


def save_to_jsonl(texts: List[Dict], domain: str, output_dir: str = "database/corpus") -> None:
    """保存语料到 JSONL 文件"""
    output_path = Path(output_dir) / domain / "texts.jsonl"
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w", encoding="utf-8") as f:
        for i, item in enumerate(texts):
            record = {
                "text": item.get("text", ""),
                "metadata": {
                    "index": i,
                    "domain": domain,
                    **{k: v for k, v in item.items() if k != "text"}
                }
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    print(f"  已保存至：{output_path}")


def load_math(sample_num: int = 500) -> List[Dict]:
    """
    加载数学推理数据集 - 使用 GSM8K + MetaMathQA 作为 MATH 替代
    MATH 数据集已不可用，使用替代数据源
    """
    corpus = []
    print("加载数学推理数据集 (GSM8K + MetaMathQA)...")

    # 加载 GSM8K (已在 gsm8k 数据集函数中处理，这里作为补充)
    try:
        ds = load_dataset("gsm8k", "main", split="train")
        count = 0
        for item in ds:
            if count >= sample_num // 2:
                break
            count += 1

            question = item.get("question", "")
            answer = item.get("answer", "")

            if question:
                corpus.append({
                    "text": f"Question: {question}\n\nAnswer: {answer}",
                    "type": "arithmetic",
                    "source": "gsm8k"
                })
    except Exception as e:
        print(f"  GSM8K 加载失败：{e}")

    # 加载 MetaMathQA 作为补充
    try:
        ds = load_dataset("MetaMathQA/MetaMathQA", split="train")
        count = 0
        for item in ds:
            if count >= sample_num // 2:
                break
            count += 1

            query = item.get("query", "")
            response = item.get("response", "")

            if query:
                corpus.append({
                    "text": f"Question: {query}\n\nAnswer: {response}",
                    "type": "math_reasoning",
                    "source": "MetaMathQA"
                })
    except Exception as e:
        print(f"  MetaMathQA 加载失败：{e}")

    print(f"  成功加载 {len(corpus)} 条数学题目")
    return corpus


def load_gsm8k(sample_num: int = 500) -> List[Dict]:
    """
    加载 GSM8K 数据集 - 多步算术推理
    https://huggingface.co/datasets/gsm8k
    """
    corpus = []
    print("加载 GSM8K 数据集...")

    try:
        ds = load_dataset("gsm8k", "main", split="test")

        for item in ds:
            if len(corpus) >= sample_num:
                break

            question = item.get("question", "")
            answer = item.get("answer", "")

            if question:
                corpus.append({
                    "text": f"Question: {question}\n\nAnswer: {answer}",
                    "type": "arithmetic"
                })

        print(f"  成功加载 {len(corpus)} 条算术题目")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_flores200(sample_num: int = 500) -> List[Dict]:
    """
    加载低资源语言数据集
    使用替代源：WMT 或 NLLB 数据集
    """
    corpus = []
    print("加载低资源语言数据集 (WMT/NLLB)...")

    # 尝试加载 NLLB 数据集 (包含多种低资源语言)
    try:
        ds = load_dataset("facebook/flores", split="dev")
        count = 0
        for item in ds:
            if count >= sample_num:
                break
            count += 1

            # 获取斯瓦希里语文本
            text = item.get("sentence_swa", "") or item.get("sentence", "")
            if text:
                corpus.append({
                    "text": text,
                    "language": "swa",
                    "source": "flores"
                })

        print(f"  成功加载 {len(corpus)} 条低资源语言文本")
    except Exception as e:
        print(f"  flores 加载失败：{e}")
        print("  尝试使用 WMT 新闻翻译数据...")

        # 降级方案：使用 WMT 数据
        try:
            ds = load_dataset("wmt16", "ro-en", split="train")
            for item in ds:
                if len(corpus) >= sample_num:
                    break

                text = item.get("translation", {}).get("ro", "")
                if text:
                    corpus.append({
                        "text": text,
                        "language": "ro",
                        "source": "wmt16"
                    })
            print(f"  [降级] 成功加载 {len(corpus)} 条文本")
        except Exception as e2:
            print(f"  降级方案也失败：{e2}")

    return corpus


def load_bigbench_hard(sample_num: int = 500) -> List[Dict]:
    """
    加载 BIG-Bench Hard 数据集 - 能力涌现探针
    使用替代源：CoT 推理数据集或 GSM8K
    """
    corpus = []
    print("加载推理数据集 (CoT reasoning / GSM8K)...")

    try:
        # 尝试加载 CoT 数据集
        ds = load_dataset("lama-lab/big-bench-hard", split="train")

        count = 0
        for item in ds:
            if count >= sample_num:
                break
            count += 1

            question = item.get("question", "")
            answer = item.get("answer", "")
            task = item.get("task", "")

            if question:
                corpus.append({
                    "text": f"Task: {task}\n\nQuestion: {question}\n\nAnswer: {answer}",
                    "task": task,
                    "type": "reasoning"
                })

        print(f"  成功加载 {len(corpus)} 条推理题目")
    except Exception as e:
        print(f"  加载失败：{e}")
        print("  尝试使用 GSM8K 作为推理探针...")

        # 降级方案：使用 GSM8K 作为推理探针
        try:
            ds = load_dataset("gsm8k", "main", split="test")
            for item in ds:
                if len(corpus) >= sample_num:
                    break

                question = item.get("question", "")
                answer = item.get("answer", "")

                if question:
                    corpus.append({
                        "text": f"Task: arithmetic_reasoning\nQuestion: {question}\n\nAnswer: {answer}",
                        "task": "arithmetic_reasoning",
                        "type": "reasoning"
                    })
            print(f"  [降级] 成功加载 {len(corpus)} 条推理题目")
        except Exception as e2:
            print(f"  降级方案也失败：{e2}")

    return corpus


def load_humaneval(sample_num: int = 200) -> List[Dict]:
    """
    加载 HumanEval 数据集 - 代码生成
    https://huggingface.co/datasets/openai_humaneval
    """
    corpus = []
    print("加载 HumanEval 数据集...")

    try:
        ds = load_dataset("openai_humaneval", split="test")

        for item in ds:
            if len(corpus) >= sample_num:
                break

            prompt = item.get("prompt", "")
            canonical_solution = item.get("canonical_solution", "")
            entry_point = item.get("entry_point", "")
            test = item.get("test", "")

            if prompt:
                corpus.append({
                    "text": f"Description: {prompt}\n\nSolution: {canonical_solution}\n\nTests: {test}",
                    "entry_point": entry_point,
                    "type": "code_generation"
                })

        print(f"  成功加载 {len(corpus)} 条代码生成题目")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


def load_alpaca_eval(sample_num: int = 200) -> List[Dict]:
    """
    加载 AlpacaEval 数据集 - 指令跟随
    使用替代源：Alpaca cleaned 数据集
    """
    corpus = []
    print("加载指令跟随数据集 (Alpaca)...")

    try:
        # 使用 Alpaca 数据集
        ds = load_dataset("yahma/alpaca-cleaned", split="train")

        count = 0
        for item in ds:
            if count >= sample_num:
                break
            count += 1

            instruction = item.get("instruction", "")
            input_text = item.get("input", "")
            output_text = item.get("output", "")

            if instruction:
                full_text = f"Instruction: {instruction}\n\nInput: {input_text}\n\nOutput: {output_text}"
                corpus.append({
                    "text": full_text,
                    "type": "instruction_following",
                    "source": "alpaca-cleaned"
                })

        print(f"  成功加载 {len(corpus)} 条指令样本")
    except Exception as e:
        print(f"  加载失败：{e}")

    return corpus


# 数据集加载器映射
DATASET_LOADERS = {
    "math": load_math,
    "gsm8k": load_gsm8k,
    "flores200": load_flores200,
    "bigbench": load_bigbench_hard,
    "humaneval": load_humaneval,
    "alpaca": load_alpaca_eval,
}


def main():
    parser = argparse.ArgumentParser(
        description="下载方向 9 实验所需数据集",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 下载所有数据集
  python database/download_direction9_datasets.py --all

  # 下载数学推理数据集
  python database/download_direction9_datasets.py --dataset math --sample-num 500

  # 下载低资源语言数据集
  python database/download_direction9_datasets.py --dataset flores200 --sample-num 1000

可用数据集:
  - math: MATH 数学推理
  - gsm8k: GSM8K 多步算术
  - flores200: flores-200 低资源语言
  - bigbench: BIG-Bench Hard 推理
  - humaneval: HumanEval 代码生成
  - alpaca: AlpacaEval 指令跟随
        """
    )

    parser.add_argument(
        "--dataset",
        type=str,
        choices=list(DATASET_LOADERS.keys()),
        help="下载指定数据集"
    )

    parser.add_argument(
        "--all",
        action="store_true",
        help="下载所有数据集"
    )

    parser.add_argument(
        "--sample-num",
        type=int,
        default=500,
        help="每个数据集的采样数量 (默认：500)"
    )

    parser.add_argument(
        "--output-dir",
        type=str,
        default="database/corpus",
        help="输出目录 (默认：database/corpus)"
    )

    args = parser.parse_args()

    print("=" * 60)
    print("方向 9 实验数据集下载")
    print("=" * 60)

    if args.all:
        datasets = list(DATASET_LOADERS.keys())
    elif args.dataset:
        datasets = [args.dataset]
    else:
        print("错误：请指定 --dataset 或 --all")
        return

    for dataset_name in datasets:
        print(f"\n[{dataset_name}]")
        loader = DATASET_LOADERS[dataset_name]

        # 根据数据集类型调整样本数
        if dataset_name in ["humaneval", "alpaca"]:
            sample_num = min(args.sample_num, 200)
        else:
            sample_num = args.sample_num

        texts = loader(sample_num)

        if texts:
            save_to_jsonl(texts, dataset_name, args.output_dir)
        else:
            print(f"  警告：未能获取有效语料")

    print("\n" + "=" * 60)
    print("下载完成!")
    print("=" * 60)


if __name__ == "__main__":
    main()