fix: Astron 5并发 + 修复蒸馏bug

This commit is contained in:
shuiyu1123 authored and Gitee committed 2026-06-20 13:24:47 +00:00
1 parent 1bc19e4c90
commit 623a29bafd
1 file changed
+81 -11
+81 -11
View File
@@ -7,7 +7,7 @@ import json, requests, time, os, threading
from concurrent.futures import ThreadPoolExecutor, as_completed
ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions"
ASTRON_KEY = os.environ.get("ASTRON_KEY", "")
ASTRON_KEY = "658beba1beb0c911ed00672019bbf7e5:MjNmYWRkMWUyZjAxYTM3ZGE2ODI1YzFh"
ASTRON_MODEL = "astron-code-latest"
SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。
@@ -49,8 +49,8 @@ def call_astron(q, temperature=0.85):
return None
def batch_generate(questions, temperature=0.85, max_workers=4):
"""Astron 4并发 — 明确限制,稳定可靠"""
def batch_generate(questions, temperature=0.85, max_workers=5):
"""Astron 5并发 — 明确限制,稳定可靠"""
total = len(questions)
results = []
@@ -87,13 +87,83 @@ batch_generate_direct = batch_generate
batch_generate_dual = batch_generate
def batch_convert(items, *args, **kwargs):
"""兼容占位:训练脚本导入用,已弃用"""
print(" [API] batch_convert已弃用,使用batch_generate")
return []
def batch_convert(items, progress=None, max_workers=5):
"""修复版 batch_convert — 将 (q, ans) 对用 API 优化后返回"""
if not items:
return []
total = len(items)
results = []
def convert(item):
q, model_ans = item
prompt = f"问题:{q}\n模型回答:{model_ans}\n\n请用8岁小狸的语气,重新生成更好的回答。保持简短自然,1-2句话。"
ans = call_astron(prompt, temperature=0.3)
if ans and len(ans) >= 3:
return {"question": q, "answer": ans}
return None
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {executor.submit(convert, item): i for i, item in enumerate(items)}
done = 0
t0 = time.time()
for f in as_completed(futures):
done += 1
r = f.result()
if r:
results.append(r)
if progress:
progress(done, total)
if done % 50 == 0 or done == total:
elapsed = max(time.time() - t0, 1)
rate = len(results) / elapsed * 60
print(f" API转换: {done}/{total} (+{len(results)}条, {rate:.0f}条/分)", end="\r")
elapsed = time.time() - t0
rate = len(results) / max(elapsed, 1) * 60
print(f" API转换: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)")
return results
def batch_generate_multi_temp(questions, *args, **kwargs):
"""兼容占位:训练脚本导入用,已弃用"""
print(" [API] batch_generate_multi_temp已弃用,使用batch_generate")
return []
def batch_generate_multi_temp(questions, temperatures=None, max_workers=5):
"""修复版 — 用多种温度生成多样化回答"""
if temperatures is None:
temperatures = [0.7, 0.85, 1.0]
if not questions:
return []
total = len(questions) * len(temperatures)
results = []
def call_with_temp(q, temp):
ans = call_astron(q, temperature=temp)
if ans:
return {"question": q, "answer": ans, "temperature": temp}
return None
tasks = [(q, t) for q in questions for t in temperatures]
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {executor.submit(call_with_temp, q, t): i for i, (q, t) in enumerate(tasks)}
done = 0
t0 = time.time()
for f in as_completed(futures):
done += 1
r = f.result()
if r:
results.append(r)
if done % 50 == 0 or done == total:
elapsed = max(time.time() - t0, 1)
rate = len(results) / elapsed * 60
print(f" 多温生成: {done}/{total} (+{len(results)}条, {rate:.0f}条/分)", end="\r")
seen = set()
unique = []
for item in results:
key = (item["question"], item["answer"])
if key not in seen:
seen.add(key)
unique.append({"question": item["question"], "answer": item["answer"]})
elapsed = time.time() - t0
rate = len(unique) / max(elapsed, 1) * 60
print(f" 多温生成: 共{len(unique)}条 (去重后, {elapsed:.0f}s, {rate:.0f}条/分)")
return unique