大整理: 仓库重组 + 新词表 + 新README
- 清理仓库: 删除所有废弃的generate_*脚本和旧版数据文件 - 目录重组: - data/ → 训练数据 - scripts/ → 数据生成脚本 - templates/ → Web UI - 重命名: - train_v2.py → train.py - vocab_v3.json → vocab.json (BPE, 6.8k tokens) - 更新 train.py 指向新词表路径 - 重写 README.md 完整文档 - 更新 .gitignore
This commit is contained in:
shuiyu1123 committed
1 parent
6e0ced4772
commit
231af72de9
26 files changed
+124
-126
No files matched your search
@@ -0,0 +1,346 @@
|
||||
#!/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 # 配对后重置，等待下一个问
|
||||
|
||||
# 最少4轮才算有效（6~12轮的长对话）
|
||||
if len(turns) >= 4:
|
||||
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(6, 12)
|
||||
|
||||
prompt = f"""请以小狸的身份，模拟一段长时间的自然闲聊对话，至少{turn_count}轮。
|
||||
|
||||
开场话题：{topic_context}
|
||||
|
||||
要求：
|
||||
1. 小狸是8岁小女孩，说话简短自然
|
||||
2. 每轮包含"问"和"答"，小狸的回答1-2句话
|
||||
3. 对话要自然展开，一问一答慢慢深入，不要强行切话题
|
||||
4. 小狸的语气：害羞但可爱，不是机器人
|
||||
5. 至少{turn_count}轮，越多越好，让对话持续下去
|
||||
|
||||
请按以下格式输出（至少{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) >= 4:
|
||||
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 = 500000 # 目标 50 万条长多轮对话
|
||||
BATCH = 20 # 每批生成 20 条（每条更长，减少批量）
|
||||
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()
|
||||
Reference in new issue
Block a user