402 lines
16 KiB
Python
402 lines
16 KiB
Python
#!/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() |