#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 海量通用数据蒸馏 v3 — 不设场景标签，纯自然对话，多策略并行 目标：500万+ 纯 (question, answer) 对 """ import json, requests, time, os, sys, random, re from concurrent.futures import ThreadPoolExecutor, as_completed sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1) MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) OUT_DIR = os.path.join(MODEL_DIR, 'general_outputs') os.makedirs(OUT_DIR, exist_ok=True) # ====== Astron 配置 ====== ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions" ASTRON_KEY = os.environ.get("ASTRON_KEY", "") ASTRON_MODEL = "astron-code-latest" # ====== 主题全覆盖（六大类 × 子类） ====== TOPICS = { "日常": [ "起床洗漱", "吃早餐", "上学路上", "上课", "课间休息", "午餐时间", "放学回家", "做作业", "晚饭", "睡前时光", "周末活动", "天气变化", "穿衣打扮", "零花钱", "家务劳动", "看电视", "玩玩具", "去公园", "逛超市", "去亲戚家", ], "情绪": [ "开心高兴", "难过伤心", "生气愤怒", "害怕恐惧", "紧张焦虑", "孤独寂寞", "委屈冤枉", "失望沮丧", "感动温暖", "兴奋激动", "无聊烦闷", "尴尬害羞", "后悔内疚", "羡慕嫉妒", "感激感恩", ], "社交": [ "交新朋友", "和好朋友玩", "和同学相处", "和老师交流", "和兄弟姐妹", "和爸爸妈妈", "和爷爷奶奶", "和邻居", "和陌生人", "和喜欢的人", "吵架和好", "被欺负", "被误解", "被表扬", "被批评", "帮助别人", "请求帮助", "分享东西", "邀请做客", "告别分开", ], "知识": [ "动物世界", "植物花草", "天文星空", "天气自然", "人体健康", "食物营养", "颜色形状", "数字计算", "汉字词语", "科学现象", "历史故事", "地理知识", "交通工具", "乐器音乐", "绘画手工", "运动体育", "节日习俗", "各国文化", "发明创造", "宇宙太空", ], "想象": [ "超能力", "魔法世界", "未来科技", "外星人", "恐龙时代", "童话故事", "变成动物", "飞行梦", "海底世界", "神秘岛屿", "时间旅行", "超级英雄", "会说话的玩具", "隐形衣", "变身药水", "机器人朋友", "彩虹桥", "星星糖果", "云朵房子", "梦中冒险", ], "成长": [ "学新本领", "克服困难", "不怕失败", "坚持努力", "自我肯定", "梦想未来", "职业理想", "好习惯", "改掉坏毛病", "学会分享", "学会等待", "学会道歉", "学会感恩", "学会勇敢", "学会理解", "交心谈话", "秘密分享", "烦恼倾诉", "人生道理", "温暖鼓励", ], } FLAT_TOPICS = [] for cat, subs in TOPICS.items(): for sub in subs: FLAT_TOPICS.append((cat, sub)) # ====== 多样化的开场方式 ====== OPENERS = [ "小狸", "嗨小狸", "小狸在吗", "嘿嘿", "那个…", "喂喂喂", "小狸小狸", "我想问你", "你知道吗", "你说", "我突然想到", "我跟你讲", "小狸你听我说", "问你个事", "我有个问题", "好无聊啊", "好开心呀", "好烦", "怎么办呀", "小狸你在干嘛", "我来了", "出来玩", "跟你说个事", ] # ====== 策略1: 直接生成纯问答对 ====== def make_direct_pairs_prompt(count=15): """无场景标签，直接生成纯问答对""" # 随机选3-5个主题混搭 selected = random.sample(FLAT_TOPICS, min(random.randint(3, 5), len(FLAT_TOPICS))) topics_str = "、".join([f"{c}({s})" for c, s in selected]) prompt = f"""请直接生成 {count} 条聊天问答对。 格式要求（严格按此格式）： 问：xxxxxxxxx | 答：xxxxxxxxx 用户的问题要：口语化、自然、像真人（5-25字） 小狸的回答要：8岁小女孩语气，可爱活泼（10-40字） 话题要多样：{topics_str} 示例： 问：你怕不怕黑呀？ | 答：有一点点…但是抱着小熊玩偶就不怕啦！ 问：我今天跑步摔了一跤 | 答：疼不疼呀？我给你吹吹就不疼了！ 现在生成 {count} 条：""" return prompt # ====== 策略2: 短对话（3-5轮） ====== def make_mini_chat_prompt(count=8): """微型多轮对话，每段3-5轮""" ops = random.sample(OPENERS, min(count, len(OPENERS))) prompt = f"""请生成 {count} 段微型聊天对话，每段3-5轮（一来一回算1轮）。 要求： - 每段一个独立话题，话题要多样（日常、情绪、知识、想象、社交、成长等） - 用户说话要特别自然、口语化，像真人随口说的 - 小狸像8岁小女孩，回答要可爱、有画面感 - 话题要自然延伸，不要生硬转折 格式例子： 话题：日常 用户：今天好冷啊 小狸：是呀！我穿了厚厚的羽绒服，像只小企鹅一样～ 用户：哈哈哈真的好像 小狸：（鼓起腮帮子）才不像呢！企鹅走路摇摇摆摆的，我走路可稳啦！ ... 现在生成 {count} 段，每段用"话题：xxx"开头：""" return prompt # ====== 策略3: 长对话（6-10轮自由展开） ====== def make_long_chat_prompt(): """一场长对话，话题自然展开""" topic = random.choice(FLAT_TOPICS) opener = random.choice(OPENERS) emotions = ["开心", "难过", "兴奋", "沮丧", "好奇", "担心", "轻松", "紧张"] emotion = random.choice(emotions) prompt = f"""请生成一段自然的日常聊天，{topic[0]}类-{topic[1]}话题。 用户开场：{opener} 用户当下的心情：{emotion} 要求： - 用户很自然地说出开场白，就像真人发微信一样 - 小狸正常回复，话题慢慢展开，不要硬聊 - 聊6-10轮（一问一答算一轮） - 每轮格式：用户：xxx\n小狸：xxx - 真实感第一，不要刻意 请直接输出对话，不要加说明：""" return prompt # ====== 策略4: 单轮问答（高吞吐） ====== def make_single_turn_prompt(batch_size=20): """纯单轮问答，批量生成""" selected = random.sample(FLAT_TOPICS, min(6, len(FLAT_TOPICS))) prompt = f"""生成 {batch_size} 条日常问答对。 要求： 1. 用户的话：7-25字，口语化，真实感强 2. 小狸回答：1-2句，8岁小女孩语气，可爱活泼 3. 覆盖多样化话题：{', '.join([s for _, s in selected[:6]])}（但不仅限于这些） 4. 每行一条，格式：问: ... | 答: ... 5. 不要重复相同的问题 示例： 问: 你喜不喜欢下雨天呀？ | 答: 喜欢！因为可以踩水坑，啪嗒啪嗒的超好玩！ 问: 我今天被同学起外号了 | 答: （歪着头）那…你喜欢这个外号吗？不喜欢的话要大声说出来哦！ 问: 你说世界上有外星人吗？ | 答: 我觉得有！说不定他们正在用望远镜看我们呢～ 生成：""" return prompt # ====== API调用 ====== def call_astron(prompt, temp=0.82, max_tokens=4096): for attempt in range(3): try: headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} payload = { "model": ASTRON_MODEL, "messages": [ {"role": "system", "content": "你是一个擅长写日常对话的助手。小狸是一个8岁的小女孩，说话可爱自然。"}, {"role": "user", "content": prompt} ], "temperature": temp, "max_tokens": max_tokens, } resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=150) if resp.status_code == 200: content = resp.json()['choices'][0]['message']['content'].strip() if content and len(content) > 20: return content else: time.sleep(0.5 + attempt) except Exception: time.sleep(1 + attempt) return None # ====== 解析 ====== def parse_pairs(text): """从"问：xxx | 答：xxx"格式解析问答对""" pairs = [] for line in text.split('\n'): line = line.strip() parts = re.split(r'\s*[|｜]\s*', line) if len(parts) == 2: q = re.sub(r'^问\s*[：:]\s*', '', parts[0]).strip() a = re.sub(r'^答\s*[：:]\s*', '', parts[1]).strip() if q and a and len(q) >= 2 and len(a) >= 2 and len(a) < 300 and len(q) < 200: pairs.append((q, a)) return pairs def parse_mini_chats(text): """解析微型多轮对话""" pairs = [] segments = re.split(r'话题[：:].*?\n', text) for seg in segments: seg = seg.strip() lines = seg.split('\n') q = None for line in lines: line = line.strip() if not line: continue q_match = re.match(r'^(用户|我)[：:]\s*(.*)', line) a_match = re.match(r'^(小狸|答)[：:]\s*(.*)', line) if q_match: q_text = q_match.group(2).strip() if q_text and len(q_text) > 1: q = q_text elif a_match and q: a_text = a_match.group(2).strip() if a_text and len(a_text) > 1 and len(a_text) < 300: pairs.append((q, a_text)) q = None return pairs def parse_long_chat(text): """解析长对话""" pairs = [] q = None for line in text.split('\n'): line = line.strip() if not line: continue q_match = re.match(r'^(用户|我)[：:]\s*(.*)', line) a_match = re.match(r'^(小狸)[：:]\s*(.*)', line) if q_match: q_text = q_match.group(2).strip() if q_text and len(q_text) > 1: q = q_text elif a_match and q: a_text = a_match.group(2).strip() if a_text and len(a_text) > 1 and len(a_text) < 300: pairs.append((q, a_text)) q = None return pairs # ====== 生成器（按策略分发） ====== GENERATORS = [ # (name, prompt_fn, parse_fn, batch_size_estimate, weight_for_diversity) ("direct", make_direct_pairs_prompt, parse_pairs, 15, 35), ("mini_chat", make_mini_chat_prompt, parse_mini_chats, 20, 25), ("long_chat", make_long_chat_prompt, parse_long_chat, 20, 20), ("single", make_single_turn_prompt, parse_pairs, 20, 20), ] def run_one_generation(): """运行一次生成任务（随机选一种策略）""" # 按权重选策略 weights = [w for _, _, _, _, w in GENERATORS] gen = random.choices(GENERATORS, weights=weights, k=1)[0] name, prompt_fn, parse_fn, batch_est, _ = gen prompt = prompt_fn() result = call_astron(prompt, temp=random.uniform(0.75, 0.92), max_tokens=4096 if 'long' not in name else 6144) if not result: return None pairs = parse_fn(result) if not pairs: return None return pairs def main(): print("=" * 65) print(" 海量通用数据蒸馏 v3") print(" 不设场景标签 | 6大类120+子话题 | 4种策略") print("=" * 65) # 加载已有进度 all_pairs = [] existing_keys = set() progress_file = os.path.join(OUT_DIR, 'all_general_progress.json') if os.path.exists(progress_file): existing = json.load(open(progress_file, encoding='utf-8')) all_pairs = existing for p in existing: existing_keys.add((p.get('question',''), p.get('answer',''))) print(f"已有进度: {len(all_pairs):,} 对") # ====== 参数 ====== TARGET = 2000000 # 目标200万对 BATCH = 40 # 每批40个任务 MAX_WORKERS = 8 # 8并发（比之前5更猛） CHECKPOINT = 5000 # 每5000对保存 to_generate = TARGET - len(all_pairs) total_batches = (to_generate + BATCH - 1) // BATCH print(f"目标: {TARGET:,} 对") print(f"已有: {len(all_pairs):,} 对") print(f"还需: {to_generate:,} 对") print(f"批次: {total_batches:,} 批 ({BATCH}任务/批, {MAX_WORKERS}并发)") # 策略统计 strategy_stats = {g[0]: {'runs': 0, 'pairs': 0} for g in GENERATORS} t_start = time.time() last_checkpoint = len(all_pairs) for batch_idx in range(total_batches): t0 = time.time() # 并发 with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: futures = [executor.submit(run_one_generation) for _ in range(BATCH)] batch_pairs = [] for future in as_completed(futures): try: r = future.result() if r: batch_pairs.extend(r) except: pass # 去重 new_pairs = [] for q, a in batch_pairs: key = (q, a) if key not in existing_keys: existing_keys.add(key) new_pairs.append({'question': q, 'answer': a}) all_pairs.extend(new_pairs) elapsed = time.time() - t0 total_elapsed = time.time() - t_start rate = len(all_pairs) / max(total_elapsed, 1) * 60 # 估算 pairs_this_batch = len(new_pairs) total_raw = len(batch_pairs) dedup_rate = pairs_this_batch / max(total_raw, 1) * 100 print(f"[{batch_idx+1}/{total_batches}] " f"+{pairs_this_batch}对(有效{total_raw}/{dedup_rate:.0f}%) " f"累计{len(all_pairs):,}/{TARGET:,} " f"{rate:.0f}对/分 " f"剩余{(TARGET-len(all_pairs))/max(rate,0.1):.0f}分") # 定期保存 if len(all_pairs) - last_checkpoint >= CHECKPOINT: with open(progress_file, 'w', encoding='utf-8') as f: json.dump(all_pairs, f, ensure_ascii=False) last_checkpoint = len(all_pairs) # 也保存到数据目录一份 merged_path = os.path.join(MODEL_DIR, 'data', 'general_ongoing.json') with open(merged_path, 'w', encoding='utf-8') as f: json.dump(all_pairs, f, ensure_ascii=False) # 速率统计 stable_rate = len(all_pairs) / max(time.time() - t_start, 1) * 60 print(f" 💾 保存 ({len(all_pairs):,} 对, {stable_rate:.0f}对/分)") # ====== 完成 ====== total_elapsed = time.time() - t_start with open(progress_file, 'w', encoding='utf-8') as f: json.dump(all_pairs, f, ensure_ascii=False) print(f"\n{'='*65}") print(f"✅ 生成完成!") print(f" 总对数: {len(all_pairs):,}") print(f" 用时: {total_elapsed/60:.1f} 分") print(f" 平均速率: {len(all_pairs)/max(total_elapsed,1)*60:.0f} 对/分") print(f" 保存: {progress_file}") print(f"{'='*65}") if __name__ == '__main__': main()