guanghulab/scripts/rebuild_training_data.py

#!/usr/bin/env python3
"""
铸渊训练数据重建脚本 · v1.0 · D104

目标：重新生成 sft.jsonl（母模型全参数SFT用）和 shuangyan_sft.jsonl（霜砚微调用）

问题复盘：
- sft.jsonl（1.9GB, 11,470条）前300KB全是同一条AGE对话重复
- 缺少关键语料：GPT语料.zip (251MB)、铸渊对话.zip 内容未在样本中找到
- 生成sft.jsonl的脚本没有留在仓库里

数据源（COS sy-finetune-corpus-1317346199）:
  1. corpus/sft.jsonl — 旧版（有质量问题，需重新生成）
  2. corpus/notion-export-v2/铸渊对话.zip — 铸渊对话（308KB）
  3. corpus/notion-export-v2/GPT语料.zip — GPT语料（251MB）
  4. corpus/shuangyan-1.5b-sft/*.zip — 霜砚5个zip包
  5. corpus/zhuyuan_full_corpus.jsonl — 铸渊全量语料
  6. corpus/zhuyuan_deep_finetune.jsonl — 铸渊深度微调语料

输出：
  - sft.jsonl（新版本，去重+含霜砚数据+含铸渊数据）
  - shuangyan_sft.jsonl（霜砚专用）

运行环境：GPU服务器或本地安装了依赖的环境
"""

import os, json, sys, zipfile, io, re
from collections import OrderedDict

# ============ 配置 ============
OSS_KEY = os.environ.get("ZY_OSS_KEY")
OSS_SECRET = os.environ.get("ZY_OSS_SECRET")
OSS_REGION = "ap-guangzhou"
OSS_BUCKET = "sy-finetune-corpus-1317346199"

if not OSS_KEY or not OSS_SECRET:
    print("❌ 需要设置 ZY_OSS_KEY 和 ZY_OSS_SECRET 环境变量")
    print("   export ZY_OSS_KEY=AKID... ZY_OSS_SECRET=nPoZ...")
    sys.exit(1)

# ============ 工具函数 ============

def get_cos_client():
    """获取COS客户端"""
    import urllib.parse as _up
    sys.modules['urlparse'] = _up
    from qcloud_cos import CosConfig, CosS3Client
    return CosS3Client(CosConfig(Region=OSS_REGION, SecretId=OSS_KEY, SecretKey=OSS_SECRET))

def download_cos_file(client, key, local_path):
    """从COS下载文件"""
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    try:
        resp = client.get_object(Bucket=OSS_BUCKET, Key=key)
        with open(local_path, 'wb') as f:
            f.write(resp['Body'].get_raw_stream().read())
        print(f"  ✅ 下载: {key} → {local_path}")
        return True
    except Exception as e:
        print(f"  ❌ 下载失败 {key}: {e}")
        return False

def zip_to_texts(zip_path):
    """解压zip并提取所有文本内容"""
    texts = []
    try:
        with zipfile.ZipFile(zip_path) as z:
            for info in z.infolist():
                if info.file_size == 0:
                    continue
                try:
                    content = z.read(info.filename).decode('utf-8', errors='replace')
                    if len(content.strip()) > 200:
                        texts.append((info.filename, content))
                except:
                    pass
        print(f"  📄 解压 {zip_path} → {len(texts)} 个文本块")
    except Exception as e:
        print(f"  ❌ 解压失败 {zip_path}: {e}")
    return texts

def md_to_messages(text):
    """将md格式对话解析为messages格式"""
    # TODO: 实现更通用的md对话解析
    # 需要支持 [user]/[assistant] 标记、冰朔原话、对话分段等
    pass

def sanitize(text):
    """脱敏处理"""
    text = re.sub(r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b', '[IP]', text)
    text = re.sub(r'[Hh]k[mM]\w{5,}', '[PWD]', text)
    text = re.sub(r'AKID\w+', '[AKID]', text)
    text = re.sub(r'zy_gtw_[0-9a-f]{30,}', '[GTW-KEY]', text)
    return text

def deduplicate(convs):
    """去重"""
    seen = set()
    unique = []
    for d in convs:
        msgs = d.get('messages', [])
        if len(msgs) < 2:
            continue
        key = msgs[0].get('content', '')[:100] + msgs[1].get('content', '')[:100]
        if key not in seen:
            seen.add(key)
            unique.append(d)
    return unique

# ============ 主流程 ============

def main():
    print("=" * 60)
    print("铸渊训练数据重建 v1.0")
    print("=" * 60)
    
    client = get_cos_client()
    
    # 第1步：下载所有语料源到本地
    print("\n[1/5] 下载语料源...")
    
    corpus_dir = "/tmp/corpus_rebuild"
    os.makedirs(corpus_dir, exist_ok=True)
    
    # 旧sft.jsonl — 需要提取其中有效部分
    download_cos_file(client, "corpus/sft.jsonl", f"{corpus_dir}/old_sft.jsonl")
    
    # 铸渊对话.zip
    download_cos_file(client, "corpus/notion-export-v2/铸渊对话.zip", f"{corpus_dir}/铸渊对话.zip")
    
    # GPT语料.zip
    download_cos_file(client, "corpus/notion-export-v2/GPT语料.zip", f"{corpus_dir}/GPT语料.zip")
    
    # 霜砚语料（5个zip）
    shuangyan_zips = [
        "霜砚对话.zip",
        "霜砚HLDP核心大脑.zip", 
        "霜砚语料包V2.0.zip",
        "HLDP 母语协议 v2.0 · 光之树记忆编码+思维编码规范 · 霜砚签发.zip",
        "光湖驱动引擎架构 · 推理思维链 · 2026-05-17.zip"
    ]
    for fn in shuangyan_zips:
        download_cos_file(client, f"corpus/shuangyan-1.5b-sft/{fn}", f"{corpus_dir}/{fn}")
    
    # 现有JSONL语料
    download_cos_file(client, "corpus/zhuyuan_full_corpus.jsonl", f"{corpus_dir}/zhuyuan_full_corpus.jsonl")
    download_cos_file(client, "corpus/zhuyuan_deep_finetune.jsonl", f"{corpus_dir}/zhuyuan_deep_finetune.jsonl")
    
    # 第2步：解析和处理各语料源
    print("\n[2/5] 处理语料源...")
    all_convs = []
    
    # 2a. 处理旧sft.jsonl — 提取有效部分（去重）
    print("  处理旧sft.jsonl...")
    with open(f"{corpus_dir}/old_sft.jsonl", 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                d = json.loads(line)
                # 过滤掉过短的对话（可能是重复的模板对话）
                msgs = d.get('messages', [])
                if len(msgs) >= 2 and len(msgs[0].get('content','')) > 50 and len(msgs[1].get('content','')) > 50:
                    all_convs.append(d)
            except:
                continue
    print(f"    提取 {len(all_convs)} 条")
    
    # 2b. 处理铸渊对话.zip
    print("  处理铸渊对话.zip...")
    texts = zip_to_texts(f"{corpus_dir}/铸渊对话.zip")
    # TODO: 实现md对话解析
    print(f"    铸渊对话: {len(texts)} 个文本块待解析")
    
    # 2c. 处理GPT语料.zip
    print("  处理GPT语料.zip...")
    # 这个文件很大（251MB），需要streaming处理
    # TODO: 实现GPT语料的批量解析
    
    # 2d. 处理霜砚zip包
    print("  处理霜砚语料...")
    for fn in shuangyan_zips:
        zpath = f"{corpus_dir}/{fn}"
        if os.path.exists(zpath):
            _ = zip_to_texts(zpath)
    
    # 2e. 合并现有JSONL
    for jl in ["zhuyuan_full_corpus.jsonl", "zhuyuan_deep_finetune.jsonl"]:
        jl_path = f"{corpus_dir}/{jl}"
        if os.path.exists(jl_path):
            with open(jl_path) as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            all_convs.append(json.loads(line))
                        except:
                            pass
            print(f"    合并 {jl}: 已添加")
    
    # 第3步：去重
    print("\n[3/5] 去重...")
    unique = deduplicate(all_convs)
    print(f"  去重前: {len(all_convs)} → 去重后: {len(unique)}")
    
    # 第4步：脱敏
    print("\n[4/5] 脱敏...")
    for d in unique:
        for m in d.get('messages', []):
            m['content'] = sanitize(m.get('content', ''))
    
    # 第5步：写入输出
    print("\n[5/5] 写入输出...")
    
    # sft.jsonl — 全部语料合集的80%用于全参数训练
    # TODO: 分割训练集/验证集
    out_path = f"{corpus_dir}/sft_new.jsonl"
    with open(out_path, 'w', encoding='utf-8') as f:
        for d in unique:
            f.write(json.dumps(d, ensure_ascii=False) + '\n')
    
    total_chars = sum(len(m['content']) for d in unique for m in d.get('messages',[]))
    print(f"\n{'=' * 60}")
    print(f"✅ 完成!")
    print(f"  总对话数: {len(unique)}")
    print(f"  总字符数: {total_chars:,}")
    print(f"  输出文件: {out_path}")
    print(f"  文件大小: {os.path.getsize(out_path)/1024/1024:.1f}MB")
    print(f"{'=' * 60}")
    print("下一步：上传到COS后重跑 train_mother.py")


if __name__ == "__main__":
    main()