Quantization - 重みを 8 bit・4 bit に詰めて小さく速くする

"""GGUF ファイルのヘッダーだけを読み、テンソルの型ごとの個数とバイト数を数える。

テンソルの中身は読まない。Python 標準ライブラリだけで動く。

使い方:
    python3 gguf_header.py <名前>=<GGUF のパス> [<名前>=<GGUF のパス> ...] [--json 結果.json]
"""
import json
import os
import struct
import sys
from collections import defaultdict

# ggml の型番号 -> (名前, 1 ブロックの要素数, 1 ブロックのバイト数)
# 値は llama.cpp の ggml/src/ggml.c の type_traits と ggml-common.h の構造体の大きさ
GGML_TYPES = {
    0: ("F32", 1, 4),
    1: ("F16", 1, 2),
    2: ("Q4_0", 32, 18),
    3: ("Q4_1", 32, 20),
    6: ("Q5_0", 32, 22),
    7: ("Q5_1", 32, 24),
    8: ("Q8_0", 32, 34),
    10: ("Q2_K", 256, 84),
    11: ("Q3_K", 256, 110),
    12: ("Q4_K", 256, 144),
    13: ("Q5_K", 256, 176),
    14: ("Q6_K", 256, 210),
    30: ("BF16", 1, 2),
}

# GGUF のメタデータの値の型番号 -> struct の書式（8: 文字列、9: 配列は別に処理する）
SCALAR = {0: "<B", 1: "<b", 2: "<H", 3: "<h", 4: "<I", 5: "<i", 6: "<f",
          7: "<?", 10: "<Q", 11: "<q", 12: "<d"}


class Reader:
    def __init__(self, f):
        self.f = f

    def unpack(self, fmt):
        size = struct.calcsize(fmt)
        return struct.unpack(fmt, self.f.read(size))[0]

    def string(self):
        n = self.unpack("<Q")
        return self.f.read(n).decode("utf-8", errors="replace")

    def value(self, vtype):
        if vtype == 8:
            return self.string()
        if vtype == 9:
            itype, count = self.unpack("<I"), self.unpack("<Q")
            return [self.value(itype) for _ in range(count)]
        return self.unpack(SCALAR[vtype])


def read_gguf(path):
    """GGUF のヘッダー（メタデータとテンソルの一覧）を返す。"""
    with open(path, "rb") as f:
        r = Reader(f)
        if f.read(4) != b"GGUF":
            raise ValueError(f"{path} is not GGUF")
        version = r.unpack("<I")
        n_tensors, n_kv = r.unpack("<Q"), r.unpack("<Q")
        meta = {}
        for _ in range(n_kv):
            key = r.string()
            meta[key] = r.value(r.unpack("<I"))
        tensors = []
        for _ in range(n_tensors):
            name = r.string()
            n_dims = r.unpack("<I")
            shape = [r.unpack("<Q") for _ in range(n_dims)]
            ttype = r.unpack("<I")
            offset = r.unpack("<Q")
            tensors.append({"name": name, "shape": shape, "type": ttype, "offset": offset})
        align = meta.get("general.alignment", 32)
        data_start = (f.tell() + align - 1) // align * align
    return version, meta, tensors, data_start


def tensor_bytes(t):
    _, block, nbytes = GGML_TYPES[t["type"]]
    n = 1
    for d in t["shape"]:
        n *= d
    return n, n // block * nbytes


def summarize(label, path):
    version, meta, tensors, data_start = read_gguf(path)
    by_type = defaultdict(lambda: {"tensors": 0, "params": 0, "bytes": 0})
    for t in tensors:
        n, b = tensor_bytes(t)
        s = by_type[GGML_TYPES[t["type"]][0]]
        s["tensors"] += 1
        s["params"] += n
        s["bytes"] += b
    params = sum(s["params"] for s in by_type.values())
    weight_bytes = sum(s["bytes"] for s in by_type.values())
    size = os.path.getsize(path)
    print(f"[{label}] {os.path.basename(path)[:19]}...")
    print(f"  file size        : {size:,} bytes")
    print(f"  GGUF version     : {version}, metadata keys: {len(meta)}, tensors: {len(tensors)}")
    print(f"  general.file_type: {meta.get('general.file_type')}")
    print(f"  params           : {params:,}")
    print(f"  weight bytes     : {weight_bytes:,} ({weight_bytes * 8 / params:.2f} bits/param)")
    print(f"  header + padding : {data_start:,} bytes")
    for name, s in sorted(by_type.items(), key=lambda x: -x[1]["bytes"]):
        print(f"    {name:5s} tensors {s['tensors']:4d}  params {s['params']:13,}"
              f"  bytes {s['bytes']:13,}  ({s['bytes'] * 8 / s['params']:.4f} bits/param)")
    return {"file_size": size, "params": params, "weight_bytes": weight_bytes,
            "file_type": meta.get("general.file_type"), "by_type": dict(by_type),
            "tensor_types": {t["name"]: GGML_TYPES[t["type"]][0] for t in tensors}}


def main():
    args = sys.argv[1:]
    out = None
    if "--json" in args:
        i = args.index("--json")
        out = args[i + 1]
        del args[i:i + 2]
    result = {}
    for a in args:
        label, path = a.split("=", 1)
        result[label] = summarize(label, path)
        print()
    if out:
        with open(out, "w") as f:
            json.dump(result, f, indent=1)


if __name__ == "__main__":
    main()

"""1 つの重み行列を自分で Q8_0 と 4 bit に丸め、FP16 との誤差を測る。

Ollama の q8_0・q4_K_M のファイルに入っている同じ行列も復元して並べる。

使い方:
    python3 block_error.py <fp16 の GGUF> <q8_0 の GGUF> <q4_K_M の GGUF> [結果の JSON]
"""
import json
import sys

import numpy as np

from gguf_header import read_gguf

TENSOR = "blk.0.ffn_gate.weight"   # 2048 x 6144。どちらの量子化でも整数に丸められている層


def load_raw(path, name):
    """GGUF から 1 つのテンソルのバイト列を、行（ne0 方向）ごとに読み出す。"""
    _, _, tensors, data_start = read_gguf(path)
    t = next(x for x in tensors if x["name"] == name)
    return t, np.memmap(path, mode="r", offset=data_start + t["offset"], dtype=np.uint8)


def load_f16(path, name):
    t, raw = load_raw(path, name)
    n = int(np.prod(t["shape"]))
    return raw[: n * 2].view(np.float16).astype(np.float32)


# --- 自分で丸める ------------------------------------------------------------

def roundf(x):
    """C の roundf と同じく、ちょうど 0.5 は 0 から遠い側に丸める（np.round は偶数側）。"""
    t = np.trunc(x)
    return t + np.where(np.abs(x - t) >= 0.5, np.sign(x), 0)


def quant_q8_0(x):
    """32 個ごとに d = max|x| / 127 を求め、x / d を -127〜127 の整数に丸める。
    llama.cpp の quantize_row_q8_0_ref と同じく、丸めには FP32 の d の逆数を使い、
    d を FP16 にするのはファイルに書く時だけ。"""
    b = x.reshape(-1, 32)
    d32 = np.abs(b).max(axis=1) / np.float32(127)
    inv = np.divide(np.float32(1), d32, out=np.zeros_like(d32), where=d32 != 0)
    q = roundf(b * inv[:, None]).astype(np.int8)
    d = d32.astype(np.float16)                                 # ファイルには FP16 で入る
    return q, d, (q * d.astype(np.float32)[:, None]).reshape(-1)


def quant_q4_0(x):
    """32 個ごとに、絶対値が最大の値を -8 に合わせた d を持ち、-8〜7 の整数に丸める。"""
    b = x.reshape(-1, 32)
    idx = np.abs(b).argmax(axis=1)
    mx = b[np.arange(len(b)), idx]                           # 符号付きの最大値
    d32 = mx / np.float32(-8)
    inv = np.divide(np.float32(1), d32, out=np.zeros_like(d32), where=d32 != 0)
    q = np.clip(np.floor(b * inv[:, None] + 8.5), 0, 15).astype(np.int8) - 8   # (int8_t)(x + 8.5f)
    d = d32.astype(np.float16)                                 # Q8_0 と同じく FP16 は保存用
    return q, d, (q * d.astype(np.float32)[:, None]).reshape(-1)


def quant_q4_min(x):
    """32 個ごとに最小値 m と幅 d を持ち、(x - m) / d を 0〜15 の整数に丸める（Q4_K の小ブロックと同じ形）。"""
    b = x.reshape(-1, 32)
    lo, hi = b.min(axis=1), b.max(axis=1)
    d = ((hi - lo) / 15).astype(np.float16).astype(np.float32)
    lo = lo.astype(np.float16).astype(np.float32)
    inv = np.divide(1.0, d, out=np.zeros_like(d), where=d != 0)
    q = np.clip(np.round((b - lo[:, None]) * inv[:, None]), 0, 15).astype(np.int8)
    return q, d, (q * d[:, None] + lo[:, None]).reshape(-1)


# --- Ollama のファイルを復元する（llama.cpp の dequantize_row_q8_0 / q4_K と同じ計算）--

def dequant_q8_0(raw, n):
    blocks = raw[: n // 32 * 34].reshape(-1, 34)
    d = blocks[:, :2].copy().view(np.float16).astype(np.float32)   # (nb, 1)
    q = blocks[:, 2:].copy().view(np.int8).astype(np.float32)
    return (q * d).reshape(-1)


def dequant_q4_k(raw, n):
    """256 個の大ブロック = FP16 の d と dmin、6 bit の scale と min が 8 組、4 bit の整数 256 個。"""
    blocks = raw[: n // 256 * 144].reshape(-1, 144)
    d = blocks[:, 0:2].copy().view(np.float16).astype(np.float32)[:, 0]
    dmin = blocks[:, 2:4].copy().view(np.float16).astype(np.float32)[:, 0]
    s = blocks[:, 4:16].astype(np.int32)
    qs = blocks[:, 16:144]
    sc = np.empty((len(blocks), 8), np.int32)
    mn = np.empty((len(blocks), 8), np.int32)
    for j in range(8):                                   # get_scale_min_k4()
        if j < 4:
            sc[:, j] = s[:, j] & 63
            mn[:, j] = s[:, j + 4] & 63
        else:
            sc[:, j] = (s[:, j + 4] & 0xF) | ((s[:, j - 4] >> 6) << 4)
            mn[:, j] = (s[:, j + 4] >> 4) | ((s[:, j] >> 6) << 4)
    out = np.empty((len(blocks), 256), np.float32)
    for k in range(4):                                   # 64 個ずつ: 下位 4 bit が前半 32 個、上位が後半
        q = qs[:, k * 32:(k + 1) * 32]
        for half, vals in ((0, q & 0xF), (1, q >> 4)):
            j = 2 * k + half
            out[:, j * 32:(j + 1) * 32] = (d * sc[:, j])[:, None] * vals - (dmin * mn[:, j])[:, None]
    return out.reshape(-1)


def rel_rmse(x, y):
    return float(np.sqrt(np.mean((x - y) ** 2)) / np.sqrt(np.mean(x ** 2)))


def main():
    fp16_path, q8_path, q4_path = sys.argv[1:4]
    out_path = sys.argv[4] if len(sys.argv) > 4 else None
    w = load_f16(fp16_path, TENSOR)
    n = w.size

    # 1 ブロック（32 個）の数値例。絶対値が最大の値（32 個目）を含む後ろの 8 個を表示する
    blk = w[:32]
    q8, d8, r8 = quant_q8_0(blk)
    q4, d4, r4 = quant_q4_0(blk)
    print(f"{TENSOR}: {n:,} weights, first block (32 weights), last 8 shown")
    print(f"  FP16        : {np.array2string(blk[24:], precision=4, floatmode='fixed')}")
    print(f"  Q8_0  d={float(d8[0]):.6f}  ints: {q8[0, 24:].tolist()}")
    print(f"        restored: {np.array2string(r8[24:], precision=4, floatmode='fixed')}")
    print(f"  Q4_0  d={float(d4[0]):.6f}  ints: {q4[0, 24:].tolist()}")
    print(f"        restored: {np.array2string(r4[24:], precision=4, floatmode='fixed')}")
    print(f"  max |error| in this block: Q8_0 {np.abs(blk - r8).max():.5f}, "
          f"Q4_0 {np.abs(blk - r4).max():.5f}  (max |w| {np.abs(blk).max():.5f})")

    # 行列全体の誤差
    _, raw8 = load_raw(q8_path, TENSOR)
    _, raw4 = load_raw(q4_path, TENSOR)
    mine8 = quant_q8_0(w)[2]
    olla8 = dequant_q8_0(raw8, n)
    rows = {
        "Q8_0 (mine)": mine8,
        "Q8_0 (Ollama q8_0 file)": olla8,
        "Q4_0 (mine)": quant_q4_0(w)[2],
        "Q4 scale+min (mine)": quant_q4_min(w)[2],
        "Q4_K (Ollama q4_K_M file)": dequant_q4_k(raw4, n),
    }
    print(f"\nwhole matrix ({n:,} weights), error relative to FP16")
    result = {"tensor": TENSOR, "n": n, "rel_rmse": {}}
    for k, v in rows.items():
        e = rel_rmse(w, v)
        result["rel_rmse"][k] = e
        print(f"  {k:27s} relative RMSE {e * 100:6.3f} %   max |error| {np.abs(w - v).max():.5f}")
    same = int(np.sum(mine8 == olla8))
    print(f"\nQ8_0 mine vs Ollama: {same:,} / {n:,} weights identical ({same / n * 100:.4f} %)")
    result["q8_identical"] = same

    # 誤差のヒストグラム用（図に使う）
    hist = {}
    edges = np.linspace(-0.02, 0.02, 81)
    for k in ("Q8_0 (Ollama q8_0 file)", "Q4_K (Ollama q4_K_M file)"):
        h, _ = np.histogram(rows[k] - w, bins=edges)
        hist[k] = h.tolist()
    result["hist_edges"] = edges.tolist()
    result["hist"] = hist
    result["block"] = {"w": blk.tolist(), "q8": r8.tolist(), "q4": r4.tolist()}
    if out_path:
        with open(out_path, "w") as f:
            json.dump(result, f, indent=1)


if __name__ == "__main__":
    main()

"""3 つのモデルで、プロンプト処理と生成の速さ（tokens/s）を繰り返し測る。

Ollama の応答の prompt_eval_count / prompt_eval_duration と
eval_count / eval_duration から速さを出す。思考モードは "think": false で無効にする。
マシンの負荷がぶれるので、モデルを 1 回ずつ順に回す周回を ROUNDS 回繰り返し、中央値を採る。
各回の前後で /api/ps（ollama ps と同じ情報）と負荷平均を記録する。

使い方:
    OLLAMA_HOST=127.0.0.1:11435 python3 bench_speed.py speed.json
"""
import json
import os
import statistics
import sys
import time
import urllib.request

MODELS = ["qwen3:1.7b-q4_K_M", "qwen3:1.7b-q8_0", "qwen3:1.7b-fp16"]
ROUNDS = 9
BASE = "http://" + os.environ.get("OLLAMA_HOST", "127.0.0.1:11435")

# プロンプト処理を測るための長めの入力（チャットの書式を含めて 950 トークン）
NOTES = "\n".join(
    f"Log {i}: at {i % 24:02d}:{(i * 7) % 60:02d} the keeper checked lamp {i % 5}, "
    f"wrote the wind speed {(i * 13) % 40} knots, and refilled {(i * 3) % 9} liters of oil."
    for i in range(1, 25))
TASK = "Using the logs above, write a long story of at least 300 words about the lighthouse keeper."


def post(path, body):
    req = urllib.request.Request(BASE + path, data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=3600) as r:   # 他の作業の待ち行列で長く待つことがある
        return json.load(r)


def loaded():
    with urllib.request.urlopen(BASE + "/api/ps", timeout=60) as r:
        return [m["name"] for m in json.load(r)["models"]]


def generate(model, tag):
    # 先頭に毎回違う文字列を置き、前回のプロンプトのキャッシュを使わせない
    prompt = f"Request {tag}.\n{NOTES}\n\n{TASK}"
    return post("/api/chat", {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "think": False,
        "stream": False,
        "keep_alive": "10m",
        "options": {"temperature": 0, "seed": 1, "num_predict": 128, "num_ctx": 4096},
    })


def main():
    out_path = sys.argv[1]
    runs = []
    for model in MODELS:                       # 読み込みの時間を計測から外す
        generate(model, f"warmup-{model}")
    for rnd in range(ROUNDS):
        for model in MODELS:
            before = loaded()
            load1 = os.getloadavg()[0]
            res = generate(model, f"{rnd}-{model}-{time.time_ns()}")
            after = loaded()
            run = {
                "round": rnd, "model": model, "loadavg_1m": round(load1, 1),
                "ps_before": before, "ps_after": after,
                "prompt_tokens": res["prompt_eval_count"],
                # キャッシュから読んだトークン（返す版だけ）を除き、計算したトークンだけで割る
                "prompt_cached": res.get("prompt_eval_cached_count", 0),
                "prompt_tps": (res["prompt_eval_count"] - res.get("prompt_eval_cached_count", 0))
                / res["prompt_eval_duration"] * 1e9,
                "eval_tokens": res["eval_count"],
                "eval_tps": res["eval_count"] / res["eval_duration"] * 1e9,
                # 思考の文が混ざっていないか（混ざると eval_count に思考のトークンが入る）
                "thinking_chars": len(res["message"].get("thinking", "")),
                "think_tag": "<think>" in res["message"]["content"],
            }
            runs.append(run)
            others = sorted(set(before + after) - set(MODELS))
            print(f"round {rnd} {model:18s} prompt {run['prompt_tokens']:4d} tok "
                  f"(cached {run['prompt_cached']:3d}) "
                  f"{run['prompt_tps']:8.1f} tok/s | gen {run['eval_tokens']:4d} tok "
                  f"{run['eval_tps']:6.1f} tok/s | thinking {run['thinking_chars']} chars"
                  f"{' <think>!' if run['think_tag'] else ''} | load {load1:6.1f} | other models: "
                  f"{', '.join(others) or '-'}", flush=True)
    print()
    summary = {}
    for model in MODELS:
        rs = [r for r in runs if r["model"] == model]
        p = [r["prompt_tps"] for r in rs]
        g = [r["eval_tps"] for r in rs]
        summary[model] = {"prompt_median": statistics.median(p), "prompt_min": min(p),
                          "prompt_max": max(p), "eval_median": statistics.median(g),
                          "eval_min": min(g), "eval_max": max(g)}
        s = summary[model]
        print(f"{model:18s} prompt median {s['prompt_median']:7.1f} tok/s "
              f"({s['prompt_min']:.1f}-{s['prompt_max']:.1f}) | gen median "
              f"{s['eval_median']:5.1f} tok/s ({s['eval_min']:.1f}-{s['eval_max']:.1f})")
    with open(out_path, "w") as f:
        json.dump({"runs": runs, "summary": summary}, f, indent=1)


if __name__ == "__main__":
    main()

"""答えが 1 つに決まる問題を 100 問作り、problems.json に書く（本記事のために自作した問題集）。

算数 40 問は乱数の seed を固定して生成し、常識 60 問は手で書いた。
使い方: python3 make_problems.py problems.json
"""
import json
import random
import sys

KNOWLEDGE = [
    # (問題, 正解として認める表記)
    ("What is the capital city of Australia?", ["canberra"]),
    ("What is the capital city of Canada?", ["ottawa"]),
    ("What is the capital city of Brazil?", ["brasilia", "brasília"]),
    ("What is the capital city of Turkey?", ["ankara"]),
    ("What is the capital city of Egypt?", ["cairo"]),
    ("What is the capital city of Kenya?", ["nairobi"]),
    ("What is the capital city of Peru?", ["lima"]),
    ("What is the capital city of Norway?", ["oslo"]),
    ("What is the capital city of Vietnam?", ["hanoi", "ha noi"]),
    ("What is the capital city of New Zealand?", ["wellington"]),
    ("What is the capital city of Morocco?", ["rabat"]),
    ("What is the capital city of Argentina?", ["buenos aires"]),
    ("What is the capital city of South Korea?", ["seoul"]),
    ("What is the capital city of Thailand?", ["bangkok"]),
    ("What is the capital city of Nigeria?", ["abuja"]),
    ("What is the chemical symbol for gold?", ["au"]),
    ("What is the chemical symbol for silver?", ["ag"]),
    ("What is the chemical symbol for sodium?", ["na"]),
    ("What is the chemical symbol for potassium?", ["k"]),
    ("What is the chemical symbol for iron?", ["fe"]),
    ("What is the chemical symbol for tungsten?", ["w"]),
    ("What is the chemical symbol for lead?", ["pb"]),
    ("What is the chemical symbol for mercury?", ["hg"]),
    ("Which chemical element has atomic number 1?", ["hydrogen"]),
    ("Which chemical element has atomic number 6?", ["carbon"]),
    ("Which chemical element has atomic number 10?", ["neon"]),
    ("Which planet is the largest in the Solar System?", ["jupiter"]),
    ("Which planet is closest to the Sun?", ["mercury"]),
    ("Which planet is known as the Red Planet?", ["mars"]),
    ("How many legs does a spider have?", ["8", "eight"]),
    ("How many sides does a hexagon have?", ["6", "six"]),
    ("At sea level, what is the boiling point of water in degrees Celsius?", ["100"]),
    ("What is the freezing point of water in degrees Fahrenheit?", ["32"]),
    ("Who wrote the play Romeo and Juliet?", ["shakespeare"]),
    ("Who wrote the novel Pride and Prejudice?", ["austen"]),
    ("Who wrote the novel Nineteen Eighty-Four?", ["orwell"]),
    ("Who painted the Mona Lisa?", ["leonardo", "da vinci"]),
    ("Who painted The Starry Night?", ["van gogh"]),
    ("Who composed the Moonlight Sonata?", ["beethoven"]),
    ("What is the largest ocean on Earth?", ["pacific"]),
    ("What is the longest river in South America?", ["amazon"]),
    ("What is the tallest mountain on Earth above sea level?", ["everest"]),
    ("What is the largest hot desert in the world?", ["sahara"]),
    ("What is the hardest natural mineral?", ["diamond"]),
    ("Which gas do plants absorb from the air for photosynthesis?", ["carbon dioxide", "co2"]),
    ("How many days are in a leap year?", ["366"]),
    ("What is the smallest prime number?", ["2", "two"]),
    ("What is the main language spoken in Brazil?", ["portuguese"]),
    ("What is the currency of Japan?", ["yen"]),
    ("Who developed the theory of general relativity?", ["einstein"]),
    ("Who formulated the three laws of motion published in the Principia?", ["newton"]),
    ("Which organ pumps blood through the human body?", ["heart"]),
    ("How many bones are in the adult human body?", ["206"]),
    ("In which year did humans first land on the Moon?", ["1969"]),
    ("In which year did the Berlin Wall fall?", ["1989"]),
    ("What is the fastest land animal?", ["cheetah"]),
    ("Which programming language was created by Guido van Rossum?", ["python"]),
    ("What is the largest country in the world by area?", ["russia"]),
    ("How many keys does a standard piano have?", ["88"]),
    ("What is the longest bone in the human body?", ["femur"]),
]


def arithmetic(rng):
    items = []
    for _ in range(10):
        a, b = rng.randint(100, 999), rng.randint(100, 999)
        items.append((f"What is {a} + {b}?", [str(a + b)]))
    for _ in range(10):
        a, b = rng.randint(1000, 9999), rng.randint(100, 999)
        items.append((f"What is {a} - {b}?", [str(a - b)]))
    for _ in range(10):
        a, b = rng.randint(11, 99), rng.randint(11, 99)
        items.append((f"What is {a} * {b}?", [str(a * b)]))
    for _ in range(10):
        boxes, per = rng.randint(12, 40), rng.randint(6, 24)
        sold = rng.randint(10, boxes * per - 10)
        items.append((f"A shop has {boxes} boxes with {per} apples in each box. "
                      f"It sells {sold} apples. How many apples are left?",
                      [str(boxes * per - sold)]))
    return items


def main():
    rng = random.Random(20261110)
    problems = []
    for i, (q, a) in enumerate(arithmetic(rng)):
        problems.append({"id": f"math{i:02d}", "kind": "math", "question": q, "answers": a})
    for i, (q, a) in enumerate(KNOWLEDGE):
        problems.append({"id": f"know{i:02d}", "kind": "knowledge", "question": q, "answers": a})
    with open(sys.argv[1], "w") as f:
        json.dump(problems, f, indent=1, ensure_ascii=False)
    print(f"{len(problems)} problems")


if __name__ == "__main__":
    main()

"""同じ問題集を 3 つのモデルに解かせ、正答率を数える。

思考モード（thinking）は "think": false で無効にし、答えだけを出させる。
温度 0・seed 固定なので、同じ環境なら同じ答えが返る想定。

使い方:
    OLLAMA_HOST=127.0.0.1:11435 python3 eval_accuracy.py problems.json accuracy.json
"""
import json
import os
import re
import sys
import urllib.request

MODELS = ["qwen3:1.7b-q4_K_M", "qwen3:1.7b-q8_0", "qwen3:1.7b-fp16"]
SYSTEM = "Answer with only the final answer. Do not explain."
BASE = "http://" + os.environ.get("OLLAMA_HOST", "127.0.0.1:11435")


def chat(model, question):
    body = {
        "model": model,
        "messages": [{"role": "system", "content": SYSTEM},
                     {"role": "user", "content": question}],
        "think": False,                       # qwen3 の思考モードを無効にする
        "stream": False,
        "options": {"temperature": 0, "seed": 1, "num_predict": 128, "num_ctx": 4096},
    }
    req = urllib.request.Request(BASE + "/api/chat", data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=600) as r:
        res = json.load(r)
    return res["message"]["content"], res["message"].get("thinking", "")


def is_correct(kind, reply, answers):
    text = reply.lower().replace(",", "")
    if kind == "math":
        nums = re.findall(r"-?\d+", text)
        return bool(nums) and nums[-1] == answers[0]       # 最後に出た整数で判定する
    return any(re.search(r"(?<![a-z0-9])" + re.escape(a) + r"(?![a-z0-9])", text)
               for a in answers)


def main():
    problems = json.load(open(sys.argv[1]))
    out_path = sys.argv[2]
    result = {"models": {}}
    for model in MODELS:
        rows = []
        for p in problems:
            reply, thinking = chat(model, p["question"])
            ok = is_correct(p["kind"], reply, p["answers"])
            rows.append({"id": p["id"], "kind": p["kind"], "reply": reply.strip(),
                         "thinking_chars": len(thinking), "correct": ok})
        result["models"][model] = rows
        for kind in ("math", "knowledge"):
            sub = [r for r in rows if r["kind"] == kind]
            print(f"{model:20s} {kind:9s} {sum(r['correct'] for r in sub):3d} / {len(sub)}")
        total = sum(r["correct"] for r in rows)
        print(f"{model:20s} {'total':9s} {total:3d} / {len(rows)}"
              f"   (replies with thinking: {sum(r['thinking_chars'] > 0 for r in rows)})")
    # fp16 と答え（文字列）が一致した問題の数
    ref = result["models"][MODELS[-1]]
    for model in MODELS[:-1]:
        same = sum(a["reply"] == b["reply"] for a, b in zip(result["models"][model], ref))
        print(f"same reply as fp16: {model:20s} {same} / {len(ref)}")
    with open(out_path, "w") as f:
        json.dump(result, f, indent=1, ensure_ascii=False)


if __name__ == "__main__":
    main()