diff --git a/gen_pretrain.py b/gen_pretrain.py index 5f5dbd8..3b4c904 100644 --- a/gen_pretrain.py +++ b/gen_pretrain.py @@ -1,401 +1,870 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- """ -小狸 AI - 通用知识预训练数据生成器 -==================================== -生成通用中文知识文本（非闲聊/非角色），用于模型预训练。 -每批并发调用 Astron API 批量生成，目标 200 万条。 +小狸 AI - 万亿参数级预训练数据生成器 (v4) +============================================== +核心架构：双AI互相对话（师生对练）+ 话题生成器 -用法: - python3 gen_pretrain.py # 启动生成 - python3 gen_pretrain.py --count 50000 # 指定本轮目标 - python3 gen_pretrain.py --resume # 断点续跑 +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 +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 -# ── 知识领域（20 大类，覆盖百科级知识） ──────────────── -KNOWLEDGE_DOMAINS = { - "天文宇宙": [ - "太阳系的行星特征", "恒星的生命周期", "黑洞的形成与性质", "银河系的结构", - "宇宙大爆炸理论", "星座与观测", "月球与潮汐", "彗星与小行星", - "太空探索历史", "国际空间站", "火星探测任务", "引力波", + +# ═══════════════════════════════════════════════════════════ +# 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}的作用机制吗？", ], - "生命科学": [ - "细胞结构与功能", "DNA与基因", "光合作用过程", "人体器官系统", - "免疫系统工作原理", "病毒与细菌区别", "生态系统食物链", "进化论", - "恐龙灭绝假说", "珍稀动物保护", "昆虫分类", "鸟类迁徙", + "L2_实验": [ + "有什么关键实验验证了{concept}？", + "检测{concept}的主要实验技术是什么？", + "相关的实验数据如何支撑{concept}？", ], - "中国历史": [ - "夏商周三代", "秦统一六国", "汉朝丝绸之路", "三国鼎立", - "唐朝盛世", "宋朝科技成就", "元朝疆域", "明朝航海", - "清朝康乾盛世", "鸦片战争", "辛亥革命", "抗日战争", + "L3_应用": [ + "{concept}在实际中有什么重要应用？", + "哪些前沿技术依赖于{concept}的原理？", + "工程实践中如何利用{concept}？", ], - "世界历史": [ - "古埃及文明", "古希腊城邦", "罗马帝国", "大航海时代", - "工业革命", "法国大革命", "第一次世界大战", "第二次世界大战", - "冷战时期", "联合国成立", "人类登月", "互联网起源", - ], - "地理探索": [ - "世界七大洲", "四大洋特征", "最长河流排名", "最高山峰", - "亚马逊雨林", "撒哈拉沙漠", "珊瑚礁生态", "火山湖形成", - "世界著名瀑布", "极光形成原理", "时区划分", "地图投影", - ], - "物理定律": [ - "牛顿三大定律", "万有引力", "电磁感应", "相对论浅说", - "量子力学入门", "光的折射反射", "声音传播特性", "热能传导方式", - "浮力与压强", "简单机械原理", "电路基础知识", "核能原理", - ], - "化学元素": [ - "元素周期表", "水的化学性质", "空气成分", "酸碱中和", - "氧化还原反应", "常见化学元素", "分子与原子", "晶体结构", - "有机化合物", "高分子材料", "催化剂作用", "化学反应类型", - ], - "数学原理": [ - "勾股定理证明", "圆周率计算史", "黄金分割比例", "微积分基础思想", - "概率论起源", "拓扑学入门", "斐波那契数列", "二进制原理", - "对称与群论", "统计与平均数", "几何发展史", "数学建模", - ], - "科技发明": [ - "印刷术发明", "蒸汽机改良", "电灯发明过程", "电话通信原理", - "无线电技术", "晶体管与芯片", "计算机发展史", "互联网诞生", - "人工智能简史", "机器人发展", "GPS定位原理", "基因编辑技术", - ], - "文化艺术": [ - "中国书法流派", "水墨画技法", "京剧角色分类", "交响乐结构", - "芭蕾舞起源", "电影发展史", "摄影构图原理", "建筑风格演变", - "敦煌莫高窟", "文艺复兴艺术", "印象派绘画", "现代艺术流派", - ], - "植物世界": [ - "植物分类系统", "开花植物特征", "树木年轮秘密", "食虫植物", - "药用植物知识", "蕨类植物特点", "多肉植物养护", "水生植物适应", - "果实种子传播", "植物激素作用", "菌类不是植物", "藻类多样性", - ], - "动物王国": [ - "哺乳动物特征", "爬行动物演化", "两栖动物生命周期", "鸟类飞行原理", - "鱼类分类", "海洋哺乳动物", "昆虫社会结构", "蛛形纲特征", - "濒危动物保护", "动物伪装术", "动物迁徙路线", "宠物饲养常识", - ], - "医学健康": [ - "人体骨骼结构", "心脏血液循环", "消化系统工作原理", "呼吸系统功能", - "营养均衡膳食", "疫苗工作原理", "抗生素使用原则", "常见疾病预防", - "睡眠科学研究", "运动健身原理", "心理健康知识", "急救基本常识", - ], - "工程技术": [ - "桥梁结构类型", "摩天大楼建造", "地铁隧道挖掘", "水力发电原理", - "风力发电技术", "太阳能利用", "高铁技术原理", "飞机制造原理", - "船舶设计", "芯片制造工艺", "光纤通信", "3D打印技术", - ], - "语言文学": [ - "汉语文字演变", "成语典故来源", "古诗词格律", "小说发展流派", - "修辞手法分类", "甲骨文发现", "四大名著简介", "世界文学名著", - "诗歌韵律基础", "写作技巧", "语言学基础", "翻译原则", - ], - "哲学思想": [ - "儒家核心思想", "道家哲学", "佛学基本概念", "西方哲学起源", - "逻辑学基础", "伦理学基本问题", "美学原理", "辩证法思想", - "知行合一", "中庸之道", "存在主义简介", "实用主义哲学", - ], - "经济常识": [ - "货币起源与职能", "供求关系原理", "市场经济学基础", "国际贸易概念", - "GDP的含义", "通货膨胀原因", "银行职能", "股票基本原理", - "保险基础知识", "个人理财入门", "经济危机历史", "数字经济趋势", - ], - "节日文化": [ - "春节传统习俗", "中秋节起源", "端午节纪念", "重阳节含义", - "圣诞节文化", "感恩节由来", "万圣节传统", "世界节日一览", - "二十四节气", "传统婚礼习俗", "茶文化历史", "饮食文化地域差异", - ], - "日常生活": [ - "垃圾分类知识", "交通安全规则", "消防安全常识", "食品安全知识", - "环保节能技巧", "家居收纳方法", "时间管理技巧", "沟通表达能力", - "礼仪基本规范", "旅行准备清单", "天气预测常识", "手机使用安全", + "L4_前沿": [ + "{concept}领域最前沿的突破是什么？", + "关于{concept}还有哪些未解之谜？", + "未来5年在{concept}方向上最值得关注的研究？", ], } -ALL_TOPICS = [] -for domain, topics in KNOWLEDGE_DOMAINS.items(): - for t in topics: - ALL_TOPICS.append((domain, t)) + +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 build_prompt(domain: str, topics: list, batch_size: int = 10) -> str: - """构建批量生成prompt""" - topic_list = "\n".join(f"{i+1}. {t}" for i, t in enumerate(topics)) - return f"""你是一个百科全书作者，请用中文撰写 {batch_size} 条高质量的科普段落。 - -领域：{domain} -子话题： -{topic_list} - -要求： -1. 每条 60-150 字，言之有物，包含具体知识 -2. 语言流畅、客观准确，适合用于训练语言模型 -3. 每条独立成段，不互相引用 -4. 不要用"首先/其次/总之"等连接词开头 -5. 末尾不要问句，用句号结尾 -6. 不要出现"小狸"等角色名 - -输出格式（严格按此格式）： -【1】第一条知识段落内容。 -【2】第二条知识段落内容。 -【3】第三条知识段落内容。 -... - -请开始生成：""" +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)) -def parse_passages(text: str) -> list: - """解析批量返回的文本""" - passages = [] - for line in text.strip().split('\n'): - line = line.strip() - if not line: - continue - # 匹配 【N】 或 [N] 或 N. 或 N、开头 - for sep in ['】', ']', '.', '、', ' ']: - if sep in line[:8]: - content = line[line.index(sep)+len(sep):].strip() - if len(content) >= 30: - passages.append(content) - break - # fallback: 如果没有匹配到格式，尝试按句号分割 - if not passages and len(text) >= 50: - # 尝试按编号提取 - import re - matches = re.findall(r'(?:【(\d+)】|\[(\d+)\]|(\d+)[\.、])\s*(.+?)(?=(?:【|\n|$))', text + '\n', re.DOTALL) - for m in matches: - content = m[3].strip() if m[3] else (m[2] if m[2] else '') - if len(content) >= 30: - passages.append(content) - return passages +# ═══════════════════════════════════════════════════════════ +# 7. 数据持久化 +# ═══════════════════════════════════════════════════════════ - -def generate_batch(domain: str, topics: list, batch_size: int = 10, - temperature: float = 0.7) -> list: - """一次 API 调用生成多条知识段落""" - prompt = build_prompt(domain, topics, batch_size) - text = call_astron(prompt, temperature=temperature) - if not text: - return [] - return parse_passages(text) - - -def format_check(text: str) -> float: - """对生成文本做简单质量评分 0-1""" - score = 0.5 - # 长度惩罚 - if len(text) < 30: - score -= 0.3 - if len(text) > 200: - score -= 0.1 - # 质量信号 - has_period = text.rstrip().endswith(('。', '！', '？')) - if has_period: - score += 0.2 - no_role = '小狸' not in text and '喵' not in text - if no_role: - score += 0.2 - no_qmark = '?' not in text and '？' not in text - if no_qmark: - score += 0.1 - return min(1.0, max(0.0, score)) - - -# ── 主生成器 ──────────────────────────────────────────── -class PretrainDataGenerator: +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_knowledge.json' + 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 = set() - self.existing_texts = [] + 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: - self.existing_texts = json.load(f) - for item in self.existing_texts: - h = hashlib.md5(item['text'][:80].encode()).hexdigest() - self.existing.add(h) - - self.total_generated = len(self.existing_texts) - - # 已保存的临时批次 - self._load_temp_batches() - - print(f"📊 已有数据: {self.total_generated} 条") - - def _load_temp_batches(self): - for f in sorted(self.temp_dir.glob('batch_*.json')): - with open(f, 'r', encoding='utf-8') as fh: - batch = json.load(fh) - added = 0 - for item in batch: - h = hashlib.md5(item['text'][:80].encode()).hexdigest() - if h not in self.existing: - self.existing.add(h) - self.existing_texts.append(item) - added += 1 - self.total_generated += added - - def add_batch(self, passages: list): - if not passages: + 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 p in passages: - p = p.strip() - if len(p) < 30: - continue - score = format_check(p) - if score < 0.4: - continue - h = hashlib.md5(p[:80].encode()).hexdigest() - if h not in self.existing: - self.existing.add(h) - new_items.append({"text": p}) - + 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 = time.strftime('%Y%m%d_%H%M%S') + f'_{random.randint(1000,9999)}' - temp_file = self.temp_dir / f'batch_{batch_id}.json' + 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: - json.dump(new_items, f, ensure_ascii=False) + for item in new_items: + f.write(json.dumps(item, ensure_ascii=False) + '\n') - self.existing_texts.extend(new_items) - self.total_generated = len(self.existing_texts) + self.total_count += len(new_items) - # 每 5000 条合并一次主文件 - if self.total_generated % 5000 < 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_temps = sorted(self.temp_dir.glob('batch_*.json')) - merged = list(self.existing_texts) - - with open(self.output_path, 'w', encoding='utf-8') as f: - json.dump(merged, f, ensure_ascii=False, indent=2) - - # 清理旧临时文件 - for f in all_temps[:-100]: # 保留最近100个 - f.unlink() - - print(f" 💾 合并保存: {self.output_path.name} ({len(merged)} 条)") + 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) - def get_progress(self): - return self.total_generated +# ═══════════════════════════════════════════════════════════ +# 8. 主控制器 +# ═══════════════════════════════════════════════════════════ def main(): - parser = argparse.ArgumentParser(description="通用预训练数据生成器") - parser.add_argument("--count", type=int, default=2000000, help="目标条数") - parser.add_argument("--resume", action="store_true", help="断点续跑") - parser.add_argument("--workers", type=int, default=5, help="并发数") - parser.add_argument("--batch", type=int, default=12, help="每轮每线程生成条数") + 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_data_gen') - gen = PretrainDataGenerator(save_dir) + save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'pretrain_dialogue_data') + data_mgr = DataManager(save_dir) - if args.resume: - print(f"📌 断点续跑，当前进度: {gen.get_progress()}") - - target = args.count - remaining = target - gen.get_progress() + target = args.target + remaining = target - data_mgr.total_count if remaining <= 0: - print(f"✅ 已达成目标 {target} 条！") + print(f"✅ 已达成目标 {target:,} 条！") return - print(f"🎯 目标: {target} 条 | 还需: {remaining} 条") - print(f"⚡ 并发: {args.workers} | 每线程: {args.batch} 条/轮") - print(f"{'='*60}") + 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") - workers = args.workers - batch_per_call = args.batch - total_done = gen.get_progress() t_start = time.time() - round_num = 0 - try: - while total_done < target: - round_num += 1 + if args.local_only: + # ── 纯本地模式（高速，不消耗API） ── + print("⚡ 纯本地生成模式启动...") + while data_mgr.total_count < target: + batch_size = min(args.local_workers, target - data_mgr.total_count) + items = [] - # 分配领域和话题 - round_tasks = [] - for _ in range(workers): - domain = random.choice(list(KNOWLEDGE_DOMAINS.keys())) - topics = random.sample(KNOWLEDGE_DOMAINS[domain], - min(batch_per_call, len(KNOWLEDGE_DOMAINS[domain]))) - temp = random.choice([0.6, 0.7, 0.75, 0.8, 0.85]) - round_tasks.append((domain, topics, temp)) + domains_batch = random.choices(ALL_DOMAINS, k=batch_size) - # 5 并发 - all_passages = [] - with ThreadPoolExecutor(max_workers=workers) as executor: + with ThreadPoolExecutor(max_workers=args.local_workers) as executor: futures = { - executor.submit(generate_batch, d, t, batch_per_call, temp): i - for i, (d, t, temp) in enumerate(round_tasks) + 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): - passages = f.result() - all_passages.extend(passages) + conv = f.result() + if conv and quality_score(conv["text"], conv["turns"]) >= 0.5: + del conv["turns"] + items.append(conv) - added = gen.add_batch(all_passages) - total_done = gen.get_progress() + added = data_mgr.add_items(items) - # 统计 - elapsed = time.time() - t_start - rate = total_done / max(elapsed, 1) * 60 - eta_sec = (target - total_done) / max(rate / 60, 0.01) + 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" 第{round_num}轮: +{added}条 | " - f"累计: {total_done}/{target} ({total_done/target*100:.1f}%) | " - f"速率: {rate:.0f}条/分 | " - f"预计剩余: {eta_sec/3600:.1f}h") - - if total_done >= target: - break + 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) - except KeyboardInterrupt: - print("\n⏹️ 用户中断") + 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) - # 最终合并 - gen.merge() + # ── 最终合并 ── + data_mgr.merge() total_time = time.time() - t_start - final_rate = total_done / max(total_time, 1) * 60 + 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" 总条数: {total_done}") - print(f" 总耗时: {total_time/3600:.1f}h") - print(f" 平均速率: {final_rate:.0f}条/分") - print(f" 输出: {gen.output_path}") + 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}") - # 统计领域分布 - print("\n📊 样例展示 (前5条):") - with open(gen.output_path, 'r', encoding='utf-8') as f: - data = json.load(f) - for i, item in enumerate(data[:5]): - print(f" [{i+1}] {item['text'][:80]}...") + # 样例展示 + 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__":