generate_data.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
一键抓取所有 HuggingFace 数据集的 PDF 直链，按分类+数量分片输出多个小 JSON 文件。
"""
import json, urllib.request, urllib.error, urllib.parse, time, os, sys, io, math

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')

# 数据集配置
DATASETS = [
    {"id": "long_meishi",    "name": "长篇·美式",     "genre": "normal",  "length": "long",  "layout": "meishi",  "repo": "quejing/20000-meishi-pdf",    "subPath": None},
    {"id": "long_yingshi",   "name": "长篇·英式",     "genre": "normal",  "length": "long",  "layout": "yingshi", "repo": "quejing/20000-yingshi-pdf",   "subPath": None},
    {"id": "short_meishi",   "name": "短篇·美式",     "genre": "normal",  "length": "short", "layout": "meishi",  "repo": "quejing/small-meishi-pdf",    "subPath": None},
    {"id": "short_yingshi",  "name": "短篇·英式",     "genre": "normal",  "length": "short", "layout": "yingshi", "repo": "quejing/small-yingshi-pdf",   "subPath": None},
    {"id": "zashu_meishi",   "name": "杂书名著·美式", "genre": "zashu",    "length": "long",  "layout": "meishi",  "repo": "quejing/zashuPDF",          "subPath": "旧书美式PDF"},
    {"id": "zashu_yingshi",  "name": "杂书名著·英式", "genre": "zashu",    "length": "long",  "layout": "yingshi", "repo": "quejing/zashuPDF",          "subPath": "旧书英式PDF"},
    {"id": "zhongwen_meishi","name": "古典中文·美式", "genre": "zhongwen", "length": "long",  "layout": "meishi",  "repo": "quejing/zhongwenxiaoshuo",   "subPath": "古典美式PDF"},
    {"id": "zhongwen_yingshi","name": "古典中文·英式","genre": "zhongwen", "length": "long",  "layout": "yingshi", "repo": "quejing/zhongwenxiaoshuo",   "subPath": "古典英式PDF"},
]

BASE_URL = "https://huggingface.co"
API_BASE = "https://huggingface.co/api"
CHUNK_SIZE = 5000  # 每个分片最多5000条

def fetch_all_files(repo):
    url = f"{API_BASE}/datasets/{repo}?full=true"
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
    resp = urllib.request.urlopen(req, timeout=120)
    data = json.loads(resp.read().decode("utf-8"))
    return data.get("siblings", [])

def build_pdf_list(siblings, sub_path, ds_info):
    results = []
    for sib in siblings:
        fname = sib.get("rfilename", "")
        if not fname.lower().endswith(".pdf"):
            continue
        if sub_path and not fname.startswith(sub_path + "/"):
            continue
        base_name = fname.rsplit("/", 1)[-1]
        # 拼接直链 — 对路径中特殊字符做 URL 编码，防止 &、空格等引起问题
        encoded_fname = urllib.parse.quote(fname, safe='/%')
        direct_url = f"{BASE_URL}/datasets/{ds_info['repo']}/resolve/main/{encoded_fname}"
        results.append([base_name, direct_url])  # 紧凑格式: [文件名, 下载链接]
    return results

def main():
    out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
    os.makedirs(out_dir, exist_ok=True)

    # 索引文件：列出所有数据源及其分片信息
    index = {"sources": []}

    for idx, ds in enumerate(DATASETS):
        print(f"[{idx+1}/{len(DATASETS)}] 正在获取: {ds['name']} ({ds['repo']})", flush=True)
        try:
            siblings = fetch_all_files(ds["repo"])
            pdfs = build_pdf_list(siblings, ds["subPath"], ds)
            total = len(pdfs)
            chunks = math.ceil(total / CHUNK_SIZE) if total > 0 else 1
            print(f"  -> 找到 {total} 个 PDF，分 {chunks} 个文件", flush=True)

            # 写分片文件
            chunk_files = []
            for ci in range(chunks):
                chunk_pdfs = pdfs[ci * CHUNK_SIZE : (ci + 1) * CHUNK_SIZE]
                chunk_name = f"{ds['id']}_{ci}.json"
                chunk_path = os.path.join(out_dir, chunk_name)
                # 紧凑JSON，不缩进，不用ensure_ascii
                with open(chunk_path, "w", encoding="utf-8") as f:
                    json.dump(chunk_pdfs, f, ensure_ascii=False, separators=(',', ':'))
                fsize = os.path.getsize(chunk_path)
                print(f"    -> {chunk_name}: {len(chunk_pdfs)} 条, {fsize:,} bytes", flush=True)
                chunk_files.append({"f": chunk_name, "n": len(chunk_pdfs)})

            index["sources"].append({
                "id": ds["id"],
                "name": ds["name"],
                "genre": ds["genre"],
                "length": ds["length"],
                "layout": ds["layout"],
                "repo": ds["repo"],
                "total": total,
                "chunks": chunk_files,
            })

        except Exception as e:
            print(f"  !! 错误: {e}", flush=True)
            index["sources"].append({
                "id": ds["id"],
                "name": ds["name"],
                "genre": ds["genre"],
                "length": ds["length"],
                "layout": ds["layout"],
                "repo": ds["repo"],
                "total": 0,
                "chunks": [],
                "error": str(e),
            })

        time.sleep(0.3)

    # 写索引文件
    idx_path = os.path.join(out_dir, "index.json")
    with open(idx_path, "w", encoding="utf-8") as f:
        json.dump(index, f, ensure_ascii=False, separators=(',', ':'))

    # 汇总
    total = sum(s["total"] for s in index["sources"])
    total_chunks = sum(len(s["chunks"]) for s in index["sources"])
    print(f"\n{'='*60}", flush=True)
    print(f"完成! {len(index['sources'])} 数据源, {total} PDF, {total_chunks} 个分片文件", flush=True)
    for s in index["sources"]:
        print(f"  {s['name']}: {s['total']} 个 ({len(s['chunks'])} 分片)", flush=True)

    # 清理旧的单文件
    old = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pdf_data.json")
    if os.path.exists(old):
        os.remove(old)
        print(f"\n已删除旧文件: pdf_data.json", flush=True)

    print(f"\n数据目录: {out_dir}", flush=True)

if __name__ == "__main__":
    main()