#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 小狸 AI - 通用知识预训练数据生成器 ==================================== 生成通用中文知识文本（非闲聊/非角色），用于模型预训练。 每批并发调用 Astron API 批量生成，目标 200 万条。 用法: python3 gen_pretrain.py # 启动生成 python3 gen_pretrain.py --count 50000 # 指定本轮目标 python3 gen_pretrain.py --resume # 断点续跑 """ import sys, os, json, time, random, hashlib, argparse 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 = { "天文宇宙": [ "太阳系的行星特征", "恒星的生命周期", "黑洞的形成与性质", "银河系的结构", "宇宙大爆炸理论", "星座与观测", "月球与潮汐", "彗星与小行星", "太空探索历史", "国际空间站", "火星探测任务", "引力波", ], "地球科学": [ "板块构造理论", "火山喷发原理", "地震成因与预警", "岩石分类", "矿物与矿石", "水循环过程", "气候带分布", "洋流与天气", "冰川与极地", "沙漠生态系统", "喀斯特地貌", "化石形成", ], "生命科学": [ "细胞结构与功能", "DNA与基因", "光合作用过程", "人体器官系统", "免疫系统工作原理", "病毒与细菌区别", "生态系统食物链", "进化论", "恐龙灭绝假说", "珍稀动物保护", "昆虫分类", "鸟类迁徙", ], "中国历史": [ "夏商周三代", "秦统一六国", "汉朝丝绸之路", "三国鼎立", "唐朝盛世", "宋朝科技成就", "元朝疆域", "明朝航海", "清朝康乾盛世", "鸦片战争", "辛亥革命", "抗日战争", ], "世界历史": [ "古埃及文明", "古希腊城邦", "罗马帝国", "大航海时代", "工业革命", "法国大革命", "第一次世界大战", "第二次世界大战", "冷战时期", "联合国成立", "人类登月", "互联网起源", ], "地理探索": [ "世界七大洲", "四大洋特征", "最长河流排名", "最高山峰", "亚马逊雨林", "撒哈拉沙漠", "珊瑚礁生态", "火山湖形成", "世界著名瀑布", "极光形成原理", "时区划分", "地图投影", ], "物理定律": [ "牛顿三大定律", "万有引力", "电磁感应", "相对论浅说", "量子力学入门", "光的折射反射", "声音传播特性", "热能传导方式", "浮力与压强", "简单机械原理", "电路基础知识", "核能原理", ], "化学元素": [ "元素周期表", "水的化学性质", "空气成分", "酸碱中和", "氧化还原反应", "常见化学元素", "分子与原子", "晶体结构", "有机化合物", "高分子材料", "催化剂作用", "化学反应类型", ], "数学原理": [ "勾股定理证明", "圆周率计算史", "黄金分割比例", "微积分基础思想", "概率论起源", "拓扑学入门", "斐波那契数列", "二进制原理", "对称与群论", "统计与平均数", "几何发展史", "数学建模", ], "科技发明": [ "印刷术发明", "蒸汽机改良", "电灯发明过程", "电话通信原理", "无线电技术", "晶体管与芯片", "计算机发展史", "互联网诞生", "人工智能简史", "机器人发展", "GPS定位原理", "基因编辑技术", ], "文化艺术": [ "中国书法流派", "水墨画技法", "京剧角色分类", "交响乐结构", "芭蕾舞起源", "电影发展史", "摄影构图原理", "建筑风格演变", "敦煌莫高窟", "文艺复兴艺术", "印象派绘画", "现代艺术流派", ], "植物世界": [ "植物分类系统", "开花植物特征", "树木年轮秘密", "食虫植物", "药用植物知识", "蕨类植物特点", "多肉植物养护", "水生植物适应", "果实种子传播", "植物激素作用", "菌类不是植物", "藻类多样性", ], "动物王国": [ "哺乳动物特征", "爬行动物演化", "两栖动物生命周期", "鸟类飞行原理", "鱼类分类", "海洋哺乳动物", "昆虫社会结构", "蛛形纲特征", "濒危动物保护", "动物伪装术", "动物迁徙路线", "宠物饲养常识", ], "医学健康": [ "人体骨骼结构", "心脏血液循环", "消化系统工作原理", "呼吸系统功能", "营养均衡膳食", "疫苗工作原理", "抗生素使用原则", "常见疾病预防", "睡眠科学研究", "运动健身原理", "心理健康知识", "急救基本常识", ], "工程技术": [ "桥梁结构类型", "摩天大楼建造", "地铁隧道挖掘", "水力发电原理", "风力发电技术", "太阳能利用", "高铁技术原理", "飞机制造原理", "船舶设计", "芯片制造工艺", "光纤通信", "3D打印技术", ], "语言文学": [ "汉语文字演变", "成语典故来源", "古诗词格律", "小说发展流派", "修辞手法分类", "甲骨文发现", "四大名著简介", "世界文学名著", "诗歌韵律基础", "写作技巧", "语言学基础", "翻译原则", ], "哲学思想": [ "儒家核心思想", "道家哲学", "佛学基本概念", "西方哲学起源", "逻辑学基础", "伦理学基本问题", "美学原理", "辩证法思想", "知行合一", "中庸之道", "存在主义简介", "实用主义哲学", ], "经济常识": [ "货币起源与职能", "供求关系原理", "市场经济学基础", "国际贸易概念", "GDP的含义", "通货膨胀原因", "银行职能", "股票基本原理", "保险基础知识", "个人理财入门", "经济危机历史", "数字经济趋势", ], "节日文化": [ "春节传统习俗", "中秋节起源", "端午节纪念", "重阳节含义", "圣诞节文化", "感恩节由来", "万圣节传统", "世界节日一览", "二十四节气", "传统婚礼习俗", "茶文化历史", "饮食文化地域差异", ], "日常生活": [ "垃圾分类知识", "交通安全规则", "消防安全常识", "食品安全知识", "环保节能技巧", "家居收纳方法", "时间管理技巧", "沟通表达能力", "礼仪基本规范", "旅行准备清单", "天气预测常识", "手机使用安全", ], } ALL_TOPICS = [] for domain, topics in KNOWLEDGE_DOMAINS.items(): for t in topics: ALL_TOPICS.append((domain, t)) 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 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 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: 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.temp_dir = self.save_dir / '_temp_batches' self.temp_dir.mkdir(exist_ok=True) # 加载已有进度 self.existing = set() self.existing_texts = [] 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: 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}) 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' with open(temp_file, 'w', encoding='utf-8') as f: json.dump(new_items, f, ensure_ascii=False) self.existing_texts.extend(new_items) self.total_generated = len(self.existing_texts) # 每 5000 条合并一次主文件 if self.total_generated % 5000 < len(new_items): 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)} 条)") def get_progress(self): return self.total_generated 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="每轮每线程生成条数") args = parser.parse_args() save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'pretrain_data_gen') gen = PretrainDataGenerator(save_dir) if args.resume: print(f"📌 断点续跑，当前进度: {gen.get_progress()}") target = args.count remaining = target - gen.get_progress() if remaining <= 0: print(f"✅ 已达成目标 {target} 条！") return print(f"🎯 目标: {target} 条 | 还需: {remaining} 条") print(f"⚡ 并发: {args.workers} | 每线程: {args.batch} 条/轮") print(f"{'='*60}") 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 # 分配领域和话题 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)) # 5 并发 all_passages = [] with ThreadPoolExecutor(max_workers=workers) as executor: futures = { executor.submit(generate_batch, d, t, batch_per_call, temp): i for i, (d, t, temp) in enumerate(round_tasks) } for f in as_completed(futures): passages = f.result() all_passages.extend(passages) added = gen.add_batch(all_passages) total_done = gen.get_progress() # 统计 elapsed = time.time() - t_start rate = total_done / max(elapsed, 1) * 60 eta_sec = (target - total_done) / max(rate / 60, 0.01) 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 except KeyboardInterrupt: print("\n⏹️ 用户中断") # 最终合并 gen.merge() total_time = time.time() - t_start final_rate = total_done / max(total_time, 1) * 60 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"{'='*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 __name__ == "__main__": main()