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