scripts/gen_distill_full.py

#!/usr/bin/env python3
"""
NeuroFlow 全面蒸馏数据生成 — 从全学科语料提取 prompt，教师生成 completion

特点:
  - 扫描 corpus 所有目录，按学科分层采样
  - 从实际语料文件中提取文本作为 prompt
  - 每个 prompt 调用 DeepSeek API 生成教师回答
  - 覆盖: 小学(5科) + 初中(9科) + 高中(10科) + 大学(14学科+)

用法:
  export ANTHROPIC_AUTH_TOKEN=sk-...
  python3 scripts/gen_distill_full.py \
    --corpus ~/corpus --output ~/distill_data \
    --max-tokens 128 --timeout-min 240
"""

import argparse, json, os, random, sys, time, requests
from pathlib import Path

API_URL = os.environ.get('ANTHROPIC_BASE_URL',
                         'https://api.deepseek.com/anthropic') + '/v1/messages'
API_KEY = os.environ.get('ANTHROPIC_AUTH_TOKEN', '')

# ═══ 领域映射 ═══
SUBJECT_MAP = {
    '语文': 'chinese', '外语': 'english', '数学': 'math', '科学': 'science',
    '物理': 'physics', '化学': 'chemistry', '生物': 'biology',
    '地理': 'geography', '历史': 'history', '政治': 'politics',
    '哲学': 'philosophy', '经济学': 'economics', '法学': 'law',
    '教育学': 'education', '文学': 'literature', '历史学': 'history',
    '理学': 'science', '工学': 'engineering', '农学': 'agriculture',
    '医学': 'medicine', '军事学': 'military', '管理学': 'management',
    '艺术学': 'arts', '交叉学科': 'interdisciplinary',
    '计算机': 'cs', '代码': 'code', '编程': 'programming',
}


def call_api(prompt: str, max_tokens: int = 128) -> dict:
    """返回 {text, success, error}"""
    try:
        resp = requests.post(API_URL, headers={
            'Authorization': f'Bearer {API_KEY}',
            'Content-Type': 'application/json',
        }, json={
            'model': 'deepseek-chat',
            'messages': [{'role': 'user', 'content': prompt}],
            'max_tokens': max_tokens, 'temperature': 0.7,
        }, timeout=120)
        if resp.status_code == 200:
            data = resp.json()
            for c in data.get('content', []):
                if c.get('type') == 'text':
                    return {'text': c['text'], 'success': True}
            if 'choices' in data:
                return {'text': data['choices'][0]['message']['content'], 'success': True}
            return {'success': False, 'error': f'No text in: {str(data)[:200]}'}
        return {'success': False, 'error': f'HTTP {resp.status_code}: {resp.text[:200]}'}
    except Exception as e:
        return {'success': False, 'error': str(e)}


def extract_prompts_from_file(filepath: Path, max_prompts: int = 10) -> list:
    """从文件中提取文本片段作为 prompt"""
    prompts = []
    try:
        ext = filepath.suffix.lower()
        if ext == '.jsonl':
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                for line in f:
                    try:
                        rec = json.loads(line.strip())
                        for k in ('text', 'content', 'title', 'question'):
                            if k in rec and isinstance(rec[k], str) and len(rec[k]) >= 10:
                                prompts.append(rec[k][:200])
                                break
                    except (json.JSONDecodeError, KeyError):
                        pass
                    if len(prompts) >= max_prompts:
                        break
        elif ext == '.json':
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read(10 * 1024 * 1024)  # First 10MB only
            for pos in range(0, len(content) - 10, max(len(content) // (max_prompts * 2), 100)):
                snippet = content[pos:pos + 200].strip()
                if len(snippet) >= 20 and not snippet.startswith('{'):
                    prompts.append(snippet)
            prompts = random.sample(prompts, min(len(prompts), max_prompts)) if prompts else []
        elif ext in ('.txt', '.md'):
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read(1024 * 1024)
            paragraphs = [p.strip() for p in content.split('\n') if len(p.strip()) >= 20]
            random.shuffle(paragraphs)
            prompts = paragraphs[:max_prompts]
        elif ext in ('.csv', '.tsv'):
            delim = '\t' if ext == '.tsv' else ','
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                lines = f.readlines(10000)
            for line in lines:
                fields = line.strip().split(delim)
                for fld in fields:
                    if len(fld) >= 20:
                        prompts.append(fld[:200])
            random.shuffle(prompts)
            prompts = prompts[:max_prompts]
    except Exception:
        pass
    return prompts


def scan_corpus(corpus_root: str) -> dict:
    """扫描语料库，按学科分组，返回 {subject: [file_paths]}"""
    root = Path(corpus_root)
    if not root.exists():
        print(f"❌ 语料目录不存在: {corpus_root}")
        sys.exit(1)

    subjects = {}
    for filepath in root.rglob('*'):
        if not filepath.is_file():
            continue
        ext = filepath.suffix.lower()
        if ext not in ('.txt', '.json', '.jsonl', '.csv', '.tsv', '.md'):
            continue

        # 推断学科: 看路径中的目录名
        subj = '通用'
        for part in filepath.parts:
            for key in SUBJECT_MAP:
                if key in part:
                    subj = key
                    break
        # 也看文件名
        fname = filepath.stem
        for key in SUBJECT_MAP:
            if key in fname:
                subj = key
                break

        if subj not in subjects:
            subjects[subj] = []
        subjects[subj].append(filepath)

    return subjects


def main():
    parser = argparse.ArgumentParser(description='NeuroFlow 全面蒸馏数据生成')
    parser.add_argument('--corpus', default='/home/administrator/corpus', help='语料根目录')
    parser.add_argument('--output', default='/home/administrator/distill_data', help='输出目录')
    parser.add_argument('--max-tokens', type=int, default=128, help='教师生成长度')
    parser.add_argument('--samples-per-subject', type=int, default=0,
                        help='每学科样本上限(0=自动)')
    parser.add_argument('--timeout-min', type=int, default=480, help='总超时(分钟)')
    args = parser.parse_args()

    if not API_KEY:
        print("❌ 请设置环境变量 ANTHROPIC_AUTH_TOKEN")
        sys.exit(1)

    os.makedirs(args.output, exist_ok=True)

    # 0. 断点续传: 加载已完成的 prompt
    out_file = os.path.join(args.output, 'teacher_data.jsonl')
    completed = set()
    if os.path.exists(out_file):
        print("📂 加载已完成数据...")
        with open(out_file, 'r', encoding='utf-8') as f:
            for line in f:
                try:
                    rec = json.loads(line.strip())
                    completed.add((rec.get('subject', ''), rec['prompt']))
                except (json.JSONDecodeError, KeyError):
                    pass
        print(f"   已完成: {len(completed)} 样本，断点续传")

    # 1. 扫描语料
    print("🔍 扫描语料库...")
    subjects = scan_corpus(args.corpus)
    total_files = sum(len(v) for v in subjects.values())
    print(f"   {len(subjects)} 个学科, {total_files:,} 个文件")
    for name, files in sorted(subjects.items(), key=lambda x: -len(x[1])):
        print(f"     {name}: {len(files)} 文件")

    # 2. 每学科采样 prompt (固定种子确保断点可复现)
    print("\n📝 提取 prompts...")
    random.seed(42)  # 固定种子，确保断点重入提取相同prompts
    all_prompts = {}
    total_expected = 0
    for subj, files in sorted(subjects.items()):
        limit = args.samples_per_subject or max(50, len(files) // 10)
        prompts = []
        random.shuffle(files)
        for fp in files[:min(len(files), max(100, limit))]:
            extracted = extract_prompts_from_file(fp, max_prompts=5)
            prompts.extend(extracted)
            if len(prompts) >= limit:
                break
        all_prompts[subj] = prompts[:limit]
        total_expected += len(all_prompts[subj])
        print(f"   {subj}: {len(all_prompts[subj])} prompts (目标{limit})")
    print(f"   总计: {total_expected} prompts")

    # 3. 调用 API 生成教师数据
    print(f"\n🤖 调用 DeepSeek API ({len(completed)} 已跳过)...")
    t0 = time.time()
    deadline = t0 + args.timeout_min * 60
    out_file = os.path.join(args.output, 'teacher_data.jsonl')
    stats_file = os.path.join(args.output, 'stats.json')

    total = 0
    success = 0
    total_chars = 0

    skipped = 0
    with open(out_file, 'a', encoding='utf-8') as fout:  # append mode
        for subj, prompts in sorted(all_prompts.items()):
            for i, prompt in enumerate(prompts):
                # 断点续传: 跳过已完成
                if (subj, prompt) in completed:
                    skipped += 1
                    continue
                if time.time() > deadline:
                    print(f"\n⏰ 超时 ({args.timeout_min}min)，已生成 {total} 样本")
                    break

                result = call_api(prompt, args.max_tokens)
                total += 1

                if result['success'] and len(result['text']) >= 20:
                    rec = {
                        'subject': subj,
                        'prompt': prompt[:300],
                        'completion': result['text'],
                    }
                    fout.write(json.dumps(rec, ensure_ascii=False) + '\n')
                    fout.flush()
                    success += 1
                    total_chars += len(result['text'])
                else:
                    err = result.get('error', 'unknown')[:100]
                    print(f"   ❌ {subj}[{i}]: {err}")

                if total % 20 == 0:
                    elapsed = time.time() - t0
                    rate = total * 60 / elapsed if elapsed > 0 else 0
                    print(f"   [{total}] {subj}: {success} ok, {rate:.0f}/min, "
                          f"{elapsed:.0f}s elapsed")

                # Rate limit
                if total % 80 == 0:
                    time.sleep(0.5)

            if time.time() > deadline:
                break

    elapsed = time.time() - t0

    # 4. 统计
    print(f"\n{'='*60}")
    print(f"✅ 蒸馏数据生成完成!")
    print(f"{'='*60}")
    print(f"   样本数: {success}/{total}")
    print(f"   字符数: {total_chars:,}")
    print(f"   耗时: {elapsed:.0f}s ({elapsed/60:.1f}min)")
    print(f"   速率: {success*60/elapsed:.0f} samples/min")

    fees = success * args.max_tokens / 1e6 * 0.28 + total * 200 / 1e6 * 0.14
    print(f"   费用: ~${fees:.2f}")
    print(f"   输出: {out_file}")

    with open(stats_file, 'w') as f:
        json.dump({
            'samples': success, 'total_calls': total,
            'chars': total_chars, 'elapsed': elapsed,
            'subjects': {k: len(v) for k, v in all_prompts.items()},
        }, f, indent=2, ensure_ascii=False)
    print(f"   统计: {stats_file}")

    if success == 0:
        print("\n⚠️ 没有成功生成任何样本! 检查 API Key 和网络")


if __name__ == '__main__':
    main()