From f8fac5946bc633977a0e47a3dded976d54cb6893 Mon Sep 17 00:00:00 2001 From: shuiyu1123 Date: Sat, 20 Jun 2026 13:52:34 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20=E9=80=9A=E7=94=A8=E7=9F=A5=E8=AF=86?= =?UTF-8?q?=E9=A2=84=E8=AE=AD=E7=BB=83=E6=95=B0=E6=8D=AE=E7=94=9F=E6=88=90?= =?UTF-8?q?=E5=99=A8=2020=E5=A4=A7=E7=B1=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- gen_pretrain.py | 402 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 402 insertions(+) create mode 100644 gen_pretrain.py diff --git a/gen_pretrain.py b/gen_pretrain.py new file mode 100644 index 0000000..5f5dbd8 --- /dev/null +++ b/gen_pretrain.py @@ -0,0 +1,402 @@ +#!/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() \ No newline at end of file