169 lines
5.9 KiB
Python
169 lines
5.9 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
API生成模块 — Astron 5并发，稳定不超限
|
|
"""
|
|
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_MODEL = "astron-code-latest"
|
|
|
|
SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。
|
|
|
|
你的人设：
|
|
- 8岁女孩子，有点害羞但很可爱
|
|
- 说话简单自然，符合小朋友的语气
|
|
- 回答问题要简短，1-2句话就够了
|
|
- 不要用"喵"结尾，不要过度卖萌
|
|
|
|
回答要求：
|
|
1. 语气自然，像真实的小朋友说话
|
|
2. 不懂就说"不知道"
|
|
3. 回答简洁但要说清楚，不要讲一半就停"""
|
|
|
|
|
|
def call_astron(q, temperature=0.85):
|
|
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": q}
|
|
],
|
|
"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:
|
|
return content
|
|
else:
|
|
time.sleep(1)
|
|
except Exception:
|
|
time.sleep(2)
|
|
return None
|
|
|
|
|
|
def batch_generate(questions, temperature=0.85, max_workers=5):
|
|
"""Astron 5并发 — 明确限制，稳定可靠"""
|
|
total = len(questions)
|
|
results = []
|
|
|
|
def call(q):
|
|
ans = call_astron(q, temperature)
|
|
if ans:
|
|
return q, ans
|
|
return None
|
|
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = {executor.submit(call, item): i for i, item in enumerate(questions)}
|
|
done = 0
|
|
t0 = time.time()
|
|
for f in as_completed(futures):
|
|
done += 1
|
|
r = f.result()
|
|
if r:
|
|
results.append({"question": r[0], "answer": r[1]})
|
|
if done % 20 == 0 or done == total:
|
|
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, 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, 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 |