484 lines
20 KiB
Python
484 lines
20 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
小狸 AI - 预训练数据生成器 (v5)
|
||
==================================
|
||
两阶段生成：
|
||
|
||
Phase 1 - 种子直出:
|
||
5并发 API 直接生成完整多轮对话
|
||
每API调用 = 1条高质量对话 (15-25轮, 2000-4000字)
|
||
速率: ~600条/h, 17h出1万条种子
|
||
|
||
Phase 2 - 语义扩展:
|
||
基于种子主题，API重新"编剧"出新对话
|
||
每条种子 → 10个变体（换角度/深度/角色）
|
||
速率: ~600条/h, ~7天扩展至10万+条
|
||
|
||
数据格式: JSONL, 每条: {"text":"...", "domain":"...", "concept":"..."}
|
||
"""
|
||
import sys, os, json, time, random, hashlib, argparse, re
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
from api_correction import call_astron
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 1. 角色定义
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
TEACHER_SYSTEM = "你是一位AI导师，回答专业详细，包含具体数据和原理。每轮回答150-500字。"
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 2. 深度学习领域图谱
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
DOMAINS = {
|
||
"量子物理": ["量子叠加态","量子纠缠","量子隧穿","海森堡测不准原理","波粒二象性","量子计算","量子场论","狄拉克方程"],
|
||
"相对论": ["狭义相对论","质能方程","时空弯曲","引力波","黑洞物理","虫洞理论","暗物质","暗能量"],
|
||
"人工智能": ["深度学习原理","Transformer","强化学习","大语言模型","扩散模型","多模态学习","注意力机制","模型量化"],
|
||
"计算机科学": ["计算理论","操作系统","分布式系统","数据库","编译器","密码学","网络协议","计算机视觉"],
|
||
"分子生物学": ["DNA复制","CRISPR","中心法则","蛋白质折叠","表观遗传学","RNA干扰","免疫系统","合成生物学"],
|
||
"数学": ["微积分","傅里叶分析","微分方程","抽象代数","数论","拓扑学","概率论","图论"],
|
||
"宇宙学": ["大爆炸","恒星演化","星系形成","宇宙微波背景","暗物质粒子","暴胀模型","引力波天文学","哈勃常数危机"],
|
||
"神经科学": ["神经元","突触可塑性","大脑皮层","记忆机制","睡眠与意识","视觉系统","脑机接口","情感神经科学"],
|
||
"化学": ["化学键","反应动力学","热力学","催化原理","光谱学","纳米化学","电化学","高分子化学"],
|
||
"工程科技": ["航天工程","半导体","通信工程","机器人学","电力系统","材料工程","机械设计","核工程"],
|
||
"经济学": ["一般均衡","博弈论","信息不对称","行为经济学","货币理论","计量经济学","拍卖理论","机制设计"],
|
||
"哲学": ["先验哲学","权力意志","存在主义","语言游戏","证伪主义","范式理论","正义论","功利主义"],
|
||
"天文学": ["赫罗图","超新星","中子星","系外行星","星系分类","脉冲星","恒星形成","宇宙射线"],
|
||
"地球科学": ["板块构造","地震学","火山学","古生物学","气候系统","海洋环流","大气动力学","行星地质"],
|
||
"生物技术": ["基因测序","PCR","冷冻电镜","核磁共振","光遗传学","类器官","纳米医学","基因敲除"],
|
||
}
|
||
|
||
ALL_DOMAINS = list(DOMAINS.keys())
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 3. Phase 1: API直接生成种子
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
def build_seed_prompt(domain: str, concept: str) -> str:
|
||
return f"""请围绕【{domain}】领域的【{concept}】概念，生成一段师生深度对话。
|
||
|
||
要求：
|
||
1. 学生：从基础概念问起，逐步深入，包含追问、质疑、要求举例
|
||
2. 导师：每轮回答150-500字，专业准确，有数据支撑
|
||
3. 对话演进：基础定义→核心机制→实验证据→应用场景→前沿进展
|
||
4. 至少15轮对话
|
||
|
||
每轮格式（严格）：
|
||
学生：[问题]
|
||
导师：[回答]
|
||
|
||
不要额外说明，直接开始对话。"""
|
||
|
||
|
||
def generate_seed(domain: str, concept: str) -> dict:
|
||
"""一次API调用生成一条种子对话"""
|
||
prompt = build_seed_prompt(domain, concept)
|
||
text = call_astron(prompt, temperature=random.uniform(0.7, 0.9),
|
||
system_prompt=TEACHER_SYSTEM, max_tokens=8192)
|
||
if not text:
|
||
return None
|
||
|
||
# 解析
|
||
turns = []
|
||
for line in text.strip().split('\n'):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
if re.match(r'^学生[：:]', line):
|
||
content = re.sub(r'^学生[：:]\s*', '', line)
|
||
turns.append(('user', content))
|
||
elif re.match(r'^导师[：:]', line):
|
||
content = re.sub(r'^导师[：:]\s*', '', line)
|
||
turns.append(('assistant', content))
|
||
|
||
if len(turns) < 10:
|
||
return None
|
||
|
||
lines = []
|
||
for role, content in turns:
|
||
lines.append(f"<|im_start|>{role}\n{content}<|im_end|>")
|
||
|
||
return {
|
||
"text": '\n'.join(lines),
|
||
"domain": domain,
|
||
"concept": concept,
|
||
}
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 4. Phase 2: 语义扩展
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
VARIATION_PROMPTS = [
|
||
"用更通俗的方式重新讲一遍，多用日常类比",
|
||
"从更高级的视角讲，假设学生已有基础，重点讲深层机制和数学原理",
|
||
"换个角度，从历史发展脉络讲，突出关键实验和突破性发现",
|
||
"从工程应用角度讲，重点讲实际怎么用、有什么技术瓶颈",
|
||
"从批判性视角讲，有哪些争议、局限性、未解决的问题",
|
||
"假设学生是高中水平，用最通俗的语言和大量类比来解释",
|
||
"从跨学科视角讲，这个概念与其他领域有什么联系",
|
||
"以问答竞赛的方式讲，学生提尖锐问题，导师用证据回击",
|
||
]
|
||
|
||
|
||
def generate_variation(seed: dict, seed_text: str, prompt_template: str) -> dict:
|
||
"""基于种子主题生成一个语义变体"""
|
||
domain = seed["domain"]
|
||
concept = seed["concept"]
|
||
instruction = prompt_template
|
||
|
||
prompt = f"""请围绕【{domain}】领域的【{concept}】概念，生成一段师生深度对话。
|
||
|
||
{instruction}
|
||
|
||
要求：
|
||
1. 至少15轮对话
|
||
2. 导师回答150-500字，专业详细
|
||
3. 对话要有递进：定义→机制→证据→应用→前沿
|
||
|
||
每轮格式（严格）：
|
||
学生：[问题]
|
||
导师：[回答]
|
||
|
||
直接开始对话，不要额外说明。"""
|
||
|
||
text = call_astron(prompt, temperature=random.uniform(0.75, 0.95),
|
||
system_prompt=TEACHER_SYSTEM, max_tokens=8192)
|
||
if not text:
|
||
return None
|
||
|
||
turns = []
|
||
for line in text.strip().split('\n'):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
if re.match(r'^学生[：:]', line):
|
||
content = re.sub(r'^学生[：:]\s*', '', line)
|
||
turns.append(('user', content))
|
||
elif re.match(r'^导师[：:]', line):
|
||
content = re.sub(r'^导师[：:]\s*', '', line)
|
||
turns.append(('assistant', content))
|
||
|
||
if len(turns) < 10:
|
||
return None
|
||
|
||
lines = []
|
||
for role, content in turns:
|
||
lines.append(f"<|im_start|>{role}\n{content}<|im_end|>")
|
||
|
||
return {
|
||
"text": '\n'.join(lines),
|
||
"domain": domain,
|
||
"concept": concept,
|
||
"variation": prompt_template[:30],
|
||
}
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 5. 数据管理
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
class DataStore:
|
||
def __init__(self, save_dir: str):
|
||
self.save_dir = Path(save_dir)
|
||
self.save_dir.mkdir(parents=True, exist_ok=True)
|
||
self.output_path = self.save_dir / 'dialogues.jsonl'
|
||
self.seed_path = self.save_dir / 'seeds.jsonl'
|
||
|
||
self.existing_hashes = set()
|
||
self.count = 0
|
||
self.load()
|
||
|
||
def load(self):
|
||
for path in [self.output_path, self.seed_path]:
|
||
if path.exists():
|
||
with open(path, 'r') as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line:
|
||
try:
|
||
item = json.loads(line)
|
||
h = hashlib.md5(item['text'][:100].encode()).hexdigest()
|
||
self.existing_hashes.add(h)
|
||
self.count += 1
|
||
except:
|
||
pass
|
||
print(f"📂 已有: {self.count:,} 条")
|
||
|
||
def add_seeds(self, items: list) -> int:
|
||
if not items:
|
||
return 0
|
||
new = []
|
||
for item in items:
|
||
h = hashlib.md5(item['text'][:100].encode()).hexdigest()
|
||
if h not in self.existing_hashes:
|
||
self.existing_hashes.add(h)
|
||
new.append(item)
|
||
if not new:
|
||
return 0
|
||
with open(self.seed_path, 'a') as f:
|
||
for item in new:
|
||
f.write(json.dumps(item, ensure_ascii=False) + '\n')
|
||
self.count += 1
|
||
return len(new)
|
||
|
||
def add_variations(self, items: list) -> int:
|
||
if not items:
|
||
return 0
|
||
new = []
|
||
for item in items:
|
||
h = hashlib.md5(item['text'][:100].encode()).hexdigest()
|
||
if h not in self.existing_hashes:
|
||
self.existing_hashes.add(h)
|
||
new.append(item)
|
||
if not new:
|
||
return 0
|
||
with open(self.output_path, 'a') as f:
|
||
for item in new:
|
||
f.write(json.dumps(item, ensure_ascii=False) + '\n')
|
||
self.count += 1
|
||
return len(new)
|
||
|
||
def get_seeds(self):
|
||
"""读取所有种子"""
|
||
seeds = []
|
||
if self.seed_path.exists():
|
||
with open(self.seed_path, 'r') as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line:
|
||
seeds.append(json.loads(line))
|
||
return seeds
|
||
|
||
def total_count(self):
|
||
return self.count
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 6. 质量评分
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
def score(text: str) -> float:
|
||
turns = (text.count('<|im_start|>user') + text.count('<|im_start|>assistant')) // 2
|
||
chars = len(text)
|
||
s = 0.0
|
||
if turns >= 20: s += 0.20
|
||
elif turns >= 12: s += 0.15
|
||
elif turns >= 8: s += 0.10
|
||
else: s += 0.05
|
||
if chars >= 4000: s += 0.20
|
||
elif chars >= 2000: s += 0.15
|
||
elif chars >= 1200: s += 0.10
|
||
else: s += 0.05
|
||
avg = chars / max(turns, 1)
|
||
if avg >= 200: s += 0.15
|
||
elif avg >= 100: s += 0.10
|
||
uc = text.count('<|im_start|>user')
|
||
ac = text.count('<|im_start|>assistant')
|
||
ratio = min(uc, ac) / max(uc, ac, 1)
|
||
if ratio >= 0.8: s += 0.10
|
||
unique = len(set(text))
|
||
if unique >= 150: s += 0.10
|
||
return max(0, min(1, s))
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 7. 卷统计
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
def print_stats(store: DataStore, t_start: float, phase: str, target: int):
|
||
elapsed = max(time.time() - t_start, 1)
|
||
rate = store.count / elapsed * 3600
|
||
eta_h = (target - store.count) / max(rate, 0.01)
|
||
|
||
seed_count = 0
|
||
if store.seed_path.exists():
|
||
with open(store.seed_path) as f:
|
||
seed_count = sum(1 for _ in f if _.strip())
|
||
|
||
tokens_est = store.count * 6000
|
||
|
||
print(f" [{time.strftime('%H:%M:%S')}] "
|
||
f"Phase {phase} | "
|
||
f"累计: {store.count:,}/{target:,} "
|
||
f"({store.count/max(target,1)*100:.1f}%) | "
|
||
f"速率: {rate:.0f}条/h | "
|
||
f"种子: {seed_count} | "
|
||
f"tokens: {tokens_est/1e9:.2f}B | "
|
||
f"预计: {eta_h/24:.1f}d", flush=True)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 8. 主函数
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--target", type=int, default=100000, help="目标条数")
|
||
parser.add_argument("--phase", type=int, default=1, choices=[1, 2],
|
||
help="1=生成种子, 2=语义扩展")
|
||
parser.add_argument("--concurrency", type=int, default=5, help="API并发数")
|
||
args = parser.parse_args()
|
||
|
||
save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'pretrain_data')
|
||
store = DataStore(save_dir)
|
||
|
||
target = args.target
|
||
if store.count >= target:
|
||
print(f"✅ 已达目标 {target:,} 条")
|
||
return
|
||
|
||
print(f"\n{'='*50}")
|
||
print(f" Phase {args.phase}: {'种子生成' if args.phase == 1 else '语义扩展'}")
|
||
print(f" 目标: {target:,} 条")
|
||
print(f" 并发: {args.concurrency}")
|
||
print(f" 当前: {store.count:,}")
|
||
print(f"{'='*50}\n")
|
||
|
||
t_start = time.time()
|
||
|
||
if args.phase == 1:
|
||
# Phase 1: 纯API种子生成
|
||
seeds = store.get_seeds()
|
||
# 从种子中的(domain, concept)对去重
|
||
used = set()
|
||
for s in seeds:
|
||
used.add((s.get("domain",""), s.get("concept","")))
|
||
print(f" 已有 {len(used)} 个不同(领域,概念)种子\n")
|
||
|
||
while store.count < target:
|
||
batch_items = []
|
||
tasks = []
|
||
for _ in range(args.concurrency):
|
||
# 选一个未用过的(领域,概念)
|
||
for attempt in range(20):
|
||
d = random.choice(ALL_DOMAINS)
|
||
c = random.choice(DOMAINS[d])
|
||
if (d, c) not in used or attempt > 15:
|
||
break
|
||
used.add((d, c))
|
||
tasks.append((d, c))
|
||
|
||
with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
|
||
futures = {ex.submit(generate_seed, d, c): (d, c) for d, c in tasks}
|
||
for f in as_completed(futures):
|
||
result = f.result()
|
||
if result and score(result["text"]) >= 0.5:
|
||
batch_items.append(result)
|
||
|
||
if batch_items:
|
||
added = store.add_seeds(batch_items)
|
||
n = len(batch_items)
|
||
print(f" +{n}条种子 | 领域: {batch_items[0]['domain']}, "
|
||
f"概念: {batch_items[0]['concept']}...", flush=True)
|
||
|
||
print_stats(store, t_start, "1", target)
|
||
|
||
if store.count >= target:
|
||
break
|
||
|
||
elif args.phase == 2:
|
||
# Phase 2: 语义扩展
|
||
seeds = store.get_seeds()
|
||
if not seeds:
|
||
print("❌ 没有种子数据，先跑 Phase 1")
|
||
return
|
||
print(f" 种子库: {len(seeds)} 条\n")
|
||
|
||
used_seeds = set()
|
||
expansion_tasks = []
|
||
|
||
while store.count < target:
|
||
if not expansion_tasks:
|
||
# 把所有种子打乱，生成扩展任务
|
||
random.shuffle(seeds)
|
||
for s in seeds:
|
||
if s["text"] in used_seeds:
|
||
continue
|
||
used_seeds.add(s["text"])
|
||
for vp in VARIATION_PROMPTS:
|
||
expansion_tasks.append((s, vp))
|
||
print(f" 生成 {len(expansion_tasks)} 个扩展任务\n")
|
||
|
||
# 取一批任务
|
||
batch = []
|
||
while len(batch) < args.concurrency and expansion_tasks:
|
||
batch.append(expansion_tasks.pop(0))
|
||
|
||
if not batch:
|
||
break
|
||
|
||
with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
|
||
futures = {ex.submit(generate_variation, s, s["text"], vp): (s, vp)
|
||
for s, vp in batch}
|
||
for f in as_completed(futures):
|
||
result = f.result()
|
||
if result and score(result["text"]) >= 0.5:
|
||
store.add_variations([result])
|
||
|
||
# 报告
|
||
seed_count = 0
|
||
if store.seed_path.exists():
|
||
with open(store.seed_path) as f:
|
||
seed_count = sum(1 for _ in f if _.strip())
|
||
elapsed = max(time.time() - t_start, 1)
|
||
rate = store.count / elapsed * 3600
|
||
eta_h = (target - store.count) / max(rate, 0.01) if rate > 0 else 0
|
||
tokens_est = store.count * 6000
|
||
done_tasks = max(1, len(used_seeds) * len(VARIATION_PROMPTS) - len(expansion_tasks))
|
||
print(f" [{time.strftime('%H:%M:%S')}] "
|
||
f"累计: {store.count:,}/{target:,} "
|
||
f"({store.count/max(target,1)*100:.1f}%) | "
|
||
f"速率: {rate:.0f}条/h | "
|
||
f"种子: {seed_count} | "
|
||
f"扩展任务进度: {done_tasks} | "
|
||
f"tokens: {tokens_est/1e9:.2f}B | "
|
||
f"预计: {eta_h/24:.1f}d", flush=True)
|
||
|
||
if store.count >= target:
|
||
break
|
||
|
||
# 最终报告
|
||
total_time = time.time() - t_start
|
||
final_rate = store.count / max(total_time, 1) * 3600
|
||
tokens = store.count * 6000
|
||
path = store.output_path if args.phase == 2 else store.seed_path
|
||
size_mb = os.path.getsize(path) / 1024 / 1024 if path.exists() else 0
|
||
|
||
print(f"\n{'='*50}")
|
||
print(f" ✅ Phase {args.phase} 完成!")
|
||
print(f" 总条数: {store.count:,}")
|
||
print(f" 文件: {path}")
|
||
print(f" 大小: {size_mb:.1f}MB")
|
||
print(f" 估算tokens: {tokens/1e9:.2f}B")
|
||
print(f" 耗时: {total_time/3600:.1f}h")
|
||
print(f" 平均速率: {final_rate:.0f}条/h")
|
||
print(f"{'='*50}")
|
||
|
||
# 样例
|
||
print("\n📊 样例:")
|
||
with open(path) as f:
|
||
for i, line in enumerate(f):
|
||
if i >= 2:
|
||
break
|
||
if line.strip():
|
||
item = json.loads(line.strip())
|
||
text = item["text"]
|
||
turns = text.count("<|im_start|>") // 2
|
||
print(f"\n 条目{i+1}: {turns}轮 | {len(text):,}字 | {item.get('domain','')}")
|
||
lines = text.split("\n")[:4]
|
||
for l in lines:
|
||
if l.strip():
|
||
role = "👤" if "user" in l else "🧑🏫"
|
||
c = l.split("<|im_start|>")[1].split("<|im_end|>")[0].strip()
|
||
print(f" {role} {c[:100]}")
|
||
print(f" ...（{turns}轮）")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |