165 lines
7.8 KiB
Python
165 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
||
"""exp12 — генерация редакторских армов A1′,A2,A3,A4,A5,A6 (платно, кап $5).
|
||
|
||
Черновики (A0) и финал glm-5-моно (A1) — из store, $0 (не тут). Здесь ре-редактируем
|
||
ТЕ ЖЕ черновики разными редакторами/режимами. Промпты ЗАМОРОЖЕНЫ (пре-регистрация §1.2).
|
||
|
||
Faithful render: system(editor.md) → system(инъекция-блок) → user(черновик[+исходник]).
|
||
Layout повторяет render.go MessagesWithInjection.
|
||
|
||
Запуск: eval/.venv/bin/python eval/exp12_arms.py [--dry]
|
||
Выход: diag/arms/<ARM>/<ch>_<chunk>.txt + diag/arms/results.json
|
||
"""
|
||
from __future__ import annotations
|
||
import json, os, sys, time, urllib.request, urllib.error
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||
from refusal_bench import load_env_file
|
||
load_env_file()
|
||
|
||
DIAG = Path("/home/ubuntu/books/gu-zhenren/diag")
|
||
ARMS = DIAG / "arms"
|
||
EDITOR_MD = Path("/home/ubuntu/projects/textmachine/backend/prompts/editor.md")
|
||
CAP_USD = 5.0
|
||
|
||
# --- render vars (book.yaml) ---
|
||
VARS = {"genre": "вебновелла", "audience": "взрослые читатели вебновелл",
|
||
"title": "蛊真人", "source_lang": "zh", "honorifics": "keep", "venuti": "0.60"}
|
||
|
||
PRICES = { # per 1M tokens (models.yaml, prices_checked 2026-07-10)
|
||
"glm-5": (1.0, 3.2), "grok-4.3": (1.25, 2.50),
|
||
"gemini-3.1-pro-preview": (2.0, 12.0),
|
||
}
|
||
|
||
# --- frozen arm specs ---
|
||
BILINGUAL_LINE = ("Тебе также дан ИСХОДНЫЙ ТЕКСТ: сверяйся с ним по смыслу и исправляй "
|
||
"искажения черновика, сохраняя полноту.")
|
||
ARMS_SPEC = {
|
||
"A1p": dict(model="glm-5", base="https://api.z.ai/api/paas/v4", keyenv="ZAI_API_KEY",
|
||
mode="mono", constraints=True, extra={"thinking": {"type": "disabled"}}, max_tokens=8000),
|
||
"A2": dict(model="glm-5", base="https://api.z.ai/api/paas/v4", keyenv="ZAI_API_KEY",
|
||
mode="biling", constraints=True, extra={"thinking": {"type": "disabled"}}, max_tokens=8000),
|
||
"A3": dict(model="grok-4.3", base="https://api.x.ai/v1", keyenv="XAI_API_KEY",
|
||
mode="mono", constraints=True, extra={"reasoning_effort": "none"}, max_tokens=8000),
|
||
"A4": dict(model="grok-4.3", base="https://api.x.ai/v1", keyenv="XAI_API_KEY",
|
||
mode="biling", constraints=True, extra={"reasoning_effort": "none"}, max_tokens=8000),
|
||
"A5": dict(model="gemini-3.1-pro-preview",
|
||
base="https://generativelanguage.googleapis.com/v1beta/openai", keyenv="GEMINI_API_KEY",
|
||
mode="biling", constraints=True, extra={}, max_tokens=16000),
|
||
"A6": dict(model="grok-4.3", base="https://api.x.ai/v1", keyenv="XAI_API_KEY",
|
||
mode="mono", constraints=False, extra={"reasoning_effort": "none"}, max_tokens=8000),
|
||
}
|
||
|
||
|
||
def render(tpl: str) -> str:
|
||
for k, v in VARS.items():
|
||
tpl = tpl.replace("{{%s}}" % k, v)
|
||
return tpl
|
||
|
||
|
||
def editor_system() -> str:
|
||
raw = EDITOR_MD.read_text(encoding="utf-8")
|
||
sys_part = raw.split("---USER---")[0].rstrip()
|
||
return render(sys_part)
|
||
|
||
|
||
def build_messages(spec: dict, source: str, draft: str, blk: dict) -> list[dict]:
|
||
sys1 = editor_system()
|
||
if spec["mode"] == "biling":
|
||
sys1 = sys1 + "\n" + BILINGUAL_LINE
|
||
msgs = [{"role": "system", "content": sys1}]
|
||
if spec["constraints"]:
|
||
inj = blk["bilingual_glossary"] if spec["mode"] == "biling" else blk["editor_constraint"]
|
||
if inj.strip():
|
||
msgs.append({"role": "system", "content": inj})
|
||
if spec["mode"] == "biling":
|
||
user = f"ИСХОДНЫЙ ТЕКСТ (zh):\n\n{source}\n\n---\n\nЧерновик перевода для редактуры:\n\n{draft}"
|
||
else:
|
||
user = f"Черновик перевода для редактуры:\n\n{draft}"
|
||
msgs.append({"role": "user", "content": user})
|
||
return msgs
|
||
|
||
|
||
def call_messages(spec: dict, messages: list[dict], timeout: int = 300):
|
||
key = os.environ.get(spec["keyenv"], "")
|
||
if not key:
|
||
return None, f"no key {spec['keyenv']}", {}
|
||
body = {"model": spec["model"], "messages": messages,
|
||
"temperature": 0.4, "max_tokens": spec["max_tokens"], **spec["extra"]}
|
||
req = urllib.request.Request(spec["base"].rstrip("/") + "/chat/completions",
|
||
data=json.dumps(body).encode(),
|
||
headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"})
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||
data = json.loads(r.read())
|
||
except urllib.error.HTTPError as e:
|
||
return None, f"HTTP {e.code}: {e.read()[:200].decode('utf-8','replace')}", {}
|
||
except Exception as e:
|
||
return None, f"ERR {str(e)[:160]}", {}
|
||
ch = data["choices"][0]
|
||
text = (ch.get("message", {}) or {}).get("content") or ""
|
||
usage = data.get("usage", {}) or {}
|
||
usage["finish_reason"] = ch.get("finish_reason", "")
|
||
return text.strip(), (None if text.strip() else f"empty|finish={ch.get('finish_reason')}"), usage
|
||
|
||
|
||
def cost_of(model: str, usage: dict) -> float:
|
||
inp = usage.get("prompt_tokens", 0) or 0
|
||
comp = usage.get("completion_tokens", 0) or 0
|
||
total = usage.get("total_tokens", 0) or 0
|
||
reason = usage.get("reasoning_tokens", 0) or 0
|
||
pin, pout = PRICES.get(model, (0, 0))
|
||
if model == "gemini-3.1-pro-preview": # additive_total: thinking only in total_tokens
|
||
reason = max(0, total - inp - comp)
|
||
return (inp * pin + (comp + reason) * pout) / 1_000_000
|
||
|
||
|
||
def main() -> None:
|
||
dry = "--dry" in sys.argv
|
||
chunks = json.loads((DIAG / "chunks.json").read_text(encoding="utf-8"))
|
||
blocks = json.loads((DIAG / "blocks.json").read_text(encoding="utf-8"))
|
||
targets = [(5, 0), (5, 1), (7, 0), (7, 1), (14, 0), (14, 1)]
|
||
|
||
ARMS.mkdir(parents=True, exist_ok=True)
|
||
results = []
|
||
spent = 0.0
|
||
for arm, spec in ARMS_SPEC.items():
|
||
adir = ARMS / arm
|
||
adir.mkdir(exist_ok=True)
|
||
for ch, ci in targets:
|
||
key = f"{ch}/{ci}"
|
||
src = chunks[key]["source"]
|
||
draft = chunks[key]["draft"]
|
||
blk = blocks[key]
|
||
outp = adir / f"{ch}_{ci}.txt"
|
||
if outp.exists() and outp.read_text(encoding="utf-8").strip():
|
||
print(f" [{arm} {key}] cached"); continue
|
||
msgs = build_messages(spec, src, draft, blk)
|
||
if dry:
|
||
approx_in = sum(len(m["content"]) for m in msgs) // 3
|
||
print(f" [{arm} {key}] mode={spec['mode']} constr={spec['constraints']} "
|
||
f"msgs={len(msgs)} ~in={approx_in}tok draft_chars={len(draft)}")
|
||
continue
|
||
if spent > CAP_USD * 0.8:
|
||
print(f"!! near cap (${spent:.3f}) — stop"); break
|
||
t0 = time.time()
|
||
text, err, usage = call_messages(spec, msgs)
|
||
c = cost_of(spec["model"], usage)
|
||
spent += c
|
||
if not text:
|
||
print(f" [{arm} {key}] ERROR: {err} (${spent:.3f})")
|
||
results.append({"arm": arm, "chunk": key, "error": err, "cost": c})
|
||
continue
|
||
outp.write_text(text, encoding="utf-8")
|
||
results.append({"arm": arm, "chunk": key, "chars": len(text),
|
||
"finish": usage.get("finish_reason"), "usage": usage, "cost": round(c, 5)})
|
||
print(f" [{arm} {key}] {time.time()-t0:.0f}s {len(text)}ch ${c:.4f} cum=${spent:.3f} "
|
||
f"finish={usage.get('finish_reason')}")
|
||
if not dry:
|
||
(ARMS / "results.json").write_text(json.dumps(results, ensure_ascii=False, indent=1), encoding="utf-8")
|
||
print(f"\nTOTAL arm spend: ${spent:.4f} (cap ${CAP_USD})")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|