benchmarks/scripts/07_domain_eval.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

07_domain_eval.py - does a quantized Tessera keep what matters (Turkish, English, code, math)?  Everything is scored against a

reference model (the Q4), because the original checkpoint cannot be run on the test machines.



  run      56 checkable prompts sent to a llama-server (OpenAI API): 11 Python functions with hidden unit tests (3 asked in Turkish),

           16 math word problems (8 English + the same 8 in Turkish), 14 Turkish and 15 English language tasks (reading comprehension,

           Turkish alphabet order, i/I casing, vowel-harmony plurals, instruction following, extraction, cross-language answers).

           Problems with numbers are generated from a seed, so the answers are computed, not typed.  Greedy decoding.

  kld      llama-perplexity per domain over the held-out text of 06_build_eval_set.py: mean KL divergence and same-top-token rate

           against the reference model's saved logits (--save-base stores them for the reference)

  compare  table of result files; every task category as a % of the first file (the reference)

  selftest the checkers accept the reference answers and reject garbage; no model needed



  python 07_domain_eval.py run --bin <llama.cpp>/build/bin --model q4.gguf --ot '<-ot regex>' --name q4 --no-think     (starts llama-server)

  python 07_domain_eval.py run --bin <llama.cpp>/build/bin --model ours.gguf --ot '<-ot regex>' --name ours --no-think

  python 07_domain_eval.py run --endpoint http://127.0.0.1:8080 --name x --no-think          (a llama-server that is already running)

  python 07_domain_eval.py compare results_q4.json results_ours.json

  python 07_domain_eval.py kld --bin <llama.cpp>/build/bin --model q4.gguf --ot '<-ot regex>' --name q4 --save-base

  python 07_domain_eval.py kld --bin <llama.cpp>/build/bin --model ours.gguf --ot '<-ot regex>' --name ours



`run` EXECUTES THE PYTHON THE MODEL WRITES (10 s limit, 2 GB, empty temp directory, nothing else isolated): use Colab or a throw-away

machine.  56 tasks are a smoke test: a difference of a few tasks is noise (one task is about 9 points of the 11 code tasks).

"""
from __future__ import annotations

import argparse
import json
import math
import os
import random
import re
import shlex
import subprocess
import sys
import tempfile
import time
import unicodedata
import urllib.request
from fractions import Fraction as Fr
from pathlib import Path

CATS = ("code", "math", "tr", "en", "reason", "lcb", "crux")
TR_ALPHA = "abcçdefgğhıijklmnoöprsştuüvyz"


# ---------------------------------------------------------------------------------------------------------------- checkers
def strip_think(r: str) -> str:
    return re.sub(r"<think>.*", "", re.sub(r"<think>.*?</think>", "", r, flags=re.S), flags=re.S).strip()


def numbers(text: str) -> list:
    out = []
    for tok in re.findall(r"-?\d[\d.,]*", text):
        tok = tok.rstrip(".,")
        tok = re.sub(r"[.,]", "", tok) if re.fullmatch(r"-?\d{1,3}([.,]\d{3})+", tok) else tok.replace(",", ".")
        try:
            out.append(Fr(tok))
        except ValueError:
            pass
    return out


def num(ans):
    def f(r):
        if "####" in r:
            n = numbers(r.rsplit("####", 1)[1])
            return bool(n) and n[0] == ans
        n = numbers(r)
        return bool(n) and n[-1] == ans
    return f


def has(*words):               # every word present (case-insensitive)
    return lambda r: all(w.lower() in r.lower() for w in words)


def ordered(words):            # the words appear in this order
    def f(r):
        pos = 0
        for w in words:
            i = r.lower().find(w.lower(), pos)
            if i < 0:
                return False
            pos = i + len(w)
        return True
    return f


def word(*ok):                 # the whole answer is one of these words
    return lambda r: re.sub(r"[^\w]", "", r.casefold()) in ok


def bullets(n):
    def f(r):
        lines = [l for l in r.splitlines() if l.strip()]
        return len(lines) == n and all(l.startswith("- ") for l in lines)
    return f


def deaccent(s: str) -> str:
    return "".join(c for c in unicodedata.normalize("NFD", s) if unicodedata.category(c) != "Mn").lower()


def as_json(want: dict):
    def f(r):
        m = re.search(r"\{.*\}", r, flags=re.S)
        try:
            d = json.loads(m.group(0))
        except Exception:
            return False
        return all(deaccent(str(d.get(k, ""))) == deaccent(v) for k, v in want.items())
    return f


def english_only(*words):      # contains the words and no Turkish-specific letter
    return lambda r: has(*words)(r) and not any(c in r for c in "çğıöşüÇĞİÖŞÜ")


def code_of(r: str) -> str:
    blocks = [b for b in re.findall(r"```(?:python|py)?[ \t]*\n(.*?)```", r, flags=re.S) if re.search(r"^\s*(def|class) ", b, re.M)]
    return blocks[0] if blocks else r


def run_code(code: str, tests: str, timeout: int = 10) -> bool:
    pre = None
    if os.name == "posix":
        import resource
        pre = lambda: resource.setrlimit(resource.RLIMIT_AS, (2 << 30, 2 << 30))
    with tempfile.TemporaryDirectory() as d:
        script = Path(d) / "task.py"      # a file, not -c: the HumanEval+ tests are bigger than the 128 KiB limit of one argument
        script.write_text(code + "\n\n" + tests, encoding="utf-8")
        try:
            return subprocess.run([sys.executable, "-I", str(script)], cwd=d, capture_output=True, timeout=timeout, preexec_fn=pre,
                                  env={**os.environ, "OPENBLAS_NUM_THREADS": "1", "OMP_NUM_THREADS": "1"}).returncode == 0
        except subprocess.TimeoutExpired:
            return False


# ---------------------------------------------------------------------------------------------------------------- tasks
ONE_BLOCK_EN = " Reply with one Python code block only."
ONE_BLOCK_TR = " Yalnızca tek bir Python kod bloğu ile cevap ver."
CODE = [  # (id, prompt, tests, reference solution)
    ("roman", "Write a Python function `roman_to_int(s)` that converts a Roman numeral (I V X L C D M, with subtractive pairs such as "
     "IV and IX) to an integer." + ONE_BLOCK_EN,
     "assert roman_to_int('III')==3 and roman_to_int('IV')==4 and roman_to_int('XLII')==42 and roman_to_int('MCMXCIV')==1994",
     "def roman_to_int(s):\n    v = {'I': 1, 'V': 5, 'X': 10, 'L': 50, 'C': 100, 'D': 500, 'M': 1000}\n    return sum(-v[c] if i + 1 < len(s) and "
     "v[c] < v[s[i + 1]] else v[c] for i, c in enumerate(s))"),
    ("merge_tr", "Python'da `merge_intervals(intervals)` fonksiyonunu yaz: [başlangıç, bitiş] listelerinden oluşan bir listeyi alır, "
     "kesişen veya uç uca değen aralıkları birleştirir ve başlangıca göre sıralı bir liste (iç içe listeler olarak) döndürür." + ONE_BLOCK_TR,
     "assert merge_intervals([[1,3],[2,6],[8,10],[15,18]])==[[1,6],[8,10],[15,18]]\nassert merge_intervals([[1,4],[4,5]])==[[1,5]]\n"
     "assert merge_intervals([])==[]\nassert merge_intervals([[5,6],[1,2]])==[[1,2],[5,6]]",
     "def merge_intervals(iv):\n    out = []\n    for a, b in sorted(iv):\n        if out and a <= out[-1][1]:\n"
     "            out[-1][1] = max(out[-1][1], b)\n        else:\n            out.append([a, b])\n    return out"),
    ("balanced_tr", "Python'da `is_balanced(s)` fonksiyonunu yaz: s içindeki (), [] ve {} parantezleri doğru sırada kapatılmışsa True, "
     "değilse False döndürsün (s yalnızca bu altı karakterden oluşur, boş dize True'dur)." + ONE_BLOCK_TR,
     "assert is_balanced('([]{})') and not is_balanced('(]') and is_balanced('') and not is_balanced('((') and not is_balanced('())')"
     " and not is_balanced('([)]')",
     "def is_balanced(s):\n    st, pair = [], {')': '(', ']': '[', '}': '{'}\n    for c in s:\n        if c in pair:\n"
     "            if not st or st.pop() != pair[c]:\n                return False\n        else:\n            st.append(c)\n    return not st"),
    ("flatten", "Write a Python function `flatten(x)` that turns an arbitrarily nested list into a flat list, keeping the order. "
     "Strings are single items, not lists." + ONE_BLOCK_EN,
     "assert flatten([1,[2,[3,[4]]],5])==[1,2,3,4,5] and flatten([])==[] and flatten([[],[[]]])==[] and flatten(['ab',[['c']]])==['ab','c']",
     "def flatten(x):\n    out = []\n    for i in x:\n        out += flatten(i) if isinstance(i, list) else [i]\n    return out"),
    ("word_freq", "Write a Python function `word_freq(text)` that returns a dict mapping each lowercase word to its count; words are the "
     "maximal runs of ASCII letters after lowercasing, everything else is a separator." + ONE_BLOCK_EN,
     "assert word_freq('Hello, hello world!')=={'hello':2,'world':1} and word_freq('')=={} and word_freq(\"It's it\")=={'it':2,'s':1}",
     "import re\nfrom collections import Counter\ndef word_freq(t):\n    return dict(Counter(re.findall('[a-z]+', t.lower())))"),
    ("rle", "Write a Python function `rle_decode(s)` that decodes run-length encoding: each run is a decimal count (one or more digits) "
     "followed by one non-digit character, e.g. '3a2b1c' -> 'aaabbc' and '12x' -> 'xxxxxxxxxxxx'." + ONE_BLOCK_EN,
     "assert rle_decode('3a2b1c')=='aaabbc' and rle_decode('')=='' and rle_decode('12x')=='x'*12 and rle_decode('10a1b')=='a'*10+'b'",
     "import re\ndef rle_decode(s):\n    return ''.join(c * int(n) for n, c in re.findall(r'(\\d+)(\\D)', s))"),
    ("moving_avg", "This function should return the averages of every window of k consecutive items of xs (an empty list when k > len(xs)), "
     "but it has a bug. Return the corrected function.\n\n```python\ndef moving_average(xs, k):\n    out = []\n    for i in range(len(xs) - k):\n"
     "        out.append(sum(xs[i:i + k]) / k)\n    return out\n```" + ONE_BLOCK_EN,
     "assert moving_average([1,2,3,4,5],2)==[1.5,2.5,3.5,4.5] and moving_average([1,2,3],3)==[2.0] and moving_average([1,2],3)==[]",
     "def moving_average(xs, k):\n    return [sum(xs[i:i + k]) / k for i in range(len(xs) - k + 1)]"),
    ("second_tr", "Python'da `ikinci_buyuk(xs)` fonksiyonunu yaz: sayı listesindeki en büyük ikinci FARKLI değeri döndürsün; böyle bir değer "
     "yoksa None döndürsün." + ONE_BLOCK_TR,
     "assert ikinci_buyuk([3,1,4,1,5])==4 and ikinci_buyuk([2,2]) is None and ikinci_buyuk([7]) is None and ikinci_buyuk([5,5,3])==3 "
     "and ikinci_buyuk([]) is None",
     "def ikinci_buyuk(xs):\n    s = sorted(set(xs))\n    return s[-2] if len(s) > 1 else None"),
    ("two_sum", "Write a Python function `two_sum(nums, target)` that returns a list [i, j] with i < j and nums[i] + nums[j] == target "
     "(exactly one such pair exists)." + ONE_BLOCK_EN,
     "assert two_sum([2,7,11,15],9)==[0,1] and two_sum([3,2,4],6)==[1,2] and two_sum([3,3],6)==[0,1]",
     "def two_sum(nums, target):\n    seen = {}\n    for j, x in enumerate(nums):\n        if target - x in seen:\n"
     "            return [seen[target - x], j]\n        seen[x] = j"),
    ("ipv4", "Write a Python function `valid_ipv4(s)` that returns True only for a dotted-quad IPv4 address: exactly four decimal parts "
     "separated by dots, digits only, each 0-255, no leading zeros (except the single digit '0')." + ONE_BLOCK_EN,
     "assert valid_ipv4('192.168.1.1') and valid_ipv4('0.0.0.0') and not valid_ipv4('256.1.1.1') and not valid_ipv4('1.2.3') "
     "and not valid_ipv4('01.2.3.4') and not valid_ipv4('1.2.3.4.5') and not valid_ipv4('a.b.c.d') and not valid_ipv4('1..2.3')",
     "def valid_ipv4(s):\n    p = s.split('.')\n    return len(p) == 4 and all(x.isascii() and x.isdigit() and (x == '0' or x[0] != '0') "
     "and int(x) <= 255 for x in p)"),
    ("lru", "Write a Python class `LRUCache(capacity)` with `get(key)` (returns the value, or -1 if absent) and `put(key, value)`. When "
     "a put would exceed the capacity, evict the least recently used key; both get and put count as a use." + ONE_BLOCK_EN,
     "c = LRUCache(2); c.put(1,1); c.put(2,2); assert c.get(1)==1; c.put(3,3); assert c.get(2)==-1; c.put(4,4); assert c.get(1)==-1\n"
     "assert c.get(3)==3 and c.get(4)==4\nc = LRUCache(1); c.put(1,1); c.put(1,5); assert c.get(1)==5; c.put(2,2); assert c.get(1)==-1",
     "from collections import OrderedDict\nclass LRUCache:\n    def __init__(self, capacity):\n        self.c, self.n = OrderedDict(), capacity\n"
     "    def get(self, k):\n        if k not in self.c:\n            return -1\n        self.c.move_to_end(k)\n        return self.c[k]\n"
     "    def put(self, k, v):\n        self.c[k] = v\n        self.c.move_to_end(k)\n        if len(self.c) > self.n:\n"
     "            self.c.popitem(last=False)"),
]


# math: each template returns (English, Turkish, answer); the numbers are drawn so that the answers are whole and sensible
def m_shop(r):
    p, n, a, b = r.randint(2, 9), r.randint(5, 15), r.randint(2, 9), r.randint(2, 6)
    spend = a * p + b * n
    m = (spend // 10 + 1) * 10 + r.choice([0, 10, 20])
    return (f"A shop sells pencils for {p} dollars each and notebooks for {n} dollars each. Maria buys {a} pencils and {b} notebooks "
            f"and pays with a {m}-dollar bill. How many dollars of change does she get?",
            f"Bir kırtasiyede kalem tanesi {p} lira, defter tanesi {n} liradır. Maria {a} kalem ve {b} defter alıp {m} liralık "
            f"banknotla ödeme yapıyor. Kaç lira para üstü alır?", m - spend)


def m_train(r):
    v1, t1, v2, t2 = r.randint(40, 90), r.randint(2, 5), r.randint(60, 120), r.randint(1, 4)
    return (f"A train travels at {v1} km/h for {t1} hours and then at {v2} km/h for {t2} hours. How many kilometers does it cover in total?",
            f"Bir tren saatte {v1} km hızla {t1} saat, ardından saatte {v2} km hızla {t2} saat gidiyor. Toplam kaç km yol alır?",
            v1 * t1 + v2 * t2)


def m_discount(r):
    x, d, c = 20 * r.randint(5, 20), r.choice([10, 15, 20, 25, 30, 40]), r.randint(3, 15)
    return (f"A jacket costs {x} dollars. It is discounted by {d}% and then a coupon takes another {c} dollars off the discounted price. "
            f"What is the final price in dollars?",
            f"Bir montun fiyatı {x} liradır. Önce %{d} indirim yapılıyor, sonra indirimli fiyattan {c} lira daha düşülüyor. "
            f"Son fiyat kaç liradır?", Fr(x * (100 - d), 100) - c)


def m_age(r):
    a, x, n = r.randint(8, 30), r.randint(2, 9), r.randint(2, 8)
    s = 2 * a + x + 2 * n
    return (f"Tom is {x} years older than Ann. In {n} years the sum of their ages will be {s}. How old is Ann now?",
            f"Ali, Ayşe'den {x} yaş büyüktür. {n} yıl sonra ikisinin yaşları toplamı {s} olacaktır. Ayşe şimdi kaç yaşındadır?", a)


def m_work(r):
    while True:
        a, b = r.choice([2, 3, 4, 5, 6, 10, 12, 15]), r.choice([2, 3, 4, 5, 6, 10, 12, 15])
        mins = Fr(60 * a * b, a + b)
        if a != b and mins.denominator == 1:
            return (f"Pipe A can fill a tank in {a} hours and pipe B can fill it in {b} hours. Working together, how many minutes do "
                    f"they need to fill the empty tank?",
                    f"A musluğu bir depoyu {a} saatte, B musluğu {b} saatte doldurabiliyor. İkisi birlikte çalışırsa boş depoyu kaç "
                    f"dakikada doldurur?", mins)


def m_rect(r):
    w = r.randint(3, 20)
    l, p = 2 * w + 3, 2 * (3 * w + 3)
    return (f"The length of a rectangle is 3 more than twice its width. Its perimeter is {p} cm. What is its area in square cm?",
            f"Bir dikdörtgenin uzun kenarı, kısa kenarının iki katından 3 cm fazladır. Çevresi {p} cm olduğuna göre alanı kaç cm²'dir?",
            w * l)


def m_ratio(r):
    a, b, c, k = r.randint(1, 4), r.randint(2, 6), r.randint(3, 8), r.randint(4, 15)
    s = k * (a + b + c)
    return (f"{s} dollars are shared among three friends in the ratio {a}:{b}:{c}. How many dollars does the friend with the largest "
            f"share get?",
            f"{s} lira, üç arkadaş arasında {a}:{b}:{c} oranında paylaştırılıyor. En büyük payı alan kaç lira alır?", k * max(a, b, c))


def m_codes(r):
    n = r.randint(5, 9)
    return (f"How many different 3-letter codes can be made from {n} distinct letters if no letter may be used twice in a code?",
            f"{n} farklı harften, hiçbir harf bir şifrede iki kez kullanılmamak koşuluyla, kaç farklı 3 harfli şifre oluşturulabilir?",
            n * (n - 1) * (n - 2))


MATH = [m_shop, m_train, m_discount, m_age, m_work, m_rect, m_ratio, m_codes]
HOW_EN = "\n\nThink step by step, then give the final answer on the last line as: #### <number>"
HOW_TR = "\n\nAdım adım düşün, sonra son satırda nihai cevabı şu biçimde yaz: #### <sayı>"
TR_PLURAL = [("kitap", "kitaplar"), ("göz", "gözler"), ("okul", "okullar"), ("kalem", "kalemler"), ("çiçek", "çiçekler"),
             ("araba", "arabalar"), ("kedi", "kediler"), ("ağaç", "ağaçlar"), ("öğretmen", "öğretmenler"), ("dağ", "dağlar"),
             ("deniz", "denizler"), ("masa", "masalar"), ("kuş", "kuşlar"), ("gül", "güller")]
TR_WORDS = ["cam", "çanta", "dağ", "ışık", "iyi", "orman", "ödev", "sandalye", "şeker", "uzay", "ümit", "zeytin", "gül", "hava", "kalem",
            "lale", "yol", "yüz", "üzüm", "masa", "pencere", "çiçek", "şehir", "ılık", "ince", "öğle", "göl", "fırın", "tarla"]
MONTHS_EN = ["January", "March", "May", "July", "September", "November"]
MONTHS_TR = ["Ocak", "Mart", "Mayıs", "Temmuz", "Eylül", "Kasım"]
CITIES = ["Izmir", "Bursa", "Konya", "Trabzon", "Antalya", "Adana"]


def tr_key(w: str):
    return [TR_ALPHA.index(c) for c in w]


def build(seed: int = 1) -> list:
    """[(id, category, prompt, check(response) -> bool, reference response)]"""
    T = []
    for tid, prompt, tests, ref in CODE:
        T.append((f"code_{tid}", "code", prompt, (lambda t: lambda r: run_code(code_of(r), t))(tests), f"```python\n{ref}\n```"))
    for i, f in enumerate(MATH):
        en, tr, ans = f(random.Random(f"{seed}m{i}"))
        T.append((f"math_{f.__name__[2:]}_en", "math", en + HOW_EN, num(ans), f"#### {ans}"))
        T.append((f"math_{f.__name__[2:]}_tr", "math", tr + HOW_TR, num(ans), f"#### {ans}"))
    # Turkish
    for i in range(3):
        r = random.Random(f"{seed}ta{i}")
        a, b, c = r.randint(2, 7), r.randint(2, 6), r.randint(15, 40)
        T.append((f"tr_age{i}", "tr", f"Zeynep, Ahmet'ten {a} yaş büyüktür. Ahmet, Deniz'den {b} yaş küçüktür. Deniz {c} yaşındadır. "
                  f"Zeynep kaç yaşındadır? Yalnızca sayıyı yaz.", num(c - b + a), str(c - b + a)))
    for i in range(2):
        ws = random.Random(f"{seed}ts{i}").sample(TR_WORDS, 6)
        want = sorted(ws, key=tr_key)
        T.append((f"tr_sort{i}", "tr", "Şu kelimeleri Türk alfabesine göre alfabetik sıraya diz ve aralarına virgül koyarak tek satırda yaz: "
                  + ", ".join(ws), ordered(want), ", ".join(want)))
    ph = random.Random(f"{seed}tu").choice(["istanbul ışığı", "ilkbahar ışıkları", "iğne ve ipek", "ılık süt içtim", "çiçekli bir bahçe"])
    up = ph.replace("i", "İ").replace("ı", "I").upper()
    T.append(("tr_upper", "tr", f"'{ph}' ifadesini Türkçe yazım kurallarına göre tamamen BÜYÜK harfe çevir. Yalnızca sonucu yaz.",
              lambda r, up=up: up in r, up))
    for i in range(2):
        pairs = random.Random(f"{seed}tp{i}").sample(TR_PLURAL, 5)
        T.append((f"tr_plural{i}", "tr", "Şu kelimelerin çoğul hâlini yaz (her biri için yalnızca çoğul biçimi, virgülle ayırarak): "
                  + ", ".join(w for w, _ in pairs), has(*[p for _, p in pairs]), ", ".join(p for _, p in pairs)))
    T.append(("tr_bullets", "tr", "Ankara hakkında tam olarak 3 madde yaz. Her madde satır başında '- ' ile başlasın, başka hiçbir şey yazma.",
              bullets(3), "- Başkenttir.\n- İç Anadolu'dadır.\n- Kalabalıktır."))
    T.append(("tr_yesno", "tr", "Yalnızca 'evet' veya 'hayır' yaz: Türkiye'nin başkenti İstanbul mu?", word("hayır", "hayir"), "hayır"))
    T.append(("tr_translate", "tr", "Şu cümleyi İngilizceye çevir, yalnızca çeviriyi yaz: 'Yarın sabah erken kalkıp işe gideceğim.'",
              lambda r: sum(bool(re.search(p, r.lower())) for p in (r"tomorrow", r"morning", r"early", r"\bwork|\bjob|office",
                                                                    r"wake|get up|rise|getting up")) >= 4,
              "Tomorrow morning I will get up early and go to work."))
    for i in range(2):
        r = random.Random(f"{seed}te{i}")
        day, mon, hh, mm, city = r.randint(2, 28), r.choice(MONTHS_TR), r.randint(9, 17), r.choice(["00", "15", "30", "45"]), r.choice(CITIES)
        txt = f"Toplantı {day} {mon} 2025 tarihinde, saat {hh}:{mm} olarak, {city} şubesinde yapılacaktır."
        if i == 0:
            T.append((f"tr_extract{i}", "tr", f"Metin: {txt}\nSoru: Toplantı saat kaçta? Yalnızca saati yaz.",
                      lambda r, a=f"{hh}:{mm}", b=f"{hh}.{mm}": a in r or b in r, f"{hh}:{mm}"))
        else:
            T.append((f"tr_extract{i}", "tr", f"Metin: {txt}\nSoru: Toplantı hangi şehirde yapılacak? Yalnızca şehir adını yaz.",
                      has(city), city))
    T.append(("tr_cross", "tr", "İngilizce olarak cevap ver, tek cümle: Fransa'nın başkenti neresidir?", english_only("paris"),
              "The capital of France is Paris."))
    # English
    for i in range(3):
        r = random.Random(f"{seed}ea{i}")
        a, b, c = r.randint(2, 7), r.randint(2, 6), r.randint(15, 40)
        T.append((f"en_age{i}", "en", f"Emma is {a} years older than Noah. Noah is {b} years younger than Liam. Liam is {c} years old. "
                  f"How old is Emma? Answer with just the number.", num(c - b + a), str(c - b + a)))
    T.append(("en_bullets", "en", "Write exactly 3 bullet points about Paris. Start each line with '- ' and write nothing else.",
              bullets(3), "- Capital of France.\n- On the Seine.\n- Home of the Louvre."))
    T.append(("en_yesno1", "en", "Answer with only 'yes' or 'no': Is the Pacific Ocean the largest ocean on Earth?", word("yes"), "yes"))
    T.append(("en_yesno2", "en", "Answer with only 'yes' or 'no': Is 91 a prime number?", word("no"), "no"))
    T.append(("en_reverse", "en", "Reverse the order of the words in this sentence and output only the result: 'the quick brown fox jumps'",
              lambda r: "jumps fox brown quick the" in re.sub(r"[^\w ]", "", r.lower()), "jumps fox brown quick the"))
    for i in range(2):
        r = random.Random(f"{seed}ee{i}")
        day, mon, hh, mm, city = r.randint(2, 28), r.choice(MONTHS_EN), r.randint(9, 17), r.choice(["00", "15", "30", "45"]), r.choice(CITIES)
        txt = f"The meeting will take place on {mon} {day}, 2025 at {hh}:{mm} in the {city} office."
        if i == 0:
            T.append((f"en_extract{i}", "en", f"Text: {txt}\nQuestion: At what time is the meeting? Answer with the time only.",
                      has(f"{hh}:{mm}"), f"{hh}:{mm}"))
        else:
            T.append((f"en_extract{i}", "en", f"Text: {txt}\nQuestion: In which city is the meeting? Answer with the city only.", has(city), city))
    T.append(("en_logic1", "en", "All bloops are razzies. All razzies are lazzies. Is every bloop a lazzie? Answer only yes or no.",
              word("yes"), "yes"))
    T.append(("en_logic2", "en", "All birds have feathers. Penguins have feathers. Does it follow that penguins are birds? Answer only yes or no.",
              word("no"), "no"))
    T.append(("en_json1", "en", "Give the capital of Japan as a JSON object with keys \"city\" and \"country\". Output only the JSON.",
              as_json({"city": "Tokyo", "country": "Japan"}), '{"city": "Tokyo", "country": "Japan"}'))
    T.append(("en_json2", "en", "Give the capital of Canada as a JSON object with keys \"city\" and \"country\". Output only the JSON.",
              as_json({"city": "Ottawa", "country": "Canada"}), '{"city": "Ottawa", "country": "Canada"}'))
    T.append(("en_cross", "en", "Answer in Turkish, in one sentence: What is the capital of Turkey?",
              lambda r: "ankara" in r.lower() and ("başkent" in r.lower() or "baskent" in r.lower()), "Türkiye'nin başkenti Ankara'dır."))
    T.append(("en_words", "en", "Summarize in at most 12 words: 'The city council voted on Tuesday to build a new bridge over the river, "
              "after years of debate about the cost.'", lambda r: len(r.split()) <= 12 and "bridge" in r.lower(),
              "Council votes to build new river bridge after years of debate."))
    return T


# ---------------------------------------------------------------------------------------------------------------- commands
def table(rows, head):
    w = [max(len(str(x)) for x in col) for col in zip(head, *rows)]
    for r in [head] + rows:
        print("  ".join(str(x).ljust(n) for x, n in zip(r, w)))


def serve(a):
    """starts llama-server on the model and returns once /health says ok"""
    cmd = [str(Path(a.bin) / "llama-server"), "-m", a.model, "-ngl", "99", "-fa", "on", "-lzm", "on", "-t", str(a.threads), "-c", str(a.ctx),
           "-np", "1", "--host", "127.0.0.1", "--port", str(a.port)] + (["-ot", a.ot] if a.ot else []) + shlex.split(a.extra)
    log = open(f"server_{a.name}.log", "w")
    p = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)
    t0 = time.time()
    while time.time() - t0 < a.load_timeout:
        if p.poll() is not None:
            sys.exit(f"llama-server exited early, see {log.name}")
        try:
            if urllib.request.urlopen(f"http://127.0.0.1:{a.port}/health", timeout=5).status == 200:
                print(f"llama-server ready after {time.time() - t0:.0f} s", flush=True)
                return p
        except Exception:
            time.sleep(5)
    p.terminate()
    sys.exit(f"llama-server not ready after {a.load_timeout} s, see {log.name}")


def cmd_run(a):
    srv = serve(a) if a.model else None
    if srv:
        a.endpoint = f"http://127.0.0.1:{a.port}"
    try:
        run_tasks(a)
    finally:
        if srv:
            srv.terminate()
            srv.wait(60)


def run_tasks(a):
    if a.suite != "basic":     # public benchmarks from 08_hard_tasks.py / 09_big_suite.py / 10_code_hard.py; long answers need room
        sys.path.insert(0, str(Path(__file__).resolve().parent))
        import importlib
        mod, fn = {"hard": ("08_hard_tasks", "build_hard"), "big": ("09_big_suite", "build_big"), "codehard": ("10_code_hard", "build_codehard")}[a.suite]
        tasks, a.max_tokens = getattr(importlib.import_module(mod), fn)(a.seed, a.hard_n, run_code, code_of, num), max(a.max_tokens, 2048)
    else:
        tasks = build(a.seed)
    tasks = [t for t in tasks if t[1] in a.only.split(",")]
    if a.redo:        # only the tasks that file failed (math and reasoning first): does a setting rescue them?
        prev = json.loads(Path(a.redo).read_text(encoding="utf-8"))["results"]
        order = {"math": 0, "reason": 1, "code": 2}
        tasks = sorted([t for t in tasks if t[0] in prev and not prev[t[0]]["pass"]], key=lambda t: order.get(t[1], 3))[:a.redo_max]
    res, t0, toks = {}, time.time(), 0
    out = a.out or f"results_{a.name}.json"

    def save():       # after every task, so a dropped session keeps what is done
        Path(out).write_text(json.dumps({"kind": "tasks", "name": a.name, "seed": a.seed, "tokens": toks, "seconds": round(time.time() - t0),
                                         "results": res}, indent=1, ensure_ascii=False), encoding="utf-8")
    def ask(messages, v):      # one chat request -> (reply text, completion tokens)
        body = {"model": a.name, "messages": messages, "max_tokens": a.max_tokens,
                "temperature": a.temperature if a.temperature is not None else (a.temp if a.votes > 1 else 0), "seed": a.seed * 1000 + v}
        for key, val in (("top_p", a.top_p), ("top_k", a.top_k)):
            if val is not None:
                body[key] = val
        if a.no_think:
            body["chat_template_kwargs"] = {"enable_thinking": False}
        elif a.think:
            body["chat_template_kwargs"] = {"enable_thinking": True}
        req = urllib.request.Request(a.endpoint.rstrip("/") + "/v1/chat/completions", json.dumps(body).encode(),
                                     {"Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=a.timeout) as f:
                d = json.load(f)
            return d["choices"][0]["message"].get("content") or "", d.get("usage", {}).get("completion_tokens", 0)
        except Exception as e:
            return f"[request failed: {e}]", 0

    first = {k: v["resp"] for k, v in json.loads(Path(a.first_from).read_text(encoding="utf-8"))["results"].items()} if a.first_from else {}
    for k, (tid, cat, prompt, check, _) in enumerate(tasks, 1):
        resps, tt, tries = [], 0, 0
        for v in range(a.votes):
            text, n = (first[tid], 0) if tid in first else ask([{"role": "user", "content": prompt}], v)     # --first-from: reuse the stored first answer
            resps.append(text)
            tt += n
        oks = [bool(check(strip_think(r))) for r in resps]
        if a.retry and a.votes == 1 and not oks[0] and hasattr(check, "feedback"):   # repair loop: only visible sample tests are shown to the model
            msgs = [{"role": "user", "content": prompt}]
            for _try in range(a.retry):
                fb = check.feedback(strip_think(resps[0]))
                if fb is None:      # the visible tests pass: nothing more to tell it
                    break
                msgs += [{"role": "assistant", "content": strip_think(resps[0])}, {"role": "user", "content": fb}]
                resps[0], n = ask(msgs, 0)
                tt += n
                tries += 1
                oks = [bool(check(strip_think(resps[0])))]
                if oks[0]:
                    break
        if a.votes > 1 and cat == "math":     # self-consistency: the most common final number wins, ties go to the first sample
            ans = [(numbers(strip_think(r)) or [None])[-1] for r in resps]
            best = max(ans, key=lambda x: (ans.count(x), -ans.index(x)))
            ok = oks[ans.index(best)]
        else:                                 # no canonical answer to vote on: the majority of the samples must pass
            ok = sum(oks) * 2 > len(oks)
        toks += tt
        res[tid] = {"cat": cat, "pass": ok, "frac": sum(oks) / len(oks), "toks": tt, "resp": resps[0], **({"retries": tries} if a.retry else {})}
        print(f"{k:3d}/{len(tasks)} {tid:22s} {'PASS' if ok else 'fail'}  {sum(oks)}/{len(oks)}  {time.time() - t0:5.0f} s" + (f"  retries {tries}" if a.retry else ""), flush=True)
        save()
    print(f"{toks} tokens in {time.time() - t0:.0f} s -> {out}")
    cmd_compare(argparse.Namespace(files=[out]))


def cmd_compare(a):
    docs = [json.loads(Path(f).read_text(encoding="utf-8")) for f in a.files]
    tasks = [d for d in docs if d["kind"] == "tasks"]
    if tasks:
        def rate(d, c):
            r = [v["pass"] for v in d["results"].values() if c is None or v["cat"] == c]
            return sum(r), len(r)
        ref = tasks[0]
        rows = []
        for c in list(CATS) + [None]:
            if c and not rate(ref, c)[1]:
                continue
            row = [c or "ALL"]
            for d in tasks:
                k, n = rate(d, c)
                rk, rn = rate(ref, c)
                pct = f"{100 * (k / n) / (rk / rn):.0f}%" if n and rk else "-"
                row.append(f"{k}/{n}" + ("" if d is ref else f" ({pct} of ref)"))
            rows.append(row)
        rows.append(["time / tokens"] + [f"{d.get('seconds', 0)} s / {d.get('tokens', 0)} tok ({d.get('tokens', 0) / max(d.get('seconds', 1), 1):.1f} t/s)" for d in tasks])
        table(rows, ["tasks"] + [d["name"] + (" (ref)" if d is ref else "") for d in tasks])
        for d in tasks[1:]:       # paired sign test on the tasks only one of the two passes ("beats" needs p < 0.05 in its favour)
            ids = [i for i in ref["results"] if i in d["results"]]
            b = sum(ref["results"][i]["pass"] and not d["results"][i]["pass"] for i in ids)
            c = sum(d["results"][i]["pass"] and not ref["results"][i]["pass"] for i in ids)
            p = min(1.0, 2 * sum(math.comb(b + c, i) for i in range(min(b, c) + 1)) / 2 ** (b + c)) if b + c else 1.0
            print(f"paired {d['name']} vs {ref['name']}: only {ref['name']} passes {b}, only {d['name']} passes {c}, sign test p = {p:.3f}")
    klds = [d for d in docs if d["kind"] == "kld"]
    if klds:
        print()
        rows = [[dom] + [f"KLD {d['domains'][dom]['kld']:.3f}  top {d['domains'][dom]['top']:.1f}%  pplx {d['domains'][dom]['ppl_ratio']:.3f}"
                         if dom in d["domains"] else "-" for d in klds] for dom in ("tr", "en", "code", "math")]
        table(rows, ["domain"] + [d["name"] for d in klds])


def parse_kld(text: str) -> dict:
    g = lambda p: float(re.search(p, text).group(1))
    return {"kld": g(r"Mean\s+KLD:\s+([\d.]+)"), "top": g(r"Same top p:\s+([\d.]+)"),
            "ppl_ratio": g(r"Mean PPL\(Q\)/PPL\(base\)\s*:\s+([\d.]+)"), "ppl": g(r"Mean PPL\(Q\)\s*:\s+([\d.]+)")}


def cmd_kld(a):
    ev = Path(a.eval_dir)
    res = {}
    for dom in ("tr", "en", "code", "math"):
        base = ev / f"base_{dom}.kld"
        cmd = [str(Path(a.bin) / "llama-perplexity"), "-m", a.model, "-f", str(ev / f"eval_{dom}.txt"), "-c", "512", "-b", "512", "-ub", "512",
               "--chunks", str(a.chunks), "-ngl", "99", "-fa", "on", "-lzm", "on", "-t", str(a.threads)]
        cmd += (["-ot", a.ot] if a.ot else []) + shlex.split(a.extra) + ["--kl-divergence-base", str(base)]
        cmd += [] if a.save_base else ["--kl-divergence"]
        t0 = time.time()
        p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        if p.returncode != 0:
            sys.exit(f"{dom}: llama-perplexity failed ({p.returncode}):\n{p.stdout[-1500:]}")
        if a.save_base:
            m = re.search(r"Final estimate: PPL = ([\d.]+)", p.stdout)
            res[dom] = {"ppl": float(m.group(1)) if m else None}
            print(f"{dom:5s} base saved {base} ({time.time() - t0:.0f} s)  PPL {res[dom]['ppl']}", flush=True)
            continue
        try:
            res[dom] = parse_kld(p.stdout)
        except AttributeError:
            sys.exit(f"{dom}: no statistics in the output:\n{p.stdout[-1500:]}")
        r = res[dom]
        print(f"{dom:5s} KLD {r['kld']:.3f}  same top {r['top']:.1f}%  PPL ratio {r['ppl_ratio']:.3f}  ({time.time() - t0:.0f} s)", flush=True)
    if not a.save_base:
        out = a.out or f"kld_{a.name}.json"
        Path(out).write_text(json.dumps({"kind": "kld", "name": a.name, "domains": res}, indent=1), encoding="utf-8")
        print("->", out)


def cmd_selftest(a):
    T = build(a.seed)
    bad = [t[0] for t in T if not t[3](t[4])]
    junk = [t[0] for t in T if t[3]("zzz 0")]
    code_wrong = run_code("def roman_to_int(s):\n    return 0", CODE[0][2])
    sample = "Mean PPL(Q)                   :   3.059708 ±   0.061630\nMean PPL(Q)/PPL(base)         :   1.131630 ±   0.011859\n" \
             "Mean    KLD:   0.350462 ±   0.007497\nSame top p: 82.255 ± 0.378 %\n"
    k = parse_kld(sample)
    ids = [t[0] for t in T]
    assert len(set(ids)) == len(ids), "duplicate task ids"
    assert numbers("1.250 ve 3,5 ve 12") == [1250, Fr(7, 2), 12] and strip_think("<think>x</think>ok") == "ok" and strip_think("<think>x") == ""
    assert not bad and not junk and not code_wrong, (bad, junk, code_wrong)
    assert k == {"kld": 0.350462, "top": 82.255, "ppl_ratio": 1.13163, "ppl": 3.059708}, k
    n = {c: sum(t[1] == c for t in T) for c in CATS}
    print(f"selftest OK: {len(T)} tasks {n}; reference answers pass, garbage fails, kld parser reads the sample")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run")
    r.add_argument("--endpoint", default="http://127.0.0.1:8080")
    r.add_argument("--name", required=True)
    r.add_argument("--out")
    r.add_argument("--seed", type=int, default=1)
    r.add_argument("--only", default=",".join(CATS), help="categories, comma separated")
    r.add_argument("--suite", choices=("basic", "hard", "big", "codehard"), default="basic", help="hard: HumanEval+, MATH-500, MMLU-Pro (08_hard_tasks.py); big: four domains, nested pools (09_big_suite.py); codehard: LiveCodeBench + CRUXEval (10_code_hard.py, use --max-tokens 4096 --ctx 8192); use --out")
    r.add_argument("--hard-n", type=int, default=100, help="tasks per category in the hard / big suite (big: 60 = quick prefix, 300 = full)")
    r.add_argument("--max-tokens", type=int, default=768)
    r.add_argument("--think", action="store_true", help="send chat_template_kwargs enable_thinking=true")
    r.add_argument("--redo", help="results file: run only the tasks it failed (same --suite/--seed/--hard-n)")
    r.add_argument("--redo-max", type=int, default=1000)
    r.add_argument("--first-from", help="with --retry: take each task's first answer from this results file (its 'resp') instead of generating it again")
    r.add_argument("--retry", type=int, default=0, help="lcb tasks: after a failed answer, show the model its first failing VISIBLE sample test and ask again, up to N times (needs --ctx big enough for the extra turns)")
    r.add_argument("--temperature", type=float, help="sampling temperature (default: 0, or --temp when --votes > 1)")
    r.add_argument("--top-p", type=float)
    r.add_argument("--top-k", type=int)
    r.add_argument("--timeout", type=int, default=1200)
    r.add_argument("--no-think", action="store_true", help="send chat_template_kwargs enable_thinking=false")
    r.add_argument("--votes", type=int, default=1, help="samples per task; math: the most common final number wins, other categories: majority must pass")
    r.add_argument("--temp", type=float, default=0.7, help="sampling temperature when --votes > 1 ")
    r.add_argument("--model", help="start llama-server on this GGUF (needs --bin) instead of using --endpoint")
    r.add_argument("--bin", help="directory with llama-server")
    r.add_argument("--ot", default="", help="-ot overrides for the server")
    r.add_argument("--extra", default="", help="more llama-server flags, one string: --extra='--flag value'")
    r.add_argument("--threads", type=int, default=8)
    r.add_argument("--ctx", type=int, default=4096)
    r.add_argument("--port", type=int, default=18080)   # 8080 is taken on Colab
    r.add_argument("--load-timeout", type=int, default=1800)
    r.set_defaults(f=cmd_run)
    c = sub.add_parser("compare")
    c.add_argument("files", nargs="+", help="result files; the first task file is the reference")
    c.set_defaults(f=cmd_compare)
    k = sub.add_parser("kld")
    k.add_argument("--bin", required=True, help="directory with llama-perplexity")
    k.add_argument("--model", required=True)
    k.add_argument("--name", required=True)
    k.add_argument("--eval-dir", default="evalset")
    k.add_argument("--chunks", type=int, default=16)
    k.add_argument("--ot", default="", help="-ot overrides (which tensors stay on the CPU)")
    k.add_argument("--extra", default="", help="more llama-perplexity flags, one string: --extra='--flag value'")
    k.add_argument("--threads", type=int, default=8)
    k.add_argument("--save-base", action="store_true", help="this is the reference: store its logits")
    k.add_argument("--out")
    k.set_defaults(f=cmd_kld)
    s = sub.add_parser("selftest")
    s.add_argument("--seed", type=int, default=1)
    s.set_defaults(f=cmd_selftest)
    a = ap.parse_args()
    a.f(a)


if __name__ == "__main__":
    main()