345 lines
13 KiB
Python
345 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
多轮对话生成 — 用 Astron 造 2~5 轮小狸闲聊对话
|
||
5并发，无限量，增量保存防丢
|
||
"""
|
||
import json, requests, time, os, sys, random, re
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
|
||
|
||
# ====== Astron 配置（和 api_correction 一致） ======
|
||
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"
|
||
|
||
# ====== 小狸基础提示词 ======
|
||
SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。
|
||
|
||
你的人设：
|
||
- 8岁女孩子，有点害羞但很可爱
|
||
- 说话简单自然，符合小朋友的语气
|
||
- 回答问题要简短，1-2句话就够了
|
||
- 不要用"喵"结尾，不要过度卖萌
|
||
|
||
回答要求：
|
||
1. 语气自然，像真实的小朋友说话
|
||
2. 不懂就说"不知道"
|
||
3. 回答简洁但要说清楚"""
|
||
|
||
# ====== 多轮对话种子（开场话题池） ======
|
||
SEEDS = [
|
||
# 日常问候
|
||
"你好呀", "你叫什么名字", "你今天开心吗", "你在干嘛呢",
|
||
"你好可爱呀", "我来看你啦", "你今天心情怎么样",
|
||
# 闲聊话题
|
||
"你喜欢什么颜色", "你最喜欢什么动物", "你爱吃什么零食",
|
||
"你最喜欢什么玩具", "你最喜欢什么动画片",
|
||
"你最好的朋友是谁", "幼儿园里好玩吗",
|
||
"你喜欢唱歌吗", "你会跳舞吗",
|
||
# 情绪表达
|
||
"我今天好开心", "我有点不开心", "我好无聊啊",
|
||
"我遇到了一件好事", "我今天有点难过",
|
||
"我好想你呀", "你能不能陪我聊聊天",
|
||
# 好奇问题
|
||
"为什么天是蓝的", "为什么月亮跟着我走",
|
||
"世界上最大的动物是什么", "你会讲故事吗",
|
||
"你会数数吗", "1+1等于几",
|
||
"你喜欢上学吗", "你今年几岁啦",
|
||
# 分享日常
|
||
"我今天去公园玩了", "我养了一只小猫",
|
||
"我今天考试考了100分", "我学会骑自行车了",
|
||
"妈妈给我买了个新玩具", "我今天摔了一跤",
|
||
# 晚安/告别
|
||
"我要睡觉了", "明天见", "晚安啦",
|
||
"我要去吃饭了", "下次再聊哦",
|
||
]
|
||
|
||
# 话题扩展 — 组合出更多种子
|
||
PREFIXES = ["小狸", "小狸小狸", "亲爱的小狸", "嘿嘿小狸", "小狸呀"]
|
||
SUFFIXES = ["呀", "吗", "呢", "", "呀", "吗", "呢", ""]
|
||
|
||
def expand_seeds(base, max_expand=5):
|
||
"""种子话题 → 多样化的开场白"""
|
||
expanded = []
|
||
for s in base:
|
||
expanded.append(s)
|
||
if '吗' not in s and not s.endswith('呀') and not s.endswith('呢'):
|
||
for _ in range(2):
|
||
prefix = random.choice(PREFIXES) if random.random() > 0.4 else ''
|
||
suffix = random.choice(SUFFIXES) if random.random() > 0.5 else ''
|
||
combo = f"{prefix}{s}{suffix}".strip()
|
||
if combo and combo not in expanded:
|
||
expanded.append(combo)
|
||
random.shuffle(expanded)
|
||
return expanded[:len(base) * max_expand]
|
||
|
||
|
||
def call_astron_multi(prompt, temperature=0.85):
|
||
"""调用 Astron 生成一条多轮对话（2~5轮）"""
|
||
for attempt in range(3):
|
||
try:
|
||
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"}
|
||
payload = {
|
||
"model": ASTRON_MODEL,
|
||
"messages": [
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
"temperature": temperature,
|
||
"max_tokens": 8192,
|
||
}
|
||
resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60)
|
||
if resp.status_code == 200:
|
||
content = resp.json()['choices'][0]['message']['content'].strip()
|
||
if content and len(content) > 10:
|
||
return content
|
||
else:
|
||
time.sleep(1)
|
||
except Exception:
|
||
time.sleep(2)
|
||
return None
|
||
|
||
|
||
def parse_multi_turn(text):
|
||
"""
|
||
解析 Astron 返回的多轮对话文本为 [(q, a), (q, a), ...]
|
||
期望返回格式（至少2轮对话）：
|
||
问：xxx
|
||
答：xxx
|
||
问：xxx
|
||
答：xxx
|
||
|
||
或
|
||
用户：xxx
|
||
小狸：xxx
|
||
用户：xxx
|
||
小狸：xxx
|
||
"""
|
||
lines = text.strip().split('\n')
|
||
turns = []
|
||
current_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:
|
||
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) < 200:
|
||
turns.append((current_q, a_text))
|
||
current_q = None # 配对后重置，等待下一个问
|
||
|
||
# 最少2轮才算有效
|
||
if len(turns) >= 2:
|
||
return turns
|
||
return None
|
||
|
||
|
||
def format_training_sample(turns):
|
||
"""
|
||
将多轮对话 [('q1','a1'), ('q2','a2'), ...] 格式化为训练样本
|
||
|
||
格式：<SOS>q1<SEP><START>a1<END>q2<SEP><START>a2<END><EOS>
|
||
"""
|
||
parts = []
|
||
for i, (q, a) in enumerate(turns):
|
||
if i == 0:
|
||
parts.append(f"<SOS>{q}<SEP><START>{a}<END>")
|
||
else:
|
||
parts.append(f"{q}<SEP><START>{a}<END>")
|
||
parts.append("<EOS>")
|
||
return "".join(parts)
|
||
|
||
|
||
def make_multi_prompt():
|
||
"""
|
||
构造多轮对话生成提示词
|
||
Astron 每次生成一条完整的4~6轮对话
|
||
"""
|
||
# 随机选开场话题（2~3个组合，让对话有展开）
|
||
topic_count = random.randint(1, 3)
|
||
topics = random.sample(all_seeds, min(topic_count, len(all_seeds)))
|
||
topic_context = "、".join(topics)
|
||
|
||
turn_count = random.randint(2, 5)
|
||
|
||
prompt = f"""请以小狸的身份，模拟一段{turn_count}轮的自然闲聊对话。
|
||
|
||
开场话题：{topic_context}
|
||
|
||
要求：
|
||
1. 小狸是8岁小女孩，说话简短自然
|
||
2. 每轮包含"问"和"答"，小狸的回答1-2句话
|
||
3. 对话要自然展开，不要强行切话题
|
||
4. 小狸的语气：害羞但可爱，不是机器人
|
||
|
||
请按以下格式输出（至少{turn_count}轮）：
|
||
问：（用户提问）
|
||
答：（小狸回答）
|
||
问：（用户续问）
|
||
答：（小狸回答）
|
||
..."""
|
||
return prompt
|
||
|
||
|
||
def generate_one_multi():
|
||
"""生成一条多轮对话"""
|
||
prompt = make_multi_prompt()
|
||
result = call_astron_multi(prompt, temperature=random.uniform(0.8, 1.0))
|
||
if result:
|
||
turns = parse_multi_turn(result)
|
||
if turns and len(turns) >= 2:
|
||
sample = format_training_sample(turns)
|
||
return {
|
||
'source': 'multiturn_astron',
|
||
'turns': len(turns),
|
||
'sample': sample,
|
||
'turns_list': turns,
|
||
'raw': result[:100]
|
||
}
|
||
return None
|
||
|
||
|
||
def main():
|
||
global all_seeds
|
||
|
||
model_dir = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
# 加载现有训练数据的问句作为种子
|
||
train_path = os.path.join(model_dir, 'xiaoli_train_v2.json')
|
||
with open(train_path, encoding='utf-8') as f:
|
||
train_data = json.load(f)
|
||
|
||
# 从现有数据中随机抽取问句作为种子
|
||
existing_qs = [item.get('question','') for item in train_data if item.get('question','')]
|
||
random.shuffle(existing_qs)
|
||
seed_from_data = existing_qs[:500]
|
||
|
||
# 合并种子话题
|
||
all_seeds = expand_seeds(SEEDS, max_expand=10) + seed_from_data
|
||
random.shuffle(all_seeds)
|
||
|
||
print(f"种子话题: {len(all_seeds)} 个 (基础{len(SEEDS)}+数据抽取{len(seed_from_data)})")
|
||
|
||
# 输出目录
|
||
out_dir = os.path.join(model_dir, 'multiturn_outputs')
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
# 检查已有进度
|
||
all_results = []
|
||
interim_files = sorted([f for f in os.listdir(out_dir)
|
||
if f.startswith('mt_batch_') and f.endswith('.json')])
|
||
if interim_files:
|
||
for fname in interim_files:
|
||
with open(os.path.join(out_dir, fname), encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
all_results.extend(data)
|
||
print(f"已有中间结果: {len(all_results)} 条 ({len(interim_files)} 批)")
|
||
|
||
# ====== 参数 ======
|
||
TARGET = 20000 # 目标生成 2 万条多轮对话
|
||
BATCH = 50 # 每批生成 50 条
|
||
MAX_WORKERS = 5 # Astron 5并发
|
||
TEMPERATURE_RANGE = (0.80, 1.0)
|
||
|
||
to_generate = TARGET - len(all_results)
|
||
total_batches = (to_generate + BATCH - 1) // BATCH
|
||
|
||
print(f"目标: {TARGET} 条, 已有: {len(all_results)} 条, 还需: {to_generate} 条")
|
||
print(f"批次: {total_batches} 批 (每批{BATCH}条, {MAX_WORKERS}并发)")
|
||
print(f"{'='*60}\n")
|
||
|
||
t_start = time.time()
|
||
|
||
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.append(r)
|
||
except Exception:
|
||
pass
|
||
|
||
# 去重（按 sample 内容）
|
||
existing_samples = set(r['sample'] for r in all_results)
|
||
deduped = []
|
||
for r in batch_results:
|
||
if r['sample'] not in existing_samples:
|
||
existing_samples.add(r['sample'])
|
||
deduped.append(r)
|
||
|
||
all_results.extend(deduped)
|
||
|
||
# 保存本批
|
||
elapsed = time.time() - t0
|
||
batch_file = os.path.join(out_dir, f'mt_batch_{batch_idx + 1:04d}.json')
|
||
with open(batch_file, 'w', encoding='utf-8') as f:
|
||
json.dump(deduped, f, ensure_ascii=False, indent=2)
|
||
|
||
# 追加保存总进度
|
||
progress_file = os.path.join(out_dir, 'all_multiturn_progress.json')
|
||
with open(progress_file, 'w', encoding='utf-8') as f:
|
||
json.dump(all_results, f, ensure_ascii=False, indent=2)
|
||
|
||
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
|
||
|
||
# 输出统计
|
||
avg_turns = sum(r['turns'] for r in deduped) / max(len(deduped), 1)
|
||
print(f"批次 {batch_idx+1}/{total_batches} | "
|
||
f"+{len(deduped)} 条 (去重后) | "
|
||
f"累计 {len(all_results)}/{TARGET} | "
|
||
f"均轮 {avg_turns:.1f} | "
|
||
f"{rate:.0f}条/分 | "
|
||
f"预计剩余 {eta:.1f}分")
|
||
|
||
# ====== 完成总结 ======
|
||
total_elapsed = time.time() - t_start
|
||
print(f"\n{'='*60}")
|
||
print(f"多轮对话生成完成!")
|
||
print(f"生成: {len(all_results)} 条 (去重后)")
|
||
avg_turns2 = sum(r['turns'] for r in all_results) / max(len(all_results), 1)
|
||
print(f"平均轮数: {avg_turns2:.1f}")
|
||
print(f"总对话轮次: {sum(r['turns'] for r in all_results)}")
|
||
print(f"用时: {total_elapsed/60:.1f} 分钟")
|
||
print(f"速率: {len(all_results)/max(total_elapsed,1)*60:.0f} 条/分")
|
||
|
||
# ====== 保存最终结果（独立存放，不追加到训练集） ======
|
||
final_file = os.path.join('/data/user/work', 'multiturn_final.json')
|
||
with open(final_file, 'w', encoding='utf-8') as f:
|
||
json.dump(all_results, f, ensure_ascii=False, indent=2)
|
||
print(f"\n最终结果已保存: {final_file} ({len(all_results)} 条)")
|
||
print(f"如需后续使用，可手动合并到训练集")
|
||
|
||
# ====== 保存多轮样本供训练时直接使用（含标记格式） ======
|
||
mt_samples_file = os.path.join(model_dir, 'multiturn_samples_v1.json')
|
||
with open(mt_samples_file, 'w', encoding='utf-8') as f:
|
||
samples = [{'turns': r['turns'], 'sample': r['sample']} for r in all_results]
|
||
json.dump(samples, f, ensure_ascii=False, indent=2)
|
||
print(f" ✅ 多轮样本保存: multiturn_samples_v1.json ({len(samples)} 条)")
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f"全部完成! 新增 {len(all_results)} 条多轮对话")
|
||
print(f"{'='*60}")
|
||
|
||
if __name__ == '__main__':
|
||
main() |