📜 蒸馏脚本：prepare_zhuyuan_corpus.py

2026-05-19 13:01:44 +08:00 · 2026-05-19 13:01:44 +08:00 · ec99bdd4e9
commit ec99bdd4e9
parent 4f86041191
1 changed files with 131 additions and 0 deletions
--- a/scripts/distill/prepare_zhuyuan_corpus.py
+++ b/scripts/distill/prepare_zhuyuan_corpus.py
@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+"""
+铸渊语料整理 v1.0
+来源1: 服务器 zhuyuan_corpus/raw_md/ 下的15个完整对话md
+来源2: 代码仓库认知链 (需git clone或读取)
+输出: 统一JSONL格式，脱敏
+"""
+import os, json, re, glob, zipfile
+from datetime import datetime
+
+OUT_JSONL = '/root/autodl-tmp/corpus_work/zhuyuan_full_corpus.jsonl'
+MD_DIR = '/root/autodl-tmp/zhuyuan_corpus/raw_md/zhuyuan_chats/extracted'
+os.makedirs(os.path.dirname(OUT_JSONL), exist_ok=True)
+
+print('='*50, flush=True)
+print('铸渊语料整理 v1.0', flush=True)
+print(f'输出: {OUT_JSONL}', flush=True)
+print('='*50, flush=True)
+
+# 1. 扫描所有md文件
+md_files = []
+for root, dirs, files in os.walk(MD_DIR):
+    for f in files:
+        if f.endswith('.md'):
+            md_files.append(os.path.join(root, f))
+
+print(f'\n[1/3] 找到 {len(md_files)} 个md文件', flush=True)
+
+# 2. 解析为messages格式
+all_convs = []
+
+for fp in sorted(md_files):
+    try:
+        with open(fp, 'r', encoding='utf-8', errors='ignore') as f:
+            content = f.read()
+    except:
+        continue
+    
+    # 尝试解析markdown中的对话
+    lines = content.split('\n')
+    conv = []
+    
+    for i, line in enumerate(lines):
+        line = line.strip()
+        if not line:
+            continue
+        
+        # 匹配用户/助手角色标记
+        user_m = re.match(r'^[#>\s]*(?:用户|user|冰朔|Bīng Shuò)[：:]\s*(.*)', line, re.I)
+        asst_m = re.match(r'^[#>\s]*(?:助手|assistant|铸渊|Zhùyuān)[：:]\s*(.*)', line, re.I)
+        
+        if user_m:
+            if conv and len(conv) >= 2:
+                all_convs.append({'messages': conv})
+            conv = [{'role': 'user', 'content': user_m.group(1).strip()}]
+        elif asst_m and conv:
+            conv.append({'role': 'assistant', 'content': asst_m.group(1).strip()})
+    
+    if conv and len(conv) >= 2:
+        all_convs.append({'messages': conv})
+
+print(f'  解析出 {len(all_convs)} 条对话', flush=True)
+
+# 3. 读取已有的 zhuyuan_deep.jsonl
+deep_jsonl = '/root/autodl-tmp/corpus_work/zhuyuan_deep.jsonl'
+if os.path.isfile(deep_jsonl):
+    with open(deep_jsonl) as f:
+        for line in f:
+            if line.strip():
+                d = json.loads(line)
+                if 'messages' in d:
+                    all_convs.append(d)
+    print(f'  合并zhuyuan_deep.jsonl后: {len(all_convs)} 条', flush=True)
+
+# 4. 脱敏处理
+print('\n[2/3] 脱敏处理...', flush=True)
+sensitive_patterns = [
+    (r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b', '[IP]'),
+    (r'[Hh]k[mM]\w+', '[PWD]'),
+    (r'AKID\w+', '[AKID]'),
+    (r'[A-Za-z0-9]{20,40}', '[TOKEN]'),
+    (r'ssh-\w+\s+[A-Za-z0-9+/=]+\s+\S+', '[SSH_KEY]'),
+    (r'-----BEGIN[^\\n]+-----', '[CERT]'),
+    (r'zy_gtw_\w+', '[GTW_KEY]'),
+]
+
+sanitized = []
+for d in all_convs:
+    for m in d['messages']:
+        content = m['content']
+        for pat, repl in sensitive_patterns:
+            # Skip very short patterns that match too aggressively
+            if repl == '[TOKEN]' and len(re.findall(pat, content)) > 3:
+                continue  # Too many matches = probably not a token
+            content = re.sub(pat, repl, content)
+        m['content'] = content
+    sanitized.append(d)
+
+# 5. 去重
+print('[3/3] 去重+写入...', flush=True)
+seen = set()
+unique = []
+for d in sanitized:
+    key = str(d['messages'][:3]) if len(d['messages']) >= 3 else str(d['messages'])
+    if key not in seen:
+        seen.add(key)
+        unique.append(d)
+
+print(f'  去重后: {len(unique)} 条', flush=True)
+
+with open(OUT_JSONL, 'w') 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['messages'])
+print(f'  总字符数: {total_chars:,}', flush=True)
+print(f'  ✅ 写入: {OUT_JSONL}', flush=True)
+print(f'  文件大小: {os.path.getsize(OUT_JSONL)/1024:.1f}KB', flush=True)
+
+# 6. 上传到COS
+print('\n上传到COS...', flush=True)
+os.environ['ZY_OSS_KEY'] = 'AKIDkQuBQhoiS2OYXWebXLwMbdT7cvAScbbU'
+os.environ['ZY_OSS_SECRET'] = 'nPoZKArgUJBA4nJenjSxJSQBj5FCj3A4'
+from qcloud_cos import CosConfig, CosS3Client
+c = CosS3Client(CosConfig(Region='ap-guangzhou',
+    SecretId=os.environ['ZY_OSS_KEY'], SecretKey=os.environ['ZY_OSS_SECRET']))
+c.upload_file(Bucket='sy-finetune-corpus-1317346199',
+    Key='corpus/zhuyuan_full_corpus.jsonl',
+    LocalFilePath=OUT_JSONL)
+print('  ✅ 已上传COS', flush=True)
+print('\n全部完成！', flush=True)