fix: scripts API同步
This commit is contained in:
1 parent
623a29bafd
commit
2ba1ee3416
1 file changed
+69
-20
+69
-20
@@ -1,13 +1,13 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
API生成模块 — Astron 5并发,稳定不超限
|
||||
API生成模块 — Astron 5并发(同步根目录版本)
|
||||
"""
|
||||
import json, requests, time, os, threading
|
||||
import json, requests, time, os
|
||||
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,17 +49,12 @@ 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):
|
||||
total = len(questions)
|
||||
results = []
|
||||
|
||||
def call(q):
|
||||
ans = call_astron(q, temperature)
|
||||
if ans:
|
||||
return q, ans
|
||||
return None
|
||||
|
||||
return (q, ans) if ans else None
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = {executor.submit(call, item): i for i, item in enumerate(questions)}
|
||||
done = 0
|
||||
@@ -73,27 +68,81 @@ def batch_generate(questions, temperature=0.85, max_workers=4):
|
||||
elapsed = max(time.time() - t0, 1)
|
||||
rate = len(results) / elapsed * 60
|
||||
print(f" Astron: {done}/{total} (已获{len(results)}条, {rate:.0f}条/分)", end="\r")
|
||||
|
||||
elapsed = time.time() - t0
|
||||
rate = len(results) / max(elapsed, 1) * 60
|
||||
print(f" Astron: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)")
|
||||
return results
|
||||
|
||||
|
||||
# 保持旧函数名兼容
|
||||
batch_generate_astron = batch_generate
|
||||
batch_generate_mimo = batch_generate
|
||||
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):
|
||||
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)
|
||||
return {"question": q, "answer": ans} if ans and len(ans) >= 3 else 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)
|
||||
return {"question": q, "answer": ans, "temperature": temp} if ans else 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
|
||||
Reference in new issue
Block a user