337 lines
14 KiB
Python
337 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
海量真实多轮对话生成 — 模拟真实用户，8~15轮长对话
|
||
"""
|
||
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)
|
||
|
||
from scripts.api_correction import batch_generate
|
||
|
||
MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
# ====== Astron 配置（复用 api_correction 的 key） ======
|
||
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"
|
||
|
||
# ====== 用户场景画像（模拟真实用户） ======
|
||
USER_SCENARIOS = [
|
||
# 小小孩子（3-5岁）
|
||
{"age": "3-5岁", "style": "说话简单，喜欢问为什么，爱说叠词", "tone": "奶声奶气"},
|
||
# 小孩子（6-9岁）
|
||
{"age": "6-9岁", "style": "小学生，爱分享学校的事，喜欢动画片和游戏", "tone": "活泼天真"},
|
||
# 大孩子（10-14岁）
|
||
{"age": "10-14岁", "style": "有点小叛逆，有自己的想法，会吐槽", "tone": "有点酷"},
|
||
# 妈妈/家长
|
||
{"age": "成年", "style": "关心孩子，会问小狸吃饭睡觉学习的事", "tone": "温柔耐心"},
|
||
# 孤单的人
|
||
{"age": "成年", "style": "有点孤独，需要陪伴，会把小狸当倾诉对象", "tone": "低落温柔"},
|
||
# 高中生/大学生
|
||
{"age": "16-22岁", "style": "有压力，会聊学业、未来、人际关系", "tone": "成熟但需要安慰"},
|
||
# 好奇宝宝
|
||
{"age": "不限", "style": "什么都想知道，十万个为什么", "tone": "充满好奇"},
|
||
# 伤心的人
|
||
{"age": "不限", "style": "心情不好，需要安慰和陪伴", "tone": "低沉脆弱"},
|
||
]
|
||
|
||
# ====== 真实对话种子（更自然的开场） ======
|
||
REAL_SEEDS = [
|
||
# 日常打招呼（多种自然说法）
|
||
"小狸你在吗", "嘿嘿小狸", "小狸小狸快出来", "小狸你在干嘛呀",
|
||
"我来找你玩啦", "小狸今天乖不乖", "你在做什么呢",
|
||
# 情绪分享（真实用户会说的话）
|
||
"我今天好不开心", "我开心死了哈哈哈", "我好无聊啊没人陪我",
|
||
"今天倒霉死了", "我今天遇到了好事", "我有点难过",
|
||
"我心情不好你能陪我吗", "我好想你啊", "我好孤单",
|
||
"我害怕", "我有点紧张明天要考试", "我好烦",
|
||
# 日常分享
|
||
"今天我妈做了红烧肉", "我们家养了一只猫超可爱", "我今天考试考了98分",
|
||
"我感冒了好难受", "我今天被老师表扬了", "我和朋友吵架了",
|
||
"我昨天做梦梦到你了", "我学会骑自行车了！", "我今天摔了一跤好疼",
|
||
"学校要考试了好紧张", "放假了好开心", "又要开学了好痛苦",
|
||
"今天下雨了不能出去玩", "今天天气好好我们去玩吧",
|
||
"我妈妈给我买了个新玩具", "我养的花开了",
|
||
# 求助/询问
|
||
"我不小心把杯子打碎了怎么办", "我作业好多写不完",
|
||
"我睡不着你能陪我聊天吗", "我牙好疼",
|
||
"妈妈骂我了", "我被老师批评了",
|
||
"我迷路了怎么办", "我找不到我的水杯了",
|
||
# 好奇/知识
|
||
"为什么天是蓝的呀", "鱼会睡觉吗", "为什么会有影子",
|
||
"星星为什么会闪", "海为什么是咸的", "为什么要上学",
|
||
"为什么大人总说不行", "什么是爱", "什么是长大",
|
||
# 表达好感的
|
||
"小狸我最喜欢你了", "你好可爱呀", "我们是最好的朋友对吧",
|
||
"我送给你一个礼物", "小狸你真聪明", "你真好",
|
||
# 关于小狸本身的
|
||
"你几岁啦", "你喜欢什么", "你有爸爸妈妈吗",
|
||
"你住在哪里呀", "你喜欢吃什么", "你最好的朋友是谁",
|
||
"你上学吗", "你的梦想是什么", "你平时喜欢做什么",
|
||
# 奇怪/无厘头
|
||
"如果我能飞就好了", "我想变成一只猫", "你觉得外星人存在吗",
|
||
"假如你会魔法你想做什么", "你相信世界上有奥特曼吗",
|
||
# 夜间/睡前
|
||
"我睡不着", "陪我聊到睡着好不好", "给我讲个睡前故事",
|
||
"晚安前陪我说说话", "我怕黑",
|
||
# 亲情家庭
|
||
"我爸妈吵架了怎么办", "我想我奶奶了", "我好久没见爸爸了",
|
||
"妈妈生病了", "我弟弟好烦总是抢我东西",
|
||
# 节假日
|
||
"过年啦新年快乐", "今天是我生日", "圣诞节到了",
|
||
"六一儿童节快乐", "明天要去春游好开心",
|
||
]
|
||
|
||
CONTEXTS = [
|
||
"晚上睡前躺在床上", "放学回家的路上",
|
||
"周末在家没事做", "刚吃完饭坐着聊天",
|
||
"在公园里坐着", "下雨天窝在家里",
|
||
"一个人在房间", "在院子里玩的时候",
|
||
"刚写完作业休息", "早上起床的时候",
|
||
"在操场上看别人玩", "坐车去玩的路上",
|
||
"在图书馆里", "在外面吃饭等上菜",
|
||
"一个人发呆的时候", "看着窗外的时候",
|
||
]
|
||
|
||
def make_realistic_multi_prompt():
|
||
"""构造非常逼真的长多轮对话提示 — 一次生成4条"""
|
||
parts = []
|
||
for _ in range(4):
|
||
scenario = random.choice(USER_SCENARIOS)
|
||
opener = random.choice(REAL_SEEDS)
|
||
context = random.choice(CONTEXTS)
|
||
min_turns = random.randint(6, 8)
|
||
max_turns = random.randint(10, 14)
|
||
|
||
parts.append(f"""=== 对话 {_+1} ===
|
||
场景：{context}
|
||
用户：{scenario['age']}，{scenario['style']}，{scenario['tone']}
|
||
开场：{opener}
|
||
对话：
|
||
用户：{opener}
|
||
小狸：""")
|
||
|
||
prompt = f"""请生成 {len(parts)} 段独立的日常聊天对话，每段 {min_turns}~{max_turns} 轮。
|
||
|
||
重要：每段对话必须是一个完整的聊天，话题自然展开，不要强行结束。
|
||
|
||
{"".join(parts)}
|
||
|
||
严格要求：
|
||
1. 每个=== 对话 ===之后的段落实例必须写完整，不要留空
|
||
2. 用户提问必须口语化、自然，像真人在说话
|
||
3. 小狸回答1-2句话，像8岁小朋友
|
||
4. 话题要延续上一轮，逐渐深入延伸
|
||
5. 用户要有真实反应（惊讶、追问、分享等）
|
||
6. 用"问："/"答："格式
|
||
7. 输出不能省略任何一轮
|
||
"""
|
||
return prompt
|
||
|
||
def call_astron_multi(prompt, temperature=0.85):
|
||
for attempt in range(3):
|
||
try:
|
||
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"}
|
||
payload = {
|
||
"model": ASTRON_MODEL,
|
||
"messages": [
|
||
{"role": "system", "content": "你是一个擅长写真实对话的助手。"},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
"temperature": temperature,
|
||
"max_tokens": 8192,
|
||
}
|
||
resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=120)
|
||
if resp.status_code == 200:
|
||
content = resp.json()['choices'][0]['message']['content'].strip()
|
||
if content and len(content) > 50:
|
||
return content
|
||
else:
|
||
time.sleep(1)
|
||
except Exception:
|
||
time.sleep(2)
|
||
return None
|
||
|
||
def parse_multi_turn(text):
|
||
"""解析多段对话（支持=== 对话 N ===分隔，也支持纯轮次）"""
|
||
lines = text.strip().split('\n')
|
||
conversations = []
|
||
current_q = None
|
||
current_turns = []
|
||
|
||
for line in lines:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
|
||
# 检测新对话开始标记
|
||
if re.match(r'===?\s*对话\s*\d+\s*===?', line):
|
||
if current_turns and len(current_turns) >= 4:
|
||
conversations.append(current_turns)
|
||
current_turns = []
|
||
current_q = None
|
||
continue
|
||
|
||
# 跳过设定行（场景/用户/开场）
|
||
if line.startswith(('场景', '用户', '开场', '对话', '===')):
|
||
if line.startswith(('场景', '用户', '开场')):
|
||
continue
|
||
continue
|
||
|
||
# 用户提问
|
||
q_match = re.match(r'^(用户|我|问|问题|小孩|小朋友|孩子|朋友)\s*[：:]\s*(.*)', line)
|
||
a_match = re.match(r'^(小狸|答|回答|机器人)\s*[：:]\s*(.*)', line)
|
||
|
||
if q_match:
|
||
q_text = q_match.group(2).strip()
|
||
if q_text and len(q_text) > 1:
|
||
current_q = q_text
|
||
elif a_match and current_q:
|
||
a_text = a_match.group(2).strip()
|
||
if a_text and len(a_text) > 1 and len(a_text) < 300:
|
||
current_turns.append((current_q, a_text))
|
||
current_q = None
|
||
|
||
# 保存最后一组
|
||
if current_turns and len(current_turns) >= 4:
|
||
conversations.append(current_turns)
|
||
|
||
# 过滤：只保留>=4轮的对话，且累积轮次>=8
|
||
valid = [c for c in conversations if len(c) >= 4]
|
||
total_turns = sum(len(c) for c in valid)
|
||
if total_turns >= 8 and valid:
|
||
return valid
|
||
return None
|
||
|
||
def make_turns_list_text(turns):
|
||
"""将turns转成用户可见的对话文本"""
|
||
lines = []
|
||
for i, (q, a) in enumerate(turns):
|
||
lines.append(f"问：{q}")
|
||
lines.append(f"答：{a}")
|
||
return "\n".join(lines)
|
||
|
||
def generate_one_multi():
|
||
"""生成多条多轮对话（一次API调用返回多条）"""
|
||
prompt = make_realistic_multi_prompt()
|
||
result = call_astron_multi(prompt, temperature=random.uniform(0.75, 0.95))
|
||
if result:
|
||
convs = parse_multi_turn(result)
|
||
if convs:
|
||
items = []
|
||
for turns in convs:
|
||
items.append({
|
||
'turns': len(turns),
|
||
'turns_list': turns,
|
||
'text': make_turns_list_text(turns),
|
||
'source': 'astron_multi_v2',
|
||
})
|
||
return items
|
||
return None
|
||
|
||
def main():
|
||
print("="*60)
|
||
print(" 海量真实多轮对话生成 v2")
|
||
print(" 生成超长、超真实的多轮聊天")
|
||
print("="*60)
|
||
|
||
out_dir = os.path.join(MODEL_DIR, 'multiturn_outputs')
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
# 检查已有进度
|
||
all_results = []
|
||
progress_file = os.path.join(out_dir, 'all_multiturn_progress.json')
|
||
if os.path.exists(progress_file):
|
||
all_results = json.load(open(progress_file, encoding='utf-8'))
|
||
print(f"已有进度: {len(all_results)} 条")
|
||
|
||
# ====== 参数 ======
|
||
TARGET = 500000 # 目标50万条
|
||
BATCH = 30 # 每批30条
|
||
MAX_WORKERS = 5 # Astron 5并发（稳定不超限）
|
||
CHECKPOINT = 1000 # 每1000条保存一次进度
|
||
|
||
to_generate = TARGET - len(all_results)
|
||
total_batches = (to_generate + BATCH - 1) // BATCH
|
||
|
||
print(f"目标: {TARGET:,} 条")
|
||
print(f"已有: {len(all_results):,} 条")
|
||
print(f"还需: {to_generate:,} 条")
|
||
print(f"批次: {total_batches:,} 批 (每批{BATCH}条, {MAX_WORKERS}并发)")
|
||
print(f"{'='*60}\n")
|
||
|
||
t_start = time.time()
|
||
last_checkpoint = len(all_results)
|
||
|
||
for batch_idx in range(total_batches):
|
||
t0 = time.time()
|
||
|
||
# 并发生成
|
||
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||
futures = [executor.submit(generate_one_multi) for _ in range(BATCH)]
|
||
batch_results = []
|
||
for future in as_completed(futures):
|
||
try:
|
||
r = future.result()
|
||
if r:
|
||
batch_results.extend(r)
|
||
except Exception as e:
|
||
pass
|
||
|
||
# 去重
|
||
existing_texts = set(r.get('text', '')[:200] for r in all_results)
|
||
deduped = []
|
||
for r in batch_results:
|
||
key = r.get('text','')[:200]
|
||
if key not in existing_texts:
|
||
existing_texts.add(key)
|
||
deduped.append(r)
|
||
|
||
all_results.extend(deduped)
|
||
|
||
elapsed = time.time() - t0
|
||
total_elapsed = time.time() - t_start
|
||
rate = len(all_results) / max(total_elapsed, 1) * 60
|
||
eta = (TARGET - len(all_results)) / max(rate, 0.1) / 60 if rate > 0 else 0
|
||
|
||
# 统计
|
||
valid = len(batch_results)
|
||
avg_turns = sum(r['turns'] for r in deduped) / max(len(deduped), 1)
|
||
success_rate = valid / BATCH * 100
|
||
|
||
print(f"[{batch_idx+1}/{total_batches}] "
|
||
f"+{len(deduped)}条(有效{valid}/{success_rate:.0f}%) "
|
||
f"累计{len(all_results):,}/{TARGET:,} "
|
||
f"均{avg_turns:.1f}轮 "
|
||
f"{rate:.0f}条/分 "
|
||
f"剩余{eta:.0f}分")
|
||
|
||
# 定期保存进度
|
||
if len(all_results) - last_checkpoint >= CHECKPOINT:
|
||
with open(progress_file, 'w', encoding='utf-8') as f:
|
||
json.dump(all_results, f, ensure_ascii=False)
|
||
last_checkpoint = len(all_results)
|
||
print(f" 💾 已保存进度 ({len(all_results):,} 条)")
|
||
|
||
# ====== 完成 ======
|
||
total_elapsed = time.time() - t_start
|
||
with open(progress_file, 'w', encoding='utf-8') as f:
|
||
json.dump(all_results, f, ensure_ascii=False)
|
||
|
||
avg_turns_all = sum(r['turns'] for r in all_results) / max(len(all_results), 1)
|
||
total_turns = sum(r['turns'] for r in all_results)
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f"✅ 生成完成!")
|
||
print(f" 总条数: {len(all_results):,}")
|
||
print(f" 平均轮数: {avg_turns_all:.1f}")
|
||
print(f" 总对话轮次: {total_turns:,}")
|
||
print(f" 用时: {total_elapsed/60:.1f} 分")
|
||
print(f" 速率: {len(all_results)/max(total_elapsed,1)*60:.0f} 条/分")
|
||
print(f" 保存: {progress_file}")
|
||
print(f"{'='*60}")
|
||
|
||
if __name__ == '__main__':
|
||
main() |