#!/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()