197 lines
8.4 KiB
Python
197 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
小狸 AI - 持续蒸馏守护脚本
|
||
============================
|
||
每轮生成 1000 条新 QA 对，去重后追加到训练集。
|
||
5 并发调用 Astron API，不超限。
|
||
|
||
使用方式:
|
||
python3 distill_daemon.py # 跑一轮
|
||
python3 distill_daemon.py --loop # 持续循环，每轮间隔 60 秒
|
||
python3 distill_daemon.py --count 5000 # 指定本轮目标条数
|
||
"""
|
||
import sys, os, json, time, random, argparse
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
from api_correction import call_astron, batch_generate
|
||
|
||
# ── 种子话题（15 大类，持续扩展） ──────────────────────
|
||
SEED_TOPICS = {
|
||
"日常问候": ["今天开心吗", "在干嘛呢", "吃了吗", "你好呀", "早上好"],
|
||
"情绪感受": ["我有点难过", "今天好开心", "我好无聊", "有点害怕", "好兴奋"],
|
||
"学校生活": ["上学好玩吗", "今天学了什么", "作业写完了吗", "喜欢上什么课", "考试考了多少"],
|
||
"朋友关系": ["你有好朋友吗", "怎么交朋友", "朋友吵架了", "一起玩什么", "好朋友搬家了"],
|
||
"家庭生活": ["喜欢爸爸妈妈吗", "有兄弟姐妹吗", "帮妈妈做家务", "周末去哪玩", "家里养宠物了吗"],
|
||
"兴趣爱好": ["喜欢唱歌吗", "会画画吗", "喜欢什么动画片", "爱看书吗", "喜欢什么游戏"],
|
||
"奇思妙想": ["为什么天是蓝的", "星星会掉下来吗", "鱼在水里怎么呼吸", "花为什么会开", "月亮上有兔子吗"],
|
||
"安慰鼓励": ["失败了怎么办", "被人欺负了", "好紧张怎么办", "觉得自己不够好", "想放弃了"],
|
||
"节日话题": ["最喜欢什么节日", "过年开心吗", "生日怎么过", "圣诞节收到什么礼物", "儿童节去哪玩"],
|
||
"自然科普": ["下雨了怎么办", "为什么会有彩虹", "树叶为什么变黄", "雪是怎么来的", "太阳下山去哪了"],
|
||
"童话故事": ["喜欢什么童话", "听过灰姑娘吗", "美人鱼的故事", "孙悟空厉害吗", "小红帽的故事"],
|
||
"美食话题": ["喜欢吃什么", "讨厌什么蔬菜", "冰淇淋好吃吗", "会做吃的吗", "妈妈做的饭好吃吗"],
|
||
"动物朋友": ["喜欢猫还是狗", "小兔子可爱吗", "小鸟会说话吗", "去过动物园吗", "恐龙是真的吗"],
|
||
"睡前话题": ["晚安怎么说", "做噩梦了", "睡不着怎么办", "想要抱抱", "讲个睡前故事"],
|
||
"成长烦恼": ["不想长大", "怕黑怎么办", "一个人在家害怕", "摔跤了很疼", "被批评了好难过"],
|
||
}
|
||
|
||
# 所有话题平铺
|
||
ALL_TOPICS = []
|
||
for cat, items in SEED_TOPICS.items():
|
||
ALL_TOPICS.extend(items)
|
||
|
||
# 对话风格模板
|
||
STYLE_TEMPLATES = [
|
||
"请用8岁小女孩小狸的语气回答，简短自然，1-2句话。",
|
||
"小狸是个有点害羞但很可爱的小女孩，请模仿她的语气回答。",
|
||
"像小朋友一样简单回答就好，要用小狸的口吻哦。",
|
||
"用小狸的口吻回答，要可爱但不要太刻意，简短就好。",
|
||
"你是一个8岁小女孩，请用小孩子的方式回答这个问题。",
|
||
]
|
||
|
||
|
||
def generate_seed_questions(count: int = 500) -> list:
|
||
"""从种子话题生成问题列表（随机组合+衍生）"""
|
||
questions = []
|
||
|
||
# 1. 直接使用种子话题
|
||
for topic in random.sample(ALL_TOPICS, min(count // 3, len(ALL_TOPICS))):
|
||
questions.append(topic)
|
||
|
||
# 2. 组合衍生（如 "今天开心吗 + 考试考了多少" → "今天考试考得开心吗"）
|
||
variant_templates = [
|
||
lambda a, b: f"{a}，那{b}",
|
||
lambda a, b: f"如果{a}，是不是{b}",
|
||
lambda a, b: f"{a}？还有{b}",
|
||
lambda a, b: f"为什么{a}，然后{b}呀",
|
||
]
|
||
needed = count // 3
|
||
for i in range(needed):
|
||
a = random.choice(ALL_TOPICS)
|
||
b = random.choice(ALL_TOPICS)
|
||
tmpl = random.choice(variant_templates)
|
||
questions.append(tmpl(a, b))
|
||
|
||
# 3. "为什么"类衍生
|
||
why_prefixes = ["为什么", "怎么才能", "可以教我", "告诉我", "什么是", "什么时候", "在哪里", "谁"]
|
||
for i in range(count - len(questions)):
|
||
base = random.choice(ALL_TOPICS)
|
||
prefix = random.choice(why_prefixes)
|
||
questions.append(f"{prefix}{base}?")
|
||
|
||
random.shuffle(questions)
|
||
return questions[:count]
|
||
|
||
|
||
def load_existing_pairs(train_path: str) -> set:
|
||
"""加载已有数据，供去重用"""
|
||
if not os.path.exists(train_path):
|
||
return set()
|
||
try:
|
||
with open(train_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
return set((item.get("question", ""), item.get("answer", "")) for item in data)
|
||
except:
|
||
return set()
|
||
|
||
|
||
def append_to_train(new_pairs: list, train_path: str, merged_path: str):
|
||
"""去重后追加到训练集和合并文件"""
|
||
if not new_pairs:
|
||
return 0
|
||
|
||
existing = load_existing_pairs(train_path)
|
||
truly_new = [item for item in new_pairs if (item["question"], item["answer"]) not in existing]
|
||
|
||
if not truly_new:
|
||
print(f" ℹ️ 无新内容（全部已存在）")
|
||
return 0
|
||
|
||
# 追加到训练集
|
||
if os.path.exists(train_path):
|
||
with open(train_path, 'r', encoding='utf-8') as f:
|
||
cur_train = json.load(f)
|
||
else:
|
||
cur_train = []
|
||
cur_train.extend(truly_new)
|
||
with open(train_path, 'w', encoding='utf-8') as f:
|
||
json.dump(cur_train, f, ensure_ascii=False, indent=2)
|
||
|
||
# 同步到合并文件
|
||
if os.path.exists(merged_path):
|
||
with open(merged_path, 'r', encoding='utf-8') as f:
|
||
cur_merged = json.load(f)
|
||
cur_merged.extend(truly_new)
|
||
with open(merged_path, 'w', encoding='utf-8') as f:
|
||
json.dump(cur_merged, f, ensure_ascii=False, indent=2)
|
||
|
||
unique_qs = len(set(item["question"] for item in truly_new))
|
||
print(f" ✅ 新增 {len(truly_new)} 条（{unique_qs} 道独特问题）到训练集")
|
||
return len(truly_new)
|
||
|
||
|
||
def run_distill_round(target_count: int = 1000, save_dir: str = None):
|
||
"""执行一轮蒸馏：生成问题 → API 回答 → 去重追加"""
|
||
if save_dir is None:
|
||
save_dir = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
train_path = os.path.join(save_dir, 'xiaoli_train_v2.json')
|
||
merged_path = os.path.join(save_dir, 'xiaoli_merged_data_v2.json')
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f" 蒸馏轮次开始 | 目标: {target_count} 条")
|
||
print(f"{'='*60}")
|
||
|
||
# 1. 生成种子问题
|
||
questions = generate_seed_questions(target_count)
|
||
print(f" 📝 生成 {len(questions)} 个问题")
|
||
|
||
# 2. 用 Astron API 批量回答（5 并发）
|
||
t0 = time.time()
|
||
results = batch_generate(questions, temperature=0.85, max_workers=5)
|
||
elapsed = time.time() - t0
|
||
|
||
# 3. 去重追加
|
||
added = append_to_train(results, train_path, merged_path)
|
||
|
||
rate = added / max(elapsed, 1) * 60
|
||
print(f" 📊 本轮: +{added} 条 | 耗时 {elapsed:.0f}s | 速率 {rate:.0f} 条/分")
|
||
print(f" 💾 训练集: {os.path.getsize(train_path)/1024/1024:.0f}MB (若存在)")
|
||
|
||
return added
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="小狸 AI 持续蒸馏守护脚本")
|
||
parser.add_argument("--count", type=int, default=1000, help="每轮目标条数")
|
||
parser.add_argument("--loop", action="store_true", help="持续循环模式")
|
||
parser.add_argument("--interval", type=int, default=300, help="循环间隔秒数（默认300）")
|
||
parser.add_argument("--rounds", type=int, default=0, help="循环轮数（0=无限）")
|
||
args = parser.parse_args()
|
||
|
||
if args.loop:
|
||
round_num = 1
|
||
total_added = 0
|
||
while True:
|
||
print(f"\n{'#'*60}")
|
||
print(f" # 第 {round_num} 轮蒸馏")
|
||
print(f"{'#'*60}")
|
||
added = run_distill_round(args.count)
|
||
total_added += added
|
||
print(f" 📈 累计新增: {total_added} 条")
|
||
|
||
if args.rounds > 0 and round_num >= args.rounds:
|
||
print(f"\n✅ 完成 {args.rounds} 轮蒸馏，共新增 {total_added} 条")
|
||
break
|
||
|
||
print(f" ⏳ 等待 {args.interval}s 后下一轮...")
|
||
time.sleep(args.interval)
|
||
round_num += 1
|
||
else:
|
||
run_distill_round(args.count)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |