9bb00d49f2
TextMachine project repository (AI translation of literary books). Includes: v2 architecture decisions, MVP plan, research 01-12, polygon experiments 01-03, backend-session revalidation verdict (03-implementation-notes.md). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
211 lines
9.3 KiB
Python
211 lines
9.3 KiB
Python
#!/usr/bin/env python3
|
||
"""Локальный стенд TextMachine (задача 3 полигона, docs/research/06-local-models.md).
|
||
|
||
Меряет на установленных ollama-моделях:
|
||
- скорость: prompt_eval tok/s и generation tok/s (из метрик /api/generate);
|
||
- память: пик VRAM (nvidia-smi) и RSS ollama во время генерации;
|
||
- качество: сохраняет переводы ja→ru / zh→ru для субъективного сравнения
|
||
(эталон DeepSeek API добавляется тем же скриптом при наличии ключа).
|
||
|
||
Запуск: eval/.venv/bin/python eval/local_bench.py [--models m1,m2] [--num-ctx 8192]
|
||
Выход: eval/data/local_bench/<model>/<fragment>.txt + results.json + таблица в stdout.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
OLLAMA = "http://localhost:11434"
|
||
|
||
# WSL: в env задан http_proxy, а NO_PROXY=<local> (WinINET-формат) urllib не понимает —
|
||
# localhost-запросы уходят на прокси и получают 403. Для ollama прокси отключаем явно.
|
||
_NO_PROXY_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))
|
||
|
||
|
||
def _load_env_file() -> None:
|
||
"""Ключи из eval/.env (gitignored, закрыт от Claude deny-правилом) или
|
||
legacy ~/.textmachine.env; только в env процесса, не печатаются."""
|
||
for cand in (Path(__file__).resolve().parent / ".env",
|
||
Path("~/.textmachine.env").expanduser()):
|
||
if not cand.exists():
|
||
continue
|
||
for line in cand.read_text().splitlines():
|
||
line = line.strip()
|
||
if not line or line.startswith("#") or "=" not in line:
|
||
continue
|
||
k, _, v = line.partition("=")
|
||
if v.strip():
|
||
os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))
|
||
return
|
||
|
||
|
||
_load_env_file()
|
||
|
||
PROMPT_TMPL = (
|
||
"Ты профессиональный литературный переводчик с {src_name} на русский. "
|
||
"Переведи фрагмент художественной прозы. Требования: точность смысла, живой "
|
||
"литературный русский без канцелярита и калек, сохранить все предложения и реплики. "
|
||
"Выведи только перевод.\n\n{text}"
|
||
)
|
||
SRC_NAMES = {"ja": "японского", "zh": "китайского"}
|
||
|
||
# Фрагменты: (lang, файл, срез символов) — ~1.2-1.5k знаков, самодостаточные сцены
|
||
FRAGMENTS = [
|
||
("ja", "akutagawa-rashomon.txt", 0, 1400),
|
||
("ja", "dazai-hashire-merosu.txt", 0, 1400),
|
||
("zh", "luxun-zhufu.txt", 0, 1500),
|
||
]
|
||
|
||
|
||
def gpu_mem_mb() -> int:
|
||
try:
|
||
out = subprocess.run(
|
||
["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
|
||
capture_output=True, text=True, timeout=5).stdout.strip()
|
||
return int(out.splitlines()[0])
|
||
except Exception:
|
||
return -1
|
||
|
||
|
||
def ram_used_gb() -> float:
|
||
with open("/proc/meminfo") as f:
|
||
mi = {l.split(":")[0]: int(l.split()[1]) for l in f}
|
||
return (mi["MemTotal"] - mi["MemAvailable"]) / 1e6
|
||
|
||
|
||
class PeakSampler(threading.Thread):
|
||
def __init__(self):
|
||
super().__init__(daemon=True)
|
||
self.peak_vram = 0
|
||
self.peak_ram = 0.0
|
||
self.stop_flag = False
|
||
|
||
def run(self):
|
||
while not self.stop_flag:
|
||
self.peak_vram = max(self.peak_vram, gpu_mem_mb())
|
||
self.peak_ram = max(self.peak_ram, ram_used_gb())
|
||
time.sleep(1)
|
||
|
||
|
||
def ollama_generate(model: str, prompt: str, num_ctx: int, timeout: int = 1800) -> dict:
|
||
req = urllib.request.Request(
|
||
f"{OLLAMA}/api/generate",
|
||
data=json.dumps({
|
||
"model": model, "prompt": prompt, "stream": False,
|
||
"options": {"num_ctx": num_ctx, "temperature": 0.3},
|
||
}).encode(),
|
||
headers={"Content-Type": "application/json"})
|
||
with _NO_PROXY_OPENER.open(req, timeout=timeout) as resp:
|
||
return json.loads(resp.read())
|
||
|
||
|
||
def strip_think(text: str) -> str:
|
||
"""Qwen3 в дефолте может отдавать <think>…</think> — вырезаем для оценки перевода."""
|
||
return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.S).strip()
|
||
|
||
|
||
def deepseek_generate(prompt: str) -> dict | None:
|
||
key = os.environ.get("DEEPSEEK_API_KEY")
|
||
if not key:
|
||
return None
|
||
req = urllib.request.Request(
|
||
"https://api.deepseek.com/v1/chat/completions",
|
||
data=json.dumps({"model": "deepseek-chat", "temperature": 0.3,
|
||
"messages": [{"role": "user", "content": prompt}]}).encode(),
|
||
headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"})
|
||
t0 = time.time()
|
||
with urllib.request.urlopen(req, timeout=600) as resp:
|
||
data = json.loads(resp.read())
|
||
return {"response": data["choices"][0]["message"]["content"],
|
||
"elapsed_s": time.time() - t0, "usage": data.get("usage", {})}
|
||
|
||
|
||
def installed_models() -> list[str]:
|
||
with _NO_PROXY_OPENER.open(f"{OLLAMA}/api/tags", timeout=10) as resp:
|
||
return [m["name"] for m in json.loads(resp.read())["models"]]
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--models", help="через запятую; по умолчанию все установленные")
|
||
ap.add_argument("--num-ctx", type=int, default=8192)
|
||
ap.add_argument("--out", default=str(ROOT / "data" / "local_bench"))
|
||
args = ap.parse_args()
|
||
|
||
models = args.models.split(",") if args.models else installed_models()
|
||
out_dir = Path(args.out)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
results = []
|
||
for model in models:
|
||
print(f"\n=== {model} (num_ctx={args.num_ctx}) ===")
|
||
subprocess.run(["ollama", "stop", model.split(":")[0]], capture_output=True)
|
||
base_vram, base_ram = gpu_mem_mb(), ram_used_gb()
|
||
for lang, fname, a, b in FRAGMENTS:
|
||
src_path = ROOT / "data" / "samples" / lang / fname
|
||
if not src_path.exists():
|
||
print(f" [skip] нет {src_path}")
|
||
continue
|
||
text = src_path.read_text()[a:b]
|
||
prompt = PROMPT_TMPL.format(src_name=SRC_NAMES[lang], text=text)
|
||
sampler = PeakSampler()
|
||
sampler.start()
|
||
t0 = time.time()
|
||
try:
|
||
r = ollama_generate(model, prompt, args.num_ctx)
|
||
except Exception as e:
|
||
sampler.stop_flag = True
|
||
print(f" {fname}: ОШИБКА {type(e).__name__}: {str(e)[:200]}")
|
||
results.append({"model": model, "fragment": fname, "error": str(e)[:500]})
|
||
continue
|
||
elapsed = time.time() - t0
|
||
sampler.stop_flag = True
|
||
sampler.join(2)
|
||
gen_tps = r.get("eval_count", 0) / max(r.get("eval_duration", 1), 1) * 1e9
|
||
pp_tps = r.get("prompt_eval_count", 0) / max(r.get("prompt_eval_duration", 1), 1) * 1e9
|
||
translation = strip_think(r.get("response", ""))
|
||
mdir = out_dir / model.replace("/", "_").replace(":", "_")
|
||
mdir.mkdir(exist_ok=True)
|
||
(mdir / f"{fname.removesuffix('.txt')}.ru.txt").write_text(translation)
|
||
rec = {"model": model, "fragment": fname, "lang": lang,
|
||
"src_chars": len(text),
|
||
"prompt_eval_tps": round(pp_tps, 1), "gen_tps": round(gen_tps, 1),
|
||
"gen_tokens": r.get("eval_count"), "elapsed_s": round(elapsed, 1),
|
||
"peak_vram_mb": sampler.peak_vram, "base_vram_mb": base_vram,
|
||
"peak_ram_gb": round(sampler.peak_ram, 1),
|
||
"base_ram_gb": round(base_ram, 1),
|
||
"out_chars": len(translation)}
|
||
results.append(rec)
|
||
print(f" {fname}: gen {gen_tps:.1f} tok/s, prompt {pp_tps:.0f} tok/s, "
|
||
f"{elapsed:.0f}s, VRAM пик {sampler.peak_vram} MB, RAM пик {sampler.peak_ram:.1f} GB")
|
||
|
||
# эталон DeepSeek (если есть ключ)
|
||
if os.environ.get("DEEPSEEK_API_KEY"):
|
||
print("\n=== deepseek-chat (эталон API) ===")
|
||
mdir = out_dir / "deepseek-api"
|
||
mdir.mkdir(exist_ok=True)
|
||
for lang, fname, a, b in FRAGMENTS:
|
||
src_path = ROOT / "data" / "samples" / lang / fname
|
||
text = src_path.read_text()[a:b]
|
||
r = deepseek_generate(PROMPT_TMPL.format(src_name=SRC_NAMES[lang], text=text))
|
||
(mdir / f"{fname.removesuffix('.txt')}.ru.txt").write_text(r["response"])
|
||
results.append({"model": "deepseek-api", "fragment": fname,
|
||
"elapsed_s": round(r["elapsed_s"], 1), "usage": r["usage"]})
|
||
print(f" {fname}: {r['elapsed_s']:.0f}s")
|
||
else:
|
||
print("\n[note] DEEPSEEK_API_KEY нет — эталонные переводы не сняты")
|
||
|
||
(out_dir / "results.json").write_text(json.dumps(results, ensure_ascii=False, indent=1))
|
||
print(f"\nJSON: {out_dir}/results.json; переводы: {out_dir}/<model>/*.ru.txt")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|