#!/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()