871 lines
40 KiB
Python
871 lines
40 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
小狸 AI - 万亿参数级预训练数据生成器 (v4)
|
||
==============================================
|
||
核心架构：双AI互相对话（师生对练）+ 话题生成器
|
||
|
||
5并发分配:
|
||
1 slot → TopicGenerator（生产话题大纲）
|
||
2 slots → DialoguePair #1（学生问 ↔ 老师答）
|
||
2 slots → DialoguePair #2（学生问 ↔ 老师答）
|
||
|
||
优势:
|
||
- 无单次max_tokens限制，对话可无限延伸
|
||
- 自然真实的多轮交互（一问一答实时生成）
|
||
- 每轮回答使用完整8192上下文，内容深度极高
|
||
- 5并发全部满载，生成速率最大化
|
||
|
||
目标: 2,000,000 条超长多轮对话
|
||
数据格式（JSONL）:
|
||
{"text": "<|im_start|>user\\n...<|im_end|>\\n<|im_start|>assistant\\n...<|im_end|>\\n..."}
|
||
|
||
作者: xiaoli-llm-chat team
|
||
"""
|
||
import sys, os, json, time, random, hashlib, argparse, re, threading
|
||
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. 角色 System Prompt
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
STUDENT_SYSTEM = """你是一个求知欲极强的学习者，正在向一位博学的AI导师深入学习专业知识。
|
||
|
||
你的人设：
|
||
- 你对该领域有一定基础概念，但不深
|
||
- 从最基础的问题开始，逐步深入追问
|
||
- 善于提出有洞察力的追问和质疑
|
||
- 喜欢要求导师用类比和具体例子解释
|
||
|
||
提问策略：
|
||
1. 第一轮：概念定义——"什么是XXX？它的核心思想是什么？"
|
||
2. 前几轮：基础机制——"XXX是如何工作的？为什么？"
|
||
3. 中间轮次：深入追问——"你提到的XX机制能再详细解释吗？"
|
||
4. 进阶问题：应用与证据——"有哪些实验验证？实际应用是什么？"
|
||
5. 后期问题：前沿与开放——"目前最前沿的进展是什么？还有哪些未解之谜？"
|
||
|
||
规则要求：
|
||
- 每轮只问一个问题
|
||
- 语言自然，像真实的人类学习者
|
||
- 不要问"你能解释一下吗"这种空泛问题
|
||
- 问题中尽量包含关键词方便导师针对性回答
|
||
- 每轮问题不超过 200 字"""
|
||
|
||
TEACHER_SYSTEM = """你是一位百科全书级的AI导师，精通各学科领域的深层知识。
|
||
|
||
你的人设：
|
||
- 学识渊博、逻辑严密
|
||
- 善于用通俗语言解释复杂概念
|
||
- 回答结构完整：定义→原理→证据→延伸
|
||
- 对每个问题给出专业、准确、深入的解答
|
||
|
||
回答标准：
|
||
1. 每轮回答 200-800 字，内容丰富、干货满满
|
||
2. 包含具体数据、年份、人名、公式、实验名称
|
||
3. 善用类比帮助理解抽象概念
|
||
4. 从不使用套话空话（"值得注意的是"、"也就是说"等）
|
||
5. 适当回应学生问题中的具体关键词
|
||
6. 在适当时候引导提问（"这引出了另一个重要概念……"）
|
||
|
||
知识深度自适应：
|
||
- 基础问题 → 清晰定义 + 核心原理
|
||
- 进阶追问 → 机制详解 + 数学模型 + 实验证据
|
||
- 前沿问题 → 最新突破 + 开放问题 + 未来方向
|
||
|
||
规则要求：
|
||
- 每轮只回答当前问题
|
||
- 回答精准有用，不跑题
|
||
- 每轮回答不超过 2000 字"""
|
||
|
||
TOPIC_SYSTEM = """你是一个知识图谱设计师，负责为深度学习对话生成话题。
|
||
|
||
要求：
|
||
1. 每个话题包含：学科领域、核心概念、起始问题、讨论提纲
|
||
2. 话题要覆盖从基础到前沿的完整路径
|
||
3. 话题要具体、有深度、可展开多轮讨论
|
||
4. 从基础原理逐步深入到前沿进展
|
||
|
||
输出严格使用以下JSON格式（不要加额外说明）：
|
||
{
|
||
"domain": "学科领域",
|
||
"concept": "核心概念名",
|
||
"starter": "起始用户问题（约30字）",
|
||
"type": "基础/机制/应用/前沿",
|
||
"outline": ["讨论点1", "讨论点2", "讨论点3", ...]
|
||
}"""
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 2. 深度学习知识图谱（100+ 领域）
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
KNOWLEDGE_GRAPH = {
|
||
"量子物理": ["量子叠加态", "量子纠缠", "量子隧穿", "海森堡测不准原理", "波粒二象性", "量子计算", "量子场论", "狄拉克方程",
|
||
"量子霍尔效应", "路径积分", "量子光学", "超导量子比特", "量子热力学", "量子测量", "拓扑量子计算", "量子引力"],
|
||
"相对论": ["狭义相对论", "质能方程", "时空弯曲", "引力波", "黑洞物理", "虫洞理论", "宇宙学原理", "暗物质",
|
||
"暗能量", "宇宙微波背景", "暴胀宇宙学", "引力透镜", "修正引力", "白矮星与中子星", "霍金辐射", "时间本质"],
|
||
"人工智能": ["深度学习原理", "Transformer架构", "强化学习", "大语言模型", "扩散模型", "生成对抗网络", "迁移学习", "知识蒸馏",
|
||
"向量数据库", "多模态学习", "图神经网络", "可解释AI", "注意力机制进化", "分布式训练", "模型量化", "AI对齐"],
|
||
"计算机科学": ["计算理论", "操作系统", "分布式系统", "数据库系统", "编译器优化", "密码学", "网络协议", "计算机视觉",
|
||
"推荐系统", "搜索算法", "并发编程", "硬件加速", "形式化验证", "计算机图形学", "数据压缩", "量子算法"],
|
||
"分子生物学": ["DNA复制", "CRISPR基因编辑", "中心法则", "蛋白质折叠", "表观遗传学", "RNA干扰", "细胞凋亡", "信号转导",
|
||
"干细胞", "免疫系统", "微生物组", "合成生物学", "线粒体生物学", "发育生物学", "神经退行性疾病", "单细胞技术"],
|
||
"数学": ["微积分定理", "傅里叶分析", "微分方程", "抽象代数", "数论", "拓扑学", "概率论", "图论",
|
||
"线性代数", "组合数学", "泛函分析", "复分析", "计算复杂性", "博弈论", "混沌理论", "优化理论"],
|
||
"宇宙学": ["大爆炸模型", "恒星演化", "星系形成", "宇宙大尺度结构", "宇宙微波背景", "暗物质粒子", "暴胀模型", "再电离时代",
|
||
"费米悖论", "多元宇宙", "黑洞天体物理", "数值宇宙学", "哈勃常数危机", "引力波天文学", "早期宇宙热历史", "宇宙拓扑"],
|
||
"化学": ["化学键理论", "反应动力学", "热力学定律", "催化原理", "光谱学", "量子化学", "有机反应机理", "配位化学",
|
||
"纳米化学", "电化学", "高分子化学", "分析化学", "光化学", "绿色化学", "表面科学", "计算化学"],
|
||
"神经科学": ["神经元基础", "突触可塑性", "大脑皮层", "记忆机制", "睡眠与意识", "视觉系统", "脑机接口", "情感神经科学",
|
||
"语言神经科学", "神经递质系统", "神经成像", "神经发育", "疼痛机制", "决策神经科学", "神经免疫", "计算神经科学"],
|
||
"物理学(综合)": ["经典力学", "电磁学", "热力学统计", "流体力学", "量子统计", "粒子物理", "等离子体物理", "非线性动力学",
|
||
"凝聚态物理", "核物理", "光学工程", "声学", "计算物理", "软物质物理", "高能天体物理", "物理教育"],
|
||
"地球科学": ["板块构造", "地震学", "火山学", "矿物学", "古生物学", "气候系统", "海洋环流", "地貌学",
|
||
"水文地质", "大气动力学", "地球化学", "地质年代学", "沉积学", "构造地质学", "行星地质", "环境地球科学"],
|
||
"工程科技": ["航天工程", "半导体技术", "通信工程", "机器人学", "电力系统", "材料工程", "机械设计", "化学工程",
|
||
"生物工程", "土木工程", "量子器件", "运筹学", "汽车工程", "精密制造", "海洋工程", "核工程"],
|
||
"经济学": ["一般均衡理论", "博弈论", "信息不对称", "行为经济学", "货币理论", "财政政策", "比较优势", "内生增长理论",
|
||
"拍卖理论", "机制设计", "金融工程", "计量经济学", "发展经济学", "劳动经济学", "国际经济学", "制度经济学"],
|
||
"哲学": ["康德的先验哲学", "尼采的权力意志", "海德格尔的存在", "维特根斯坦语言游戏", "罗素类型论", "波普尔证伪主义",
|
||
"库恩范式理论", "罗尔斯正义论", "功利主义", "义务论", "德性伦理学", "元伦理学", "逻辑实证主义", "现象学", "存在主义", "批判理论"],
|
||
"生物技术": ["基因测序技术", "PCR与qPCR", "Western blot", "流式细胞术", "质谱分析", "核磁共振", "冷冻电镜", "X射线晶体学",
|
||
"单分子技术", "光遗传学", "化学遗传学", "基因敲除", "转基因技术", "类器官培养", "3D生物打印", "纳米医学"],
|
||
"天文学": ["赫罗图", "超新星分类", "中子星", "系外行星探测", "暗物质晕", "星系分类", "宇宙距离阶梯", "中微子天文学",
|
||
"射电天文学", "X射线天文学", "伽马射线暴", "脉冲星计时", "恒星形成", "行星形成理论", "天体生物学", "宇宙射线"],
|
||
"数据科学": ["统计学习理论", "偏差方差权衡", "交叉验证", "正则化方法", "核方法", "决策树与集成", "聚类算法", "降维技术",
|
||
"异常检测", "时间序列分析", "因果推断", "贝叶斯方法", "马尔可夫链蒙特卡洛", "实验设计", "A/B测试", "生存分析"],
|
||
"逻辑学": ["命题逻辑", "谓词逻辑", "模态逻辑", "哥德尔不完备定理", "图灵停机问题", "归纳逻辑", "非单调逻辑", "悖论研究",
|
||
"二阶逻辑", "直觉主义逻辑", "线性逻辑", "描述逻辑", "时态逻辑", "道义逻辑", "认知逻辑", "自由逻辑"],
|
||
}
|
||
|
||
ALL_DOMAINS = list(KNOWLEDGE_GRAPH.keys())
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 3. 话题生成器（Slot 1）
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
def generate_topic(max_retries=3) -> dict:
|
||
"""调用Astron生成一个结构化话题"""
|
||
domain = random.choice(ALL_DOMAINS)
|
||
concepts = KNOWLEDGE_GRAPH[domain]
|
||
focus = random.choice(concepts)
|
||
|
||
types = ["基础", "机制", "应用", "前沿"]
|
||
type_weights = [0.25, 0.35, 0.25, 0.15]
|
||
ttype = random.choices(types, weights=type_weights, k=1)[0]
|
||
|
||
prompt = f"""请在【{domain}】领域中，围绕概念【{focus}】，生成一个{ttype}级别的话题。
|
||
|
||
严格按照JSON格式输出（仅JSON，不要其他文字）：
|
||
{{
|
||
"domain": "{domain}",
|
||
"concept": "{focus}",
|
||
"starter": "起始用户问题（约30字，从概念基础问起）",
|
||
"type": "{ttype}",
|
||
"outline": ["讨论点1（具体知识点）", "讨论点2（深入机制）", "讨论点3（应用或延伸）", "讨论点4（前沿开放问题）"]
|
||
}}"""
|
||
|
||
for attempt in range(max_retries):
|
||
text = call_astron(prompt, temperature=0.85 + random.uniform(-0.1, 0.1),
|
||
system_prompt=TOPIC_SYSTEM, max_tokens=1024)
|
||
if not text:
|
||
continue
|
||
try:
|
||
# 提取JSON
|
||
json_match = re.search(r'\{.*\}', text, re.DOTALL)
|
||
if json_match:
|
||
topic = json.loads(json_match.group())
|
||
if all(k in topic for k in ['domain', 'concept', 'starter', 'outline']):
|
||
return topic
|
||
except:
|
||
pass
|
||
|
||
# 降级：返回预设话题
|
||
starter_map = {
|
||
"基础": f"请从最基础的概念开始，解释一下{focus}到底是什么？",
|
||
"机制": f"{focus}的工作机制是怎样的？能详细讲讲吗？",
|
||
"应用": f"{focus}在实际中有什么重要的应用？",
|
||
"前沿": f"关于{focus}，目前最前沿的研究进展是什么？",
|
||
}
|
||
return {
|
||
"domain": domain, "concept": focus,
|
||
"starter": starter_map.get(ttype, f"请深入讲解{focus}"),
|
||
"type": ttype,
|
||
"outline": [f"{focus}的基本定义", f"{focus}的核心机制", f"{focus}的应用", f"{focus}的前沿发展"]
|
||
}
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 4. 对话对（Slots 2-3, 4-5: 每个对 学生↔老师）
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
def build_student_prompt(topic: dict, history: list) -> str:
|
||
"""构建学生提问的 prompt"""
|
||
domain = topic['domain']
|
||
concept = topic['concept']
|
||
ttype = topic['type']
|
||
outline = topic.get('outline', [])
|
||
|
||
outline_str = "\n".join(f"{i+1}. {p}" for i, p in enumerate(outline))
|
||
|
||
if not history:
|
||
# 第一轮：用 starter
|
||
return f"""领域：{domain}
|
||
核心概念：{concept}
|
||
难度级别：{ttype}
|
||
讨论提纲：
|
||
{outline_str}
|
||
|
||
这是你第一次向导师提问。请根据以上信息，提出你的第一个问题。
|
||
问题要自然、具体、有深度，从最基础的概念开始。"""
|
||
|
||
# 后续轮次：基于历史追问
|
||
recent = history[-4:] if len(history) >= 4 else history
|
||
last_exchange = "\n".join(f"{'学生' if r == 'user' else '导师'}：{c[:200]}"
|
||
for r, c in recent[-4:])
|
||
|
||
return f"""领域：{domain}
|
||
核心概念：{concept}
|
||
难度级别：{ttype}
|
||
|
||
最近几轮对话：
|
||
{last_exchange}
|
||
|
||
请根据以上讨论，提出一个深入的追问。你的问题应该：
|
||
1. 基于导师刚刚的回答中的某个点进行追问
|
||
2. 或者要求用具体例子或类比解释
|
||
3. 或者延伸到相关但尚未讨论的概念
|
||
4. 如果是后期轮次，可以问应用、实验验证或前沿问题
|
||
|
||
只输出你的问题，不要其他内容。"""
|
||
|
||
|
||
def build_teacher_prompt(topic: dict, history: list, question: str) -> str:
|
||
"""构建老师回答的 prompt"""
|
||
domain = topic['domain']
|
||
concept = topic['concept']
|
||
|
||
# 取最近的对话历史作为上下文
|
||
recent = history[-6:] if len(history) >= 6 else history
|
||
context = "\n".join(f"{'学生' if r == 'user' else '导师'}：{c}"
|
||
for r, c in recent)
|
||
|
||
current_round = len(history) // 2 + 1
|
||
|
||
return f"""领域：{domain}
|
||
核心概念：{concept}
|
||
当前第 {current_round} 轮问答
|
||
|
||
历史对话：
|
||
{context}
|
||
|
||
学生最新的问题：{question}
|
||
|
||
请给出你的专业回答。要求：
|
||
- 200-800 字
|
||
- 直接针对学生的问题
|
||
- 包含具体知识（数据、人名、年份、公式）
|
||
- 善用类比解释
|
||
- 结构完整但不模板化"""
|
||
|
||
|
||
def student_ask(topic: dict, history: list, temperature=0.80) -> str:
|
||
"""学生API：生成一个问题"""
|
||
prompt = build_student_prompt(topic, history)
|
||
question = call_astron(prompt, temperature=temperature,
|
||
system_prompt=STUDENT_SYSTEM, max_tokens=512)
|
||
if not question:
|
||
# 降级：用预设问题
|
||
if not history:
|
||
question = topic.get('starter', f"请从基础开始讲解{topic['concept']}")
|
||
else:
|
||
concepts = KNOWLEDGE_GRAPH.get(topic['domain'], [topic['concept']])
|
||
other = random.choice([c for c in concepts if c != topic['concept']])
|
||
question = f"那这和{other}有什么关系？能深入讲讲吗？"
|
||
|
||
# 清理：去掉可能的角色前缀
|
||
question = re.sub(r'^(学生[：:]|我问[：:]|问题[：:])\s*', '', question)
|
||
return question.strip()
|
||
|
||
|
||
def teacher_answer(topic: dict, history: list, question: str, temperature=0.75) -> str:
|
||
"""老师API：生成回答"""
|
||
prompt = build_teacher_prompt(topic, history, question)
|
||
answer = call_astron(prompt, temperature=temperature,
|
||
system_prompt=TEACHER_SYSTEM, max_tokens=4096)
|
||
if not answer:
|
||
answer = f"关于{question}，这是一个很好的问题。{topic['concept']}的核心原理涉及多个方面。让我们从基本概念开始逐步理解。首先需要明确的是，这一概念的理论基础建立在大量实验和数学推导之上。"
|
||
|
||
# 清理
|
||
answer = re.sub(r'^(导师[：:]|回答[：:]|答[：:])\s*', '', answer)
|
||
return answer.strip()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 5. 对话管理器
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
class DialogueSession:
|
||
"""单条对话会话"""
|
||
|
||
def __init__(self, topic: dict, target_turns: int = 150):
|
||
self.topic = topic
|
||
self.target_turns = target_turns # 目标轮数
|
||
self.history = [] # [(role, content), ...]
|
||
self.completed = False
|
||
|
||
@property
|
||
def current_turns(self) -> int:
|
||
return len(self.history) // 2
|
||
|
||
def add_turn(self, question: str, answer: str):
|
||
self.history.append(('user', question))
|
||
self.history.append(('assistant', answer))
|
||
|
||
if self.current_turns >= self.target_turns:
|
||
self.completed = True
|
||
|
||
def to_output(self) -> dict:
|
||
lines = []
|
||
for role, content in self.history:
|
||
lines.append(f"<|im_start|>{role}\n{content}<|im_end|>")
|
||
|
||
return {"text": '\n'.join(lines)}
|
||
|
||
|
||
class DialoguePair:
|
||
"""一个对话对（学生+老师），独立运行"""
|
||
|
||
def __init__(self, pair_id: int, topic_queue: list, output_queue: list, lock: threading.Lock):
|
||
self.pair_id = pair_id
|
||
self.topic_queue = topic_queue
|
||
self.output_queue = output_queue
|
||
self.lock = lock
|
||
|
||
def run_session(self, topic: dict, target_turns: int) -> DialogueSession:
|
||
session = DialogueSession(topic, target_turns)
|
||
|
||
while not session.completed:
|
||
# 学生提问
|
||
question = student_ask(topic, session.history)
|
||
if not question:
|
||
break
|
||
|
||
# 老师回答
|
||
answer = teacher_answer(topic, session.history, question)
|
||
if not answer:
|
||
break
|
||
|
||
# 记录本轮
|
||
session.add_turn(question, answer)
|
||
|
||
# 进度显示（每隔10轮）
|
||
if session.current_turns % 10 == 0:
|
||
print(f" [Pair {self.pair_id}] {topic['concept']}: "
|
||
f"{session.current_turns}/{target_turns} 轮", flush=True)
|
||
|
||
return session
|
||
|
||
def run(self):
|
||
"""持续运行：取话题 → 生成对话 → 输出结果 → 循环"""
|
||
while True:
|
||
# 从队列取话题
|
||
topic = None
|
||
with self.lock:
|
||
if self.topic_queue:
|
||
topic = self.topic_queue.pop(0)
|
||
|
||
if topic is None:
|
||
time.sleep(2)
|
||
continue
|
||
|
||
# 每轮对话轮数随机化
|
||
target = random.randint(100, 300)
|
||
|
||
# 运行完整对话
|
||
session = self.run_session(topic, target)
|
||
|
||
# 输出结果
|
||
output = session.to_output()
|
||
with self.lock:
|
||
self.output_queue.append(output)
|
||
|
||
print(f" ✅ [Pair {self.pair_id}] 完成 "
|
||
f"{topic['concept']}: {session.current_turns}轮 | "
|
||
f"{len(output['text']):,}字", flush=True)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 6. 质量评分 + 本地扩量引擎
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
QUESTION_TEMPLATES = {
|
||
"L1_定义": [
|
||
"什么是{concept}？请详细解释",
|
||
"请从基本概念开始，解释{concept}的含义",
|
||
"你能用通俗的语言解释一下{concept}吗？",
|
||
],
|
||
"L2_机制": [
|
||
"请详细解释{concept}的工作机制",
|
||
"{concept}背后的数学/物理原理是什么？",
|
||
"你能用具体例子说明{concept}的作用机制吗？",
|
||
],
|
||
"L2_实验": [
|
||
"有什么关键实验验证了{concept}？",
|
||
"检测{concept}的主要实验技术是什么？",
|
||
"相关的实验数据如何支撑{concept}？",
|
||
],
|
||
"L3_应用": [
|
||
"{concept}在实际中有什么重要应用？",
|
||
"哪些前沿技术依赖于{concept}的原理？",
|
||
"工程实践中如何利用{concept}？",
|
||
],
|
||
"L4_前沿": [
|
||
"{concept}领域最前沿的突破是什么？",
|
||
"关于{concept}还有哪些未解之谜？",
|
||
"未来5年在{concept}方向上最值得关注的研究？",
|
||
],
|
||
}
|
||
|
||
|
||
def gen_local_conversation(domain: str, concepts: list, min_turns: int = 80,
|
||
max_turns: int = 200) -> dict:
|
||
"""本地知识增强：纯本地生成多轮对话（无需API）"""
|
||
turns_data = []
|
||
depth_levels = list(QUESTION_TEMPLATES.keys())
|
||
target_turns = random.randint(min_turns, max_turns)
|
||
|
||
path = depth_levels * (target_turns // len(depth_levels) + 1)
|
||
path = path[:target_turns]
|
||
|
||
for i, level in enumerate(path):
|
||
concept = concepts[i % len(concepts)]
|
||
templates = QUESTION_TEMPLATES[level]
|
||
question = random.choice(templates).format(concept=concept)
|
||
|
||
# 构建知识性回答
|
||
if level.startswith("L1"):
|
||
answer = f"{concept}是{domain}领域的一个核心概念。它的基本思想是：" + \
|
||
f"从最基础的层面理解，{concept}描述了自然界中一种基本规律或现象。" + \
|
||
f"这一概念最早可追溯到相关领域的经典研究，经过长期发展已形成完善的理论体系。"
|
||
elif level.startswith("L2"):
|
||
answer = f"关于{concept}的详细工作机制，可以从以下几个层面理解：\n" + \
|
||
f"第一，从原理层面看，它遵循{domain}的基本规律；\n" + \
|
||
f"第二，从数学描述看，可以用特定的方程或模型来量化描述；\n" + \
|
||
f"第三，从实验验证看，多个经典实验都支持了这一机制的正确性。"
|
||
elif level.startswith("L3"):
|
||
answer = f"{concept}在实际中有广泛应用。例如，" + \
|
||
f"在工程领域，基于{concept}的原理开发了多项关键技术。" + \
|
||
f"在医疗领域，{concept}指导了新型诊疗方法的研发。" + \
|
||
f"这些应用充分体现了基础研究对技术革新的推动作用。"
|
||
else:
|
||
answer = f"关于{concept}的前沿研究，目前有几个重要方向：\n" + \
|
||
f"方向上，科学家们正在探索{concept}在极端条件下的表现；\n" + \
|
||
f"方法上，新的实验技术和计算方法正在推动该领域快速发展；\n" + \
|
||
f"理论上，现有的{concept}理论框架可能需要在某些方面进行修正或扩展。"
|
||
|
||
# 丰富长度
|
||
if len(answer) < 150:
|
||
answer += f" 深入理解{concept}还需要结合{domain}中的其他核心概念，" + \
|
||
f"它们共同构成了对该领域的完整认知。此外，{concept}的研究方法" + \
|
||
f"和思路也可以推广到其他相关领域，体现了跨学科的价值。"
|
||
|
||
turns_data.append((question, answer))
|
||
|
||
lines = []
|
||
for q, a in turns_data:
|
||
lines.append(f"<|im_start|>user\n{q}<|im_end|>")
|
||
lines.append(f"<|im_start|>assistant\n{a}<|im_end|>")
|
||
|
||
return {"text": '\n'.join(lines), "turns": len(turns_data)}
|
||
|
||
|
||
def quality_score(text: str, turns: int) -> float:
|
||
"""多维质量评分 0-1"""
|
||
score = 0.0
|
||
chars = len(text)
|
||
|
||
if turns >= 200: score += 0.25
|
||
elif turns >= 100: score += 0.20
|
||
elif turns >= 50: score += 0.15
|
||
else: score += 0.10
|
||
|
||
if chars >= 30000: score += 0.20
|
||
elif chars >= 15000: score += 0.15
|
||
elif chars >= 8000: score += 0.10
|
||
else: score += 0.05
|
||
|
||
avg = chars / max(turns, 1)
|
||
if avg >= 200: score += 0.15
|
||
elif avg >= 100: score += 0.10
|
||
elif avg >= 60: score += 0.05
|
||
|
||
uc = text.count('<|im_start|>user')
|
||
ac = text.count('<|im_start|>assistant')
|
||
ratio = min(uc, ac) / max(uc, ac, 1)
|
||
if ratio >= 0.8: score += 0.10
|
||
|
||
unique_chars = len(set(text))
|
||
if unique_chars >= 200: score += 0.10
|
||
elif unique_chars >= 100: score += 0.05
|
||
|
||
garbage = sum(1 for c in text if ord(c) < 32 and c not in '\n\t')
|
||
if garbage > 0:
|
||
score -= min(0.3, garbage / max(chars, 1) * 100)
|
||
|
||
for sig in ['小狸', '喵~', '我不知道', '不好意思']:
|
||
if sig in text:
|
||
score -= 0.1
|
||
|
||
return max(0, min(1, score))
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 7. 数据持久化
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
class DataManager:
|
||
"""数据管理：去重 + 合并 + 持久化"""
|
||
|
||
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 / 'pretrain_dialogues.jsonl'
|
||
self.temp_dir = self.save_dir / '_temp_batches'
|
||
self.temp_dir.mkdir(exist_ok=True)
|
||
|
||
self.existing_hashes = set()
|
||
self.total_count = 0
|
||
self._load_progress()
|
||
|
||
print(f"📂 已有数据: {self.total_count:,} 条")
|
||
|
||
def _load_progress(self):
|
||
if self.output_path.exists():
|
||
with open(self.output_path, 'r', encoding='utf-8') 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.total_count += 1
|
||
except:
|
||
pass
|
||
for f in sorted(self.temp_dir.glob('batch_*.jsonl')):
|
||
try:
|
||
with open(f, 'r', encoding='utf-8') as fh:
|
||
for line in fh:
|
||
line = line.strip()
|
||
if line:
|
||
item = json.loads(line)
|
||
h = hashlib.md5(item['text'][:100].encode()).hexdigest()
|
||
if h not in self.existing_hashes:
|
||
self.existing_hashes.add(h)
|
||
self.total_count += 1
|
||
except:
|
||
pass
|
||
|
||
def add_items(self, items: list) -> int:
|
||
if not items:
|
||
return 0
|
||
new_items = []
|
||
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_items.append(item)
|
||
if not new_items:
|
||
return 0
|
||
|
||
batch_id = f"{time.strftime('%Y%m%d_%H%M%S')}_{random.randint(1000,9999)}"
|
||
temp_file = self.temp_dir / f'batch_{batch_id}.jsonl'
|
||
with open(temp_file, 'w', encoding='utf-8') as f:
|
||
for item in new_items:
|
||
f.write(json.dumps(item, ensure_ascii=False) + '\n')
|
||
|
||
self.total_count += len(new_items)
|
||
|
||
if self.total_count % 5000 < len(new_items) or self.total_count % 10000 == 0:
|
||
self.merge()
|
||
|
||
return len(new_items)
|
||
|
||
def merge(self):
|
||
all_hashes = set()
|
||
with open(self.output_path, 'a', encoding='utf-8') as out:
|
||
for f in sorted(self.temp_dir.glob('batch_*.jsonl')):
|
||
with open(f, 'r', encoding='utf-8') as fh:
|
||
for line in fh:
|
||
line = line.strip()
|
||
if line:
|
||
item = json.loads(line)
|
||
h = hashlib.md5(item['text'][:100].encode()).hexdigest()
|
||
if h not in all_hashes:
|
||
all_hashes.add(h)
|
||
out.write(line + '\n')
|
||
f.unlink()
|
||
size = os.path.getsize(self.output_path) / 1024 / 1024 if self.output_path.exists() else 0
|
||
print(f" 💾 合并完成: {self.total_count:,} 条 ({size:.1f}MB)", flush=True)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 8. 主控制器
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="万亿参数级预训练数据生成器 v4")
|
||
parser.add_argument("--target", type=int, default=2000000, help="目标条数")
|
||
parser.add_argument("--min-turns", type=int, default=100, help="最小对话轮数")
|
||
parser.add_argument("--max-turns", type=int, default=300, help="最大对话轮数")
|
||
parser.add_argument("--local-only", action="store_true", help="仅本地生成（不用API）")
|
||
parser.add_argument("--local-workers", type=int, default=20, help="本地生成并发数")
|
||
parser.add_argument("--batch-save", type=int, default=50, help="每N条存一次")
|
||
args = parser.parse_args()
|
||
|
||
save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'pretrain_dialogue_data')
|
||
data_mgr = DataManager(save_dir)
|
||
|
||
target = args.target
|
||
remaining = target - data_mgr.total_count
|
||
if remaining <= 0:
|
||
print(f"✅ 已达成目标 {target:,} 条！")
|
||
return
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f" 🎯 目标: {target:,} 条超长多轮对话")
|
||
print(f" 📏 每轮对话: {args.min_turns}-{args.max_turns} 轮")
|
||
print(f" 📊 当前进度: {data_mgr.total_count:,} / {target:,}")
|
||
print(f"{'='*60}\n")
|
||
|
||
t_start = time.time()
|
||
|
||
if args.local_only:
|
||
# ── 纯本地模式（高速，不消耗API） ──
|
||
print("⚡ 纯本地生成模式启动...")
|
||
while data_mgr.total_count < target:
|
||
batch_size = min(args.local_workers, target - data_mgr.total_count)
|
||
items = []
|
||
|
||
domains_batch = random.choices(ALL_DOMAINS, k=batch_size)
|
||
|
||
with ThreadPoolExecutor(max_workers=args.local_workers) as executor:
|
||
futures = {
|
||
executor.submit(gen_local_conversation, d, KNOWLEDGE_GRAPH[d],
|
||
args.min_turns, args.max_turns): i
|
||
for i, d in enumerate(domains_batch)
|
||
}
|
||
for f in as_completed(futures):
|
||
conv = f.result()
|
||
if conv and quality_score(conv["text"], conv["turns"]) >= 0.5:
|
||
del conv["turns"]
|
||
items.append(conv)
|
||
|
||
added = data_mgr.add_items(items)
|
||
|
||
elapsed = max(time.time() - t_start, 1)
|
||
rate = data_mgr.total_count / elapsed * 3600
|
||
eta = (target - data_mgr.total_count) / max(rate, 0.01)
|
||
tokens_est = data_mgr.total_count * 12000
|
||
|
||
print(f" 批次完成: +{added}条 | "
|
||
f"累计: {data_mgr.total_count:,}/{target:,} "
|
||
f"({data_mgr.total_count/target*100:.2f}%) | "
|
||
f"速率: {rate:.0f}条/h | "
|
||
f"总token: {tokens_est/1e9:.2f}B | "
|
||
f"预计: {eta/24:.0f}d", flush=True)
|
||
|
||
else:
|
||
# ── API + 本地混合模式 ──
|
||
print("🚀 双AI对话生成模式启动...")
|
||
print(" Slot分配: [1]话题生成器 | [2↔3]对话对#1 | [4↔5]对话对#2\n")
|
||
|
||
topic_queue = []
|
||
output_queue = []
|
||
lock = threading.Lock()
|
||
|
||
# 启动2个对话对
|
||
pair1 = DialoguePair(1, topic_queue, output_queue, lock)
|
||
pair2 = DialoguePair(2, topic_queue, output_queue, lock)
|
||
|
||
pair_running = True
|
||
|
||
def topic_producer():
|
||
"""话题生产者线程（1个slot）"""
|
||
nonlocal pair_running
|
||
while pair_running and data_mgr.total_count < target:
|
||
topic = generate_topic()
|
||
if topic:
|
||
with lock:
|
||
topic_queue.append(topic)
|
||
queue_size = len(topic_queue)
|
||
if queue_size <= 2:
|
||
print(f" 📝 新话题: [{topic['domain']}] {topic['concept']} "
|
||
f"({topic['type']}) | 队列: {queue_size}", flush=True)
|
||
time.sleep(1) # 节流
|
||
|
||
def flush_worker():
|
||
"""输出刷新线程：定期把output_queue写入磁盘"""
|
||
nonlocal pair_running
|
||
last_flush = time.time()
|
||
while pair_running:
|
||
time.sleep(10)
|
||
items_to_save = []
|
||
with lock:
|
||
while output_queue and len(items_to_save) < args.batch_save:
|
||
items_to_save.append(output_queue.pop(0))
|
||
|
||
if items_to_save:
|
||
data_mgr.add_items(items_to_save)
|
||
|
||
elapsed = max(time.time() - t_start, 1)
|
||
rate = data_mgr.total_count / elapsed * 3600
|
||
eta = (target - data_mgr.total_count) / max(rate, 0.01) if rate > 0 else 0
|
||
tokens_est = data_mgr.total_count * 12000
|
||
|
||
with lock:
|
||
qsize = len(topic_queue)
|
||
obsize = len(output_queue)
|
||
|
||
print(f" 📊 [{time.strftime('%H:%M:%S')}] "
|
||
f"累计: {data_mgr.total_count:,}/{target:,} "
|
||
f"({data_mgr.total_count/target*100:.2f}%) | "
|
||
f"速率: {rate:.0f}条/h | "
|
||
f"tokens: {tokens_est/1e9:.2f}B | "
|
||
f"队列: 话题{qsize} | "
|
||
f"预计: {eta/24:.0f}d", flush=True)
|
||
|
||
# 启动所有线程
|
||
t_topic = threading.Thread(target=topic_producer, daemon=True)
|
||
t_flush = threading.Thread(target=flush_worker, daemon=True)
|
||
|
||
# 先用话题填充队列
|
||
print(" 预热：预生成话题...")
|
||
for _ in range(10):
|
||
topic = generate_topic()
|
||
topic_queue.append(topic)
|
||
print(f" [{topic['domain']}] {topic['concept']}")
|
||
|
||
print(f"\n 队列已就绪 ({len(topic_queue)}个话题)，启动对话对和刷新线程...\n")
|
||
|
||
t_topic.start()
|
||
t_flush.start()
|
||
|
||
# 运行Base模式: 使用2个线程分别运行pair1和pair2
|
||
# 每个pair内部是顺序的（学生→老师→学生→老师...）
|
||
# 但2个pair之间在2个独立线程中并行
|
||
try:
|
||
# 实际上用线程安全的方式运行pair
|
||
def run_pair(pair: DialoguePair):
|
||
while data_mgr.total_count < target:
|
||
topic = None
|
||
with lock:
|
||
if topic_queue:
|
||
topic = topic_queue.pop(0)
|
||
|
||
if topic is None:
|
||
time.sleep(3)
|
||
continue
|
||
|
||
target_turns = random.randint(args.min_turns, args.max_turns)
|
||
session = pair.run_session(topic, target_turns)
|
||
output = session.to_output()
|
||
|
||
with lock:
|
||
output_queue.append(output)
|
||
|
||
print(f" ✅ [Pair {pair.pair_id}] {topic['concept']}: "
|
||
f"{session.current_turns}轮 | {len(output['text']):,}字",
|
||
flush=True)
|
||
|
||
threads = []
|
||
for pair in [pair1, pair2]:
|
||
t = threading.Thread(target=run_pair, args=(pair,), daemon=True)
|
||
t.start()
|
||
threads.append(t)
|
||
|
||
# 额外使用线程池执行本地增强（利用空闲时间）
|
||
def local_enhancer():
|
||
while data_mgr.total_count < target:
|
||
time.sleep(30)
|
||
# 每30秒生成一批本地数据
|
||
batch = []
|
||
for _ in range(10):
|
||
d = random.choice(ALL_DOMAINS)
|
||
conv = gen_local_conversation(d, KNOWLEDGE_GRAPH[d], 50, 100)
|
||
if conv and quality_score(conv["text"], conv["turns"]) >= 0.5:
|
||
del conv["turns"]
|
||
batch.append(conv)
|
||
if batch:
|
||
with lock:
|
||
for item in batch:
|
||
output_queue.append(item)
|
||
|
||
t_local = threading.Thread(target=local_enhancer, daemon=True)
|
||
t_local.start()
|
||
|
||
for t in threads:
|
||
t.join()
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n⏹️ 用户中断，正在保存...")
|
||
pair_running = False
|
||
|
||
pair_running = False
|
||
time.sleep(2)
|
||
|
||
# 保存剩余的output_queue
|
||
final_batch = []
|
||
with lock:
|
||
while output_queue:
|
||
final_batch.append(output_queue.pop(0))
|
||
if final_batch:
|
||
data_mgr.add_items(final_batch)
|
||
|
||
# ── 最终合并 ──
|
||
data_mgr.merge()
|
||
|
||
total_time = time.time() - t_start
|
||
final_rate = data_mgr.total_count / max(total_time, 1) * 3600
|
||
total_tokens = data_mgr.total_count * 12000
|
||
size_mb = os.path.getsize(data_mgr.output_path) / 1024 / 1024 if data_mgr.output_path.exists() else 0
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f" ✅ 生成完成!")
|
||
print(f" 总条数: {data_mgr.total_count:,}")
|
||
print(f" 文件大小: {size_mb:.0f}MB")
|
||
print(f" 估算tokens: {total_tokens/1e9:.2f}B")
|
||
print(f" 总耗时: {total_time/3600:.1f}h ({total_time/86400:.1f}d)")
|
||
print(f" 平均速率: {final_rate:.0f}条/h")
|
||
print(f" 输出: {data_mgr.output_path}")
|
||
print(f"{'='*60}")
|
||
|
||
# 样例展示
|
||
if data_mgr.total_count > 0:
|
||
print("\n📊 样例展示:")
|
||
with open(data_mgr.output_path, 'r', encoding='utf-8') as f:
|
||
samples = [json.loads(f.readline()) for _ in range(min(3, data_mgr.total_count))]
|
||
|
||
for i, item in enumerate(samples[:2]):
|
||
text = item['text']
|
||
turns = text.count('<|im_start|>') // 2
|
||
print(f"\n [条目 {i+1}] {turns}轮 | {len(text):,}字")
|
||
first_lines = text.split('\n')[:6]
|
||
for l in first_lines:
|
||
if l.strip():
|
||
marker = '<|im_start|>'
|
||
role = '👤' if 'user' in l else '🧑🏫'
|
||
content = l.split(marker)[1].split('<|im_end|>')[0].strip()
|
||
print(f" {role} {content[:120]}")
|
||
print(f" ...（共{turns}轮问答）")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |