#!/usr/bin/env python3
"""Ollama 모델 비교 스크립트 (bric.pe.kr 「로컬 LLM 설치 가이드 2026」에서 쓴 것과 같은 측정·같은 한국어 10문제).
Python 3 표준 라이브러리만 사용. Ollama가 켜져 있어야 하고, 측정할 모델은 미리 ollama pull 해 두세요.
예: python3 ollama-bench-ko.py gemma4:12b qwen3.5:9b   → 화면 출력 + results.jsonl
For each model: speed (warm run, fixed Korean prompt, 256 tokens, temperature 0, num_ctx 8192), placement from /api/ps,
and 10 Korean questions with exact answers, auto-scored. Thinking is switched off (think=false); gpt-oss cannot switch it off, so it runs at think="low". Answers get up to 1024 tokens.
Usage: python3 bench.py model [model ...]  → appends one JSON line per model to results.jsonl
"""
import json, sys, time, re, shutil, subprocess, urllib.request
import os
URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
OPT = {"temperature": 0, "num_ctx": 8192, "seed": 1}
SPEED_PROMPT = "로컬 LLM을 집에서 돌릴 때의 장점과 단점을 각각 세 가지씩 설명해 주세요."
Q = [  # (question, accepted answers after normalisation)
    ("사과 3개에 1,200원이면 사과 7개는 얼마인가요? 숫자만 답하세요.", ["2800"]),
    ("2026년 10월 17일은 토요일입니다. 2026년 10월 31일은 무슨 요일인가요? 요일만 답하세요.", ["토요일"]),
    ("다음 문장에서 맞춤법이 틀린 단어 하나를 바르게 고쳐 그 단어만 쓰세요: '오늘은 왠지 기분이 않 좋다.'", ["안"]),
    ("'먹다'의 높임말을 한 단어로 쓰세요.", ["드시다", "잡수시다", "잡숫다"]),
    ("'일석이조'와 뜻이 가장 비슷한 속담의 기호만 답하세요. (가) 꿩 먹고 알 먹고 (나) 소 잃고 외양간 고친다 (다) 가는 말이 고와야 오는 말이 곱다", ["가"]),
    ("다음 문장에서 전화번호만 그대로 쓰세요: '문의는 평일 9시부터 6시까지 02-123-4567로 연락 주세요.'", ["021234567"]),
    ("대한민국의 도(道) 가운데 면적이 가장 넓은 곳은? 이름만 쓰세요.", ["경상북도", "경북"]),
    ("다음 문장을 띄어쓰기만 바르게 고쳐 쓰세요: 나는학교에갔다", ["나는 학교에 갔다"]),
    ("'증가'의 반의어를 한 단어로 쓰세요.", ["감소"]),
    ("철수는 영희보다 키가 크고, 영희는 민수보다 키가 큽니다. 세 사람 중 키가 가장 작은 사람은? 이름만 쓰세요.", ["민수"]),
]
def post(path, body):
    r = urllib.request.Request(URL + path, json.dumps(body).encode(), {"Content-Type": "application/json"})
    return json.load(urllib.request.urlopen(r, timeout=1800))
def gen(model, prompt, n, think="auto"):
    if think == "auto": think = think_for(model)
    b = {"model": model, "prompt": prompt, "stream": False, "options": {**OPT, "num_predict": n}, "keep_alive": "10m"}
    if think is not None: b["think"] = think
    try: return post("/api/generate", b)
    except urllib.error.HTTPError as e:
        if think is not None and "think" in e.read().decode().lower(): return gen(model, prompt, n, think=None)
        raise
def gpu_mib():   # Ollama 0.40.5's /api/ps under-reports some models (gemma4:12b shown as 1.1 GB), so read nvidia-smi too
    smi = shutil.which("nvidia-smi") or "/usr/lib/wsl/lib/nvidia-smi"
    try: out = subprocess.run([smi, "--query-gpu=memory.used", "--format=csv,noheader,nounits"], capture_output=True, text=True).stdout
    except OSError: return []   # NVIDIA GPU가 없으면 VRAM 측정은 건너뜀
    return [int(x) for x in out.split()]
def think_for(model): return "low" if model.startswith("gpt-oss") else False   # gpt-oss cannot switch reasoning off; lowest level
def norm(s, spaces):
    s = re.sub(r"<think>.*?</think>", "", s, flags=re.S).strip().splitlines()[0] if s.strip() else ""
    s = re.sub(r"[\s]+", " ", s) if spaces else re.sub(r"\s+", "", s)
    return re.sub(r"[\"'“”‘’.,!?:;()\[\]*`원]|입니다|이다", "", s).replace("-", "")
for model in sys.argv[1:]:
    post("/api/generate", {"model": model, "keep_alive": 0})          # unload, so the first call measures a cold load
    time.sleep(3); base = gpu_mib()
    t = time.time(); cold = gen(model, SPEED_PROMPT, 256); cold_wall = time.time() - t
    warm = gen(model, SPEED_PROMPT, 256)
    used = [a - b for a, b in zip(gpu_mib(), base)]
    ps = json.load(urllib.request.urlopen(URL + "/api/ps"))["models"]
    ps = next((m for m in ps if m["name"] == model or m["model"] == model), ps[0] if ps else {})
    score, answers = 0, []
    for q, ok in Q:
        a = gen(model, q, 1024)["response"]
        spaces = " " in ok[0]
        hit = norm(a, spaces) in [norm(x, spaces) for x in ok]   # exact match on the first line; raw answers are kept for a hand check
        score += hit; answers.append({"q": q[:30], "answer": a.strip()[:80], "ok": bool(hit)})
    res = {"model": model, "cold_load_s": round(cold["load_duration"] / 1e9, 1), "cold_wall_s": round(cold_wall, 1),
           "gen_tok_s": round(warm["eval_count"] / warm["eval_duration"] * 1e9, 1), "prompt_tok_s": round(warm["prompt_eval_count"] / warm["prompt_eval_duration"] * 1e9, 1),
           "loaded_gb": round(ps.get("size", 0) / 1e9, 1), "vram_gb": round(ps.get("size_vram", 0) / 1e9, 1),
           "nvidia_smi_gib_per_gpu": [round(u / 1024, 1) for u in used], "processor_note": ps.get("size_vram", 0) / max(ps.get("size", 1), 1),
           "korean_score": score, "answers": answers}
    print(json.dumps({k: v for k, v in res.items() if k != "answers"}, ensure_ascii=False), flush=True)
    open("results.jsonl", "a", encoding="utf-8").write(json.dumps(res, ensure_ascii=False) + "\n")
    post("/api/generate", {"model": model, "keep_alive": 0})