inferencekit-cli 0.1.0__py3-none-any.whl

inferencekit/quality/fake.py

@@ -0,0 +1,329 @@
+"""A deterministic fake OpenAI-compatible server, for tests and dry runs of the quality pipeline.
+
+It serves a toy "model" whose next-token distribution is a fixed function of the last
+two tokens, so everything (greedy text, top-k logprobs, teacher-forced prompt logprobs,
+exact KL between two fakes) is reproducible and computable in closed form.
+
+Endpoints: ``GET /v1/models``, ``POST /tokenize`` (vLLM style, ``messages`` or
+``prompt``), ``POST /v1/chat/completions`` (``logprobs``/``top_logprobs``, vLLM's
+``return_tokens_as_token_ids`` and ``return_token_ids``), ``POST /v1/completions`` with
+a token-id prompt and vLLM's ``prompt_logprobs``.
+
+``perturb`` shifts the logits by a fixed pattern: a fake with ``perturb=0`` plays the
+reference; one with ``perturb > 0`` plays a lower-precision engine. ``sharpness`` scales
+the logits (> 1 gives the peaked next-token distributions of a real LLM). ``max_logprobs``
+mimics vLLM's cap (HTTP 400 above it); ``has_tokenize=False`` drops ``/tokenize`` (404), like
+servers without vLLM's tokenizer API. ``answers`` maps a substring of the user prompt
+to a canned reply (for GSM8K tests). The toy tokenizer reads ``<tN>`` as token N, so the
+text of a generation re-tokenizes to its own ids.
+
+``logprob_mode``: "ok"; "ignore" answers without logprobs (an engine with no logprob
+support); "refuse" answers HTTP 400 when logprobs are asked for.
+
+Sampling: a chat request with ``temperature > 0`` samples, with ``top_k``, ``top_p`` and
+``seed`` applied the usual way (temperature, then top-k, then top-p on the renormalised
+rest). ``seeds``: "honour" (a draw is a function of seed, prompt and position), "ignore"
+(a fresh draw per request, not reproducible), "fixed" (every seed gives the same draw).
+``rng_salt`` changes the random stream (a second engine with a correct sampler but other
+random numbers); ``wrong_top_p`` / ``wrong_temperature`` replace the requested top_p / temperature
+when sampling (sampler bugs; the logprobs stay right). ``top_p_order="joint"`` samples as SGLang
+does: top-p over the whole distribution, then the first top_k tokens. ``batch_shift`` moves the
+logits by ``batch_shift * (requests in flight - 1)`` times a fixed pattern, like an engine whose
+numerics change with the batch; ``max_inflight`` is the most requests seen in flight at once.
+
+Server facts for concurrency probes: ``server_info`` is served at ``GET /get_server_info``
+(SGLang) and ``props`` at ``GET /props`` (llama.cpp); ``slots`` caps how many generation
+requests run at once (the rest wait) and ``delay_s`` is how long each one takes.
+
+Usage::
+
+    with serve(FakeLM()) as base_url:
+        ...
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import math
+import random
+import re
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+
+class FakeLM:
+    BOS, ASSISTANT = 1, 2
+    _MARK = re.compile(r"<t(\d+)>")
+
+    def __init__(self, vocab: int = 16, perturb: float = 0.0, max_logprobs: int = 20, ids: str = "vllm",
+                 answers: dict | None = None, name: str = "fake-model", max_model_len: int = 4096,
+                 has_tokenize: bool = True, sharpness: float = 1.0, logprob_mode: str = "ok",
+                 seeds: str = "honour", rng_salt: str = "", wrong_top_p: float | None = None,
+                 server_info: dict | None = None, props: dict | None = None, slots: int | None = None,
+                 delay_s: float = 0.0, wrong_temperature: float | None = None, batch_shift: float = 0.0,
+                 top_p_order: str = "top_k_first"):
+        self.vocab, self.perturb, self.max_logprobs, self.ids = vocab, perturb, max_logprobs, ids
+        self.sharpness = sharpness
+        self.has_tokenize = has_tokenize
+        self.answers = answers or {}
+        self.name, self.max_model_len = name, max_model_len
+        self.logprob_mode, self.seeds, self.rng_salt, self.wrong_top_p = logprob_mode, seeds, rng_salt, wrong_top_p
+        self.wrong_temperature, self.batch_shift, self.top_p_order = wrong_temperature, batch_shift, top_p_order
+        self._inflight = self.max_inflight = 0
+        self._tl = threading.local()  # requests in flight when this thread's request arrived (batch_shift)
+        self.server_info, self.props, self.delay_s = server_info, props, delay_s
+        self._slots = threading.Semaphore(slots) if slots else None
+        self._draws = 0
+        self._lock = threading.Lock()
+        self.requests: list[dict] = []
+
+    # ----------------------------------------------------------------- the toy model
+    def logprobs(self, context: list[int]) -> list[float]:
+        a = context[-1] if context else 0
+        b = context[-2] if len(context) > 1 else 0
+        logits = [3.0 * math.sin(1.7 * a + 0.37 * b * j + 0.91 * j + 0.13 * j * j) for j in range(self.vocab)]
+        if self.perturb:
+            logits = [x + self.perturb * math.cos(2.3 * j + 0.7 * a) for j, x in enumerate(logits)]
+        batch = getattr(self._tl, "inflight", 1)
+        if self.batch_shift and batch > 1:
+            logits = [x + self.batch_shift * (batch - 1) * math.cos(1.3 * j + 0.4 * a) for j, x in enumerate(logits)]
+        if self.sharpness != 1.0:
+            logits = [self.sharpness * x for x in logits]
+        m = max(logits)
+        lse = m + math.log(sum(math.exp(x - m) for x in logits))
+        return [x - lse for x in logits]
+
+    def tokenize_text(self, text: str) -> list[int]:
+        def chars(part: str) -> list[int]:
+            return [3 + (ord(c) % (self.vocab - 3)) for c in part]
+
+        ids, pos = [], 0
+        for m in self._MARK.finditer(text):
+            if int(m.group(1)) < self.vocab:
+                ids += chars(text[pos:m.start()]) + [int(m.group(1))]
+                pos = m.end()
+        return ids + chars(text[pos:])
+
+    def template(self, messages: list[dict]) -> list[int]:
+        ids = [self.BOS]
+        for m in messages:
+            ids += self.tokenize_text(f"{m.get('role')}:{m.get('content')}")
+        return ids + [self.ASSISTANT]
+
+    def greedy(self, prompt: list[int], n: int) -> list[int]:
+        ctx, out = list(prompt), []
+        for _ in range(n):
+            lp = self.logprobs(ctx)
+            t = max(range(self.vocab), key=lambda j: lp[j])
+            out.append(t)
+            ctx.append(t)
+        return out
+
+    def top(self, lp: list[float], k: int) -> list[int]:
+        return sorted(range(self.vocab), key=lambda j: -lp[j])[:k]
+
+    def _rng(self, body: dict, prompt: list[int], step: int) -> random.Random:
+        where = f"{','.join(map(str, prompt))}|{step}"
+        if self.seeds == "ignore":
+            with self._lock:
+                self._draws += 1
+                return random.Random(f"{self.rng_salt}|draw|{self._draws}")
+        if self.seeds == "fixed":
+            return random.Random(f"{self.rng_salt}|fixed|{where}")
+        return random.Random(f"{self.rng_salt}|{body.get('seed')}|{where}")
+
+    def sample(self, prompt: list[int], n: int, body: dict) -> list[int]:
+        """Seeded sampling: temperature, then top-k, then top-p on the renormalised remainder (``joint``: top-p
+        over the whole distribution, then the first top_k)."""
+        temp = float(body["temperature"] if self.wrong_temperature is None else self.wrong_temperature)
+        top_k = int(body.get("top_k") or 0)
+        top_p = self.wrong_top_p if self.wrong_top_p is not None else body.get("top_p")
+        ctx, out = list(prompt), []
+        for step in range(n):
+            lp = self.logprobs(ctx)
+            m = max(lp)
+            w = {j: math.exp((x - m) / temp) for j, x in enumerate(lp)}
+            order = sorted(w, key=lambda j: (-w[j], j))
+            joint = self.top_p_order == "joint"
+            if top_k > 0 and not joint:
+                order = order[:top_k]
+            if top_p is not None and float(top_p) < 1.0:
+                total, kept, acc = sum(w[j] for j in order), [], 0.0
+                for j in order:
+                    kept.append(j)
+                    acc += w[j] / total
+                    if acc >= float(top_p):
+                        break
+                order = kept
+            if top_k > 0 and joint:
+                order = order[:top_k]
+            u = self._rng(body, prompt, step).random() * sum(w[j] for j in order)
+            t = order[-1]
+            for j in order:
+                u -= w[j]
+                if u < 0:
+                    t = j
+                    break
+            out.append(t)
+            ctx.append(t)
+        return out
+
+    # ----------------------------------------------------------------- endpoints
+    def _tokstr(self, t: int, as_ids: bool) -> str:
+        return f"token_id:{t}" if as_ids else f"<t{t}>"
+
+    def chat(self, body: dict) -> tuple[int, dict]:
+        messages = body.get("messages") or []
+        text = "\n".join(str(m.get("content")) for m in messages)
+        for key, reply in self.answers.items():
+            if key in text:
+                return 200, {"id": "fake", "object": "chat.completion", "model": self.name,
+                             "choices": [{"index": 0, "finish_reason": "stop",
+                                          "message": {"role": "assistant", "content": reply}}],
+                             "usage": {"prompt_tokens": len(text), "completion_tokens": len(reply.split())}}
+        k = int(body.get("top_logprobs") or 0)
+        if body.get("logprobs") and self.logprob_mode == "refuse":
+            return 400, {"error": {"message": "logprobs are not supported by this server", "code": 400}}
+        if body.get("logprobs") and k > self.max_logprobs:
+            return 400, {"error": {"message": f"Requested sample logprobs of {k}, which is greater than max "
+                                              f"allowed: {self.max_logprobs}", "code": 400}}
+        prompt = self.template(messages)
+        n = int(body.get("max_tokens") or 16)
+        gen = self.sample(prompt, n, body) if float(body.get("temperature") or 0) > 0 else self.greedy(prompt, n)
+        as_ids = bool(body.get("return_tokens_as_token_ids")) and self.ids == "vllm"
+        choice: dict = {"index": 0, "finish_reason": "length",
+                        "message": {"role": "assistant", "content": "".join(f"<t{t}>" for t in gen)}}
+        if body.get("logprobs") and self.logprob_mode == "ok":
+            content, ctx = [], list(prompt)
+            for t in gen:
+                lp = self.logprobs(ctx)
+                entry = {"token": self._tokstr(t, as_ids), "logprob": lp[t], "bytes": None,
+                         "top_logprobs": [{"token": self._tokstr(j, as_ids), "logprob": lp[j], "bytes": None}
+                                          for j in self.top(lp, k)]}
+                if self.ids == "llamacpp":
+                    entry["id"] = t
+                    for alt, j in zip(entry["top_logprobs"], self.top(lp, k)):
+                        alt["id"] = j
+                content.append(entry)
+                ctx.append(t)
+            choice["logprobs"] = {"content": content}
+        resp = {"id": "fake", "object": "chat.completion", "model": self.name, "choices": [choice],
+                "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(gen),
+                          "total_tokens": len(prompt) + len(gen)}}
+        if body.get("return_token_ids") and self.ids == "vllm":
+            choice["token_ids"] = gen
+            resp["prompt_token_ids"] = prompt
+        return 200, resp
+
+    def completions(self, body: dict) -> tuple[int, dict]:
+        prompt = body.get("prompt")
+        if not (isinstance(prompt, list) and all(isinstance(t, int) for t in prompt)):
+            prompt = self.tokenize_text(str(prompt))
+        if len(prompt) + int(body.get("max_tokens") or 1) > self.max_model_len:
+            return 400, {"error": {"message": "This model's maximum context length is exceeded", "code": 400}}
+        k = body.get("prompt_logprobs")
+        if k is not None and int(k) > self.max_logprobs:
+            return 400, {"error": {"message": f"Requested prompt logprobs of {k}, which is greater than max "
+                                              f"allowed: {self.max_logprobs}", "code": 400}}
+        choice: dict = {"index": 0, "text": "", "finish_reason": "length"}
+        if k is not None:
+            plp: list = [None]
+            for i in range(1, len(prompt)):
+                lp = self.logprobs(prompt[:i])
+                ranks = {j: r + 1 for r, j in enumerate(self.top(lp, self.vocab))}
+                keep = self.top(lp, int(k)) + [prompt[i]]
+                plp.append({str(j): {"logprob": lp[j], "rank": ranks[j], "decoded_token": f"<t{j}>"}
+                            for j in keep})
+            choice["prompt_logprobs"] = plp
+        gen = self.greedy(prompt, int(body.get("max_tokens") or 1))
+        choice["text"] = "".join(f"<t{t}>" for t in gen)
+        return 200, {"id": "fake", "object": "text_completion", "model": self.name, "choices": [choice],
+                     "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(gen)}}
+
+    def generate(self, body: dict) -> tuple[int, dict]:
+        """A chat request as the server runs it: wait for a slot, take ``delay_s``, answer."""
+        with self._lock:
+            self._inflight += 1
+            self._tl.inflight = self._inflight
+            self.max_inflight = max(self.max_inflight, self._inflight)
+        try:
+            return self._generate(body)
+        finally:
+            with self._lock:
+                self._inflight -= 1
+
+    def _generate(self, body: dict) -> tuple[int, dict]:
+        if self._slots is None:
+            if self.delay_s:
+                time.sleep(self.delay_s)
+            return self.chat(body)
+        with self._slots:
+            if self.delay_s:
+                time.sleep(self.delay_s)
+            return self.chat(body)
+
+    def tokenize(self, body: dict) -> tuple[int, dict]:
+        if "messages" in body:
+            ids = self.template(body["messages"])
+        else:
+            ids = self.tokenize_text(str(body.get("prompt", "")))
+        return 200, {"tokens": ids, "count": len(ids), "max_model_len": self.max_model_len}
+
+
+def _handler(lm: FakeLM):
+    class Handler(BaseHTTPRequestHandler):
+        def log_message(self, *args):  # keep test output quiet
+            pass
+
+        def _send(self, status: int, obj: dict):
+            data = json.dumps(obj).encode()
+            self.send_response(status)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(data)))
+            self.end_headers()
+            self.wfile.write(data)
+
+        def do_GET(self):
+            path = self.path.rstrip("/")
+            if path == "/v1/models":
+                self._send(200, {"object": "list", "data": [{"id": lm.name, "object": "model"}]})
+            elif path == "/get_server_info" and lm.server_info is not None:
+                self._send(200, lm.server_info)
+            elif path == "/props" and lm.props is not None:
+                self._send(200, lm.props)
+            else:
+                self._send(404, {"error": {"message": "not found"}})
+
+        def do_POST(self):
+            body = json.loads(self.rfile.read(int(self.headers.get("Content-Length") or 0)) or b"{}")
+            lm.requests.append({"path": self.path, "body": body})
+            routes = {"/v1/chat/completions": lm.generate, "/v1/completions": lm.completions}
+            if lm.has_tokenize:
+                routes["/tokenize"] = lm.tokenize
+            fn = routes.get(self.path)
+            if fn is None:
+                self._send(404, {"error": {"message": "not found"}})
+                return
+            status, obj = fn(body)
+            self._send(status, obj)
+
+    return Handler
+
+
+class _Server(ThreadingHTTPServer):
+    request_queue_size = 128  # many short concurrent connections (sampling histograms)
+
+
+@contextlib.contextmanager
+def serve(lm: FakeLM, host: str = "127.0.0.1"):
+    """Run ``lm`` on a free port in a background thread; yields the base URL."""
+    server = _Server((host, 0), _handler(lm))
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        yield f"http://{host}:{server.server_address[1]}"
+    finally:
+        server.shutdown()
+        server.server_close()

inferencekit/quality/gsm8k.py

@@ -0,0 +1,224 @@
+"""GSM8K: the pinned test split (downloaded once, cached), the zero-shot chain-of-thought prompt,
+and answer extraction.
+
+Data: the 1,319-problem test split of GSM8K (Cobbe et al. 2021,
+github.com/openai/grade-school-math, MIT licence in ``data/gsm8k_LICENSE.txt``). It is not
+shipped: the first ``quality gsm8k`` run downloads the Hugging Face dataset ``openai/gsm8k``
+(config ``main``, split ``test``) at a pinned revision, checks the Parquet file's sha256,
+converts it to JSONL (one ``{"question", "answer"}`` object per line, ``json.dumps``
+defaults) and checks that file's sha256 too. Those JSONL bytes are identical to the
+upstream ``test.jsonl`` and to the file early versions shipped, so ``dataset_sha256`` in evidence
+records is unchanged. The cache is ``~/.cache/inferencekit/gsm8k/`` (``$XDG_CACHE_HOME``
+and ``INFERENCEKIT_CACHE_DIR`` move it). The gold answer is the number after the last
+``####`` in ``answer``.
+
+Prompt (zero-shot, one user turn): the question, then an instruction to reason step by
+step and finish with ``#### <number>``.
+
+Extraction (first rule that finds a number wins; the rule used is recorded per item):
+  1. ``hash``:  the first number after the last ``####``;
+  2. ``boxed``: the number inside the last ``\\boxed{...}`` (many chat models use it);
+  3. ``last``:  the last number anywhere in the response.
+Numbers are normalised before comparison: thousands separators, ``$`` and a trailing
+``.`` are dropped; ``-1,234.50`` becomes Decimal("-1234.5"). A response with no number
+is wrong. Correct means numerically equal to the gold answer.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import re
+import urllib.error
+import urllib.request
+from decimal import Decimal, InvalidOperation
+from pathlib import Path
+
+from inferencekit.canonical import sha256_hex
+from inferencekit.registry import OpError
+
+DATASET = "openai/gsm8k (main, test); github.com/openai/grade-school-math; MIT"
+# Pinned source: the Hugging Face dataset repo at one commit, and the hashes of what it holds.
+HF_REPO = "openai/gsm8k"
+HF_REVISION = "740312add88f781978c0658806c59bc2815b9866"
+HF_FILE = "main/test-00000-of-00001.parquet"
+PARQUET_SHA256 = "ee7b8da9e381df27b9e3f7758a159ab2bdaa4dbaa910546cbbc47e0cb44e4f59"  # = its LFS oid
+JSONL_SHA256 = "3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14"
+N_TEST = 1319
+INSTRUCTION = ("Solve the problem step by step. At the end, write the final answer as a single number "
+               "on its own line in the form \"#### <number>\".")
+
+_NUM = re.compile(r"-?\$?\d[\d,]*(?:\.\d+)?|-?\.\d+")
+_BOXED = re.compile(r"\\boxed\{([^{}]*)\}")
+
+
+def cache_dir() -> Path:
+    base = os.environ.get("INFERENCEKIT_CACHE_DIR") or os.path.join(
+        os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache"), "inferencekit")
+    return Path(base) / "gsm8k"
+
+
+def source_url() -> str:
+    endpoint = (os.environ.get("HF_ENDPOINT") or "https://huggingface.co").rstrip("/")
+    return f"{endpoint}/datasets/{HF_REPO}/resolve/{HF_REVISION}/{HF_FILE}"
+
+
+def _parquet_path() -> Path:
+    return cache_dir() / Path(HF_FILE).name
+
+
+def manual_download_command() -> str:
+    p = _parquet_path()
+    return f"mkdir -p '{p.parent}' && curl -fL -o '{p}' '{source_url()}'"
+
+
+def parquet_to_jsonl(parquet: str | Path | bytes) -> bytes:
+    """The GSM8K Parquet table as JSONL bytes: {"question", "answer"} per line, json.dumps defaults."""
+    from inferencekit.quality.parquet import read_table
+
+    table = read_table(parquet)
+    if "question" not in table or "answer" not in table:
+        raise OpError(f"GSM8K Parquet file without question/answer columns (has {sorted(table)})")
+    rows = []
+    for q, a in zip(table["question"], table["answer"]):
+        if q is None or a is None:
+            raise OpError("GSM8K Parquet file with an empty question or answer")
+        rows.append(json.dumps({"question": q.decode("utf-8"), "answer": a.decode("utf-8")}) + "\n")
+    return "".join(rows).encode("utf-8")
+
+
+def _download(dest: Path, timeout: float = 120.0) -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_name(dest.name + f".part-{os.getpid()}")
+    req = urllib.request.Request(source_url(), headers={"User-Agent": "inferencekit (quality gsm8k)"})
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp, open(tmp, "wb") as f:  # noqa: S310
+            while chunk := resp.read(1 << 16):
+                f.write(chunk)
+        os.replace(tmp, dest)
+    finally:
+        tmp.unlink(missing_ok=True)
+
+
+def _sha(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def data_path(fetch: bool = True) -> Path:
+    """The cached GSM8K test split as JSONL (sha256-checked), downloading it once if needed.
+
+    Order: the cached JSONL; else the cached Parquet file (downloaded here, or placed by hand
+    with ``manual_download_command()``); else a download from Hugging Face at HF_REVISION.
+    ``fetch=False`` or ``HF_HUB_OFFLINE=1`` never touches the network."""
+    jsonl = cache_dir() / f"gsm8k-test-{HF_REVISION[:12]}.jsonl"
+    if jsonl.is_file() and _sha(jsonl.read_bytes()) == JSONL_SHA256:
+        return jsonl
+    pq = _parquet_path()
+    if not (pq.is_file() and _sha(pq.read_bytes()) == PARQUET_SHA256):
+        offline = os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in ("1", "true", "yes", "on")
+        bad_cache = pq.is_file()
+        reason = ("HF_HUB_OFFLINE is set" if offline else "fetching is off" if not fetch else "")
+        if bad_cache:
+            reason = f"the cached {pq.name} has the wrong sha256" + (f" and {reason}" if reason else "")
+        if fetch and not offline:  # (re)download; a cached file with the wrong hash is replaced
+            reason = ""
+            try:
+                _download(pq)
+            except (urllib.error.URLError, OSError, TimeoutError) as e:
+                reason = f"download failed: {getattr(e, 'reason', None) or e}"
+        if not reason and _sha(pq.read_bytes()) != PARQUET_SHA256:
+            reason = f"downloaded file sha256 {_sha(pq.read_bytes())} is not the pinned {PARQUET_SHA256}"
+            pq.unlink(missing_ok=True)
+        if reason:
+            raise OpError(
+                f"GSM8K test split is not available ({reason}). It is downloaded once from the Hugging Face "
+                f"dataset {HF_REPO} (config main, split test) at revision {HF_REVISION} and cached in "
+                f"{cache_dir()}. Fetch it by hand (on another machine if this one is offline, then copy it to "
+                f"{_parquet_path()} here):\n  {manual_download_command()}\nexpected sha256 {PARQUET_SHA256}. "
+                "Or pass --data with that Parquet file or any GSM8K-format JSONL.")
+    data = parquet_to_jsonl(pq)
+    if _sha(data) != JSONL_SHA256:
+        raise OpError(f"GSM8K: the JSONL converted from {pq} has sha256 {_sha(data)}, not the pinned "
+                      f"{JSONL_SHA256} (Parquet reader problem; please report it)")
+    tmp = jsonl.with_name(jsonl.name + f".part-{os.getpid()}")
+    tmp.write_bytes(data)
+    os.replace(tmp, jsonl)
+    return jsonl
+
+
+def load(path: str | Path | None = None, limit: int | None = None) -> tuple[list[dict], str]:
+    """[{idx, question, gold}] and the sha256 of the JSONL bytes.
+
+    ``path``: a GSM8K-format JSONL file, or a Parquet file with question/answer columns
+    (converted as for the pinned split). None: the pinned test split (downloaded once)."""
+    p = Path(path) if path else data_path()
+    if not p.is_file():
+        raise OpError(f"GSM8K data file not found: {p}")
+    try:
+        raw = parquet_to_jsonl(p) if p.suffix == ".parquet" else p.read_bytes()
+        text = raw.decode("utf-8")
+    except (OSError, ValueError) as e:  # ParquetError, not UTF-8
+        raise OpError(f"{p}: {e}") from None
+    items = []
+    for i, line in enumerate(text.splitlines()):
+        if not line.strip():
+            continue
+        try:
+            obj = json.loads(line)
+            gold = obj.get("gold")
+            if gold is None:
+                gold = obj["answer"].split("####")[-1].strip()
+            items.append({"idx": obj.get("idx", i), "question": obj["question"], "gold": str(gold)})
+        except (ValueError, KeyError, TypeError, AttributeError) as e:
+            raise OpError(f"{p}:{i + 1}: not a GSM8K row ({type(e).__name__}: {e}); need "
+                          "{\"question\", \"answer\" ending in '#### <n>'} or {\"question\", \"gold\"}") from None
+    if limit is not None and limit >= 0:
+        items = items[:limit]
+    return items, sha256_hex(raw)
+
+
+def messages(question: str) -> list[dict]:
+    return [{"role": "user", "content": f"{question}\n\n{INSTRUCTION}"}]
+
+
+def normalize(text: str | None) -> Decimal | None:
+    if text is None:
+        return None
+    s = text.replace(",", "").replace("$", "").replace(" ", "").strip().rstrip(".")
+    try:
+        return Decimal(s).normalize() if s else None
+    except InvalidOperation:
+        return None
+
+
+def _first_number(text: str) -> str | None:
+    m = _NUM.search(text)
+    return m.group(0) if m else None
+
+
+def _last_number(text: str) -> str | None:
+    found = _NUM.findall(text)
+    return found[-1] if found else None
+
+
+def extract(response: str | None) -> tuple[str | None, str]:
+    """(number string, rule) from a model response; (None, "none") when no number is found."""
+    if not response:
+        return None, "none"
+    if "####" in response:
+        n = _first_number(response.rsplit("####", 1)[1])
+        if n is not None:
+            return n, "hash"
+    boxed = _BOXED.findall(response)
+    if boxed:
+        n = _last_number(boxed[-1])
+        if n is not None:
+            return n, "boxed"
+    n = _last_number(response)
+    return (n, "last") if n is not None else (None, "none")
+
+
+def is_correct(extracted: str | None, gold: str) -> bool:
+    a, b = normalize(extracted), normalize(gold)
+    return a is not None and b is not None and a == b

inferencekit/quality/logprobs.py

@@ -0,0 +1,106 @@
+"""Turn server logprob payloads into token ids.
+
+Chat completions (OpenAI format) name tokens by *string*, which is ambiguous. Token ids
+are recovered, in this order, and the method is recorded:
+  ``token_id_strings``  vLLM with ``return_tokens_as_token_ids``: tokens read "token_id:1234";
+  ``id_field``          servers that add an integer ``id`` to each logprob entry (llama-server);
+  ``return_token_ids``  vLLM with ``return_token_ids``: ``choices[0].token_ids`` (chosen tokens only);
+  ``strings``           none of the above: keys stay token strings (scoring then re-tokenizes).
+
+Teacher forcing uses vLLM's ``prompt_logprobs`` on ``/v1/completions`` (a vLLM extension):
+``choices[0].prompt_logprobs`` is a list with one entry per prompt token; entry 0 is
+null (the first token has no context) and entry i maps token id -> {logprob, rank,
+decoded_token} for the top-k tokens at position i plus the token actually there.
+"""
+
+from __future__ import annotations
+
+import re
+
+_TOKEN_ID = re.compile(r"^token_id:(\d+)$")
+
+
+def _entry_id(entry: dict):
+    tok = entry.get("token")
+    if isinstance(tok, str):
+        m = _TOKEN_ID.match(tok)
+        if m:
+            return int(m.group(1)), "token_id_strings"
+    for key in ("id", "token_id"):
+        if isinstance(entry.get(key), int):
+            return entry[key], "id_field"
+    return None, None
+
+
+def parse_chat_logprobs(response: dict) -> dict:
+    """{token_ids, tokens, logprobs, top: [{key: logprob}], method, top_k_returned, text, finish_reason}.
+
+    ``top[i]`` always contains the chosen token as well as the returned alternatives.
+    Keys are ints when ids could be recovered, else ``"s:" + token string``.
+    """
+    choice = (response.get("choices") or [{}])[0]
+    content = ((choice.get("logprobs") or {}).get("content")) or []
+    side_ids = choice.get("token_ids")
+    ids, tokens, lps, tops, methods = [], [], [], [], set()
+    for i, entry in enumerate(content):
+        tid, how = _entry_id(entry)
+        if tid is None and isinstance(side_ids, list) and i < len(side_ids):
+            tid, how = side_ids[i], "return_token_ids"
+        methods.add(how or "strings")
+        key = tid if tid is not None else "s:" + str(entry.get("token"))
+        top = {}
+        for alt in entry.get("top_logprobs") or []:
+            aid, _ = _entry_id(alt)
+            akey = aid if aid is not None else "s:" + str(alt.get("token"))
+            if akey not in top:
+                top[akey] = alt.get("logprob")
+        top_k = len(top)
+        top.setdefault(key, entry.get("logprob"))
+        ids.append(tid)
+        tokens.append(entry.get("token"))
+        lps.append(entry.get("logprob"))
+        tops.append({"k": top_k, "top": top})
+    if isinstance(side_ids, list) and side_ids and len(side_ids) == len(ids):
+        mismatch = [i for i, (a, b) in enumerate(zip(ids, side_ids)) if a is not None and a != b]
+    else:
+        mismatch = []
+    method = "strings" if "strings" in methods or not methods else sorted(methods)[0]
+    if None in ids:
+        method = "strings"
+    msg = choice.get("message") or {}
+    return {"token_ids": ids if method != "strings" else None, "tokens": tokens, "logprobs": lps,
+            "top": [t["top"] for t in tops], "method": method if content else "none",
+            "top_k_returned": min((t["k"] for t in tops), default=0),
+            "id_mismatch_positions": mismatch, "text": msg.get("content") if isinstance(msg, dict) else None,
+            "finish_reason": choice.get("finish_reason"), "side_token_ids": side_ids}
+
+
+def parse_prompt_logprobs(response: dict) -> list[dict | None]:
+    """[{token_id: (logprob, rank)} | None] per prompt position, from a vLLM /v1/completions response."""
+    choice = (response.get("choices") or [{}])[0]
+    plp = choice.get("prompt_logprobs")
+    if plp is None:
+        plp = response.get("prompt_logprobs")
+    if not isinstance(plp, list):
+        return []
+    out: list[dict | None] = []
+    for entry in plp:
+        if not entry:
+            out.append(None)
+            continue
+        pos = {}
+        for k, v in entry.items():
+            if isinstance(v, dict):
+                pos[int(k)] = (v.get("logprob"), v.get("rank"))
+            else:  # a bare float
+                pos[int(k)] = (float(v), None)
+        out.append(pos)
+    return out
+
+
+def argmax_token(position: dict) -> int:
+    """The reference's most likely token at a position: rank 1 if ranks are given, else max logprob."""
+    ranked = [k for k, (_, rank) in position.items() if rank == 1]
+    if ranked:
+        return ranked[0]
+    return max(position, key=lambda k: position[k][0])