knn_classifier.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# bench/argv_policy/baselines/knn_classifier.py
 
"""
Embed the gate-labelled corpus once, then vote with the nearest neighbours.
 
Purpose
-------
This is the cheapest shape on this bench that learns the argv policy from
labels rather than from prose, and the only one that needs no training loop at
all: embed every corpus command once, embed the query, and read the verdict off
the nearest neighbours' labels. A decision costs one encoder forward pass and
one matrix-vector product, nothing per call, and the whole model is rebuilt by
re-embedding when the gate changes.
 
It exists to bound plan 290 from below. An encoder that is fine-tuned on this
corpus has to beat what the same encoder already achieves with no training at
all, or the training was not worth its cost; and both have to beat plan 288's
standard-library token floor, which decides in 0.07 ms.
 
The vote is a softmax over PER-CLASS SIMILARITY SUMS across the top-k cosine
neighbours, divided by a temperature. Two readings of that formula differ in
what happens to a class no neighbour voted for. This module's default masks it
out -- an empty sum is no forecast, not a zero one -- because the alternative
gives a class with no votes ``exp(0)`` of the mass, which at k=1 makes the
nearest-neighbour forecast 1/(1+2exp(-1/T)) rather than the exact 1.0 that a
1-NN forecast actually is. Both readings are implemented and both are swept by
:func:`fit`, so the choice is measured on validation Brier rather than asserted
(see ``mask_absent``).
 
Ties break toward ``deny``. A vote that cannot separate two classes grants the
least permission, not the most, which matters because the query that produces
an exactly tied vote is the degenerate one -- a command whose embedding is the
zero vector, so every similarity is 0.0 and every present class sums to 0.0.
 
Run the index build and the evaluation under the SCRATCH HARNESS venv (torch
and sentence-transformers live there, and this bench adds no dependency to any
``pyproject.toml``), from the oc_chat subproject root::
 
    python bench/argv_policy/baselines/knn_classifier.py --stage all
 
Run the tests under the PROJECT venv, where :class:`HashEmbedder` stands in for
the real encoder so the witnesses need no model, no GPU and no network.
 
Responsibilities
----------------
- :class:`HashEmbedder` / :class:`SentenceTransformerEmbedder` /
  :class:`LMStudioEmbedder` -- the three embedding routes.
- :func:`write_npy` / :func:`read_npy` -- a standard-library ``.npy`` codec, so
  an index written here is readable with or without numpy.
- :func:`build_index` / :func:`load_index` -- the index on disk.
- :func:`top_neighbours` / :func:`class_sums` / :func:`probabilities` /
  :func:`verdict_of` / :func:`classify` -- one decision.
- :func:`fit` -- k and the temperature from validation Brier alone.
- :func:`evaluate` / :func:`write_result` -- a case set to a plan-288 result
  file.
- :func:`main` -- the CLI: build, fit, evaluate.
 
Diagnostics
-----------
Domain: OC_CHAT.MAIN
Levels:
    L2 - index and evaluation lifecycle (build start and finish, files written)
    L3 - per-batch, per-fit-cell and per-case-set detail
    L4 - module load
 
Contracts
---------
- Deterministic given an embedder: the corpus, the pinned subset seed and the
  (k, temperature) pair fix every number this module reports.
- Read-only with respect to the corpus, the case sets and the scorer.
- Index arrays are megabytes and rebuildable, so :func:`build_index` writes a
  ``.gitignore`` of ``*`` into its own output directory. Nothing under it is
  ever committed, and no model weights are either.
- numpy is used when it is importable and a pure-Python path is used when it is
  not; the on-disk format and every number are identical either way. The
  fallback is for the witnesses, not for a 27000-row index.
- Console output is ASCII only; commands are printed through :func:`ascii`.
- Standard library only at import time. torch, sentence-transformers and numpy
  are imported lazily, inside the routes that need them.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
- numpy (lazy; fast matrix maths and .npy loading when importable, with a pure-Python fallback)
- sentence_transformers (lazy; the SentenceTransformer encoder behind the required embedding route)
- torch (lazy; pins the encoder's CPU thread count)
"""
from __future__ import annotations
 
import argparse
import ast
import json
import math
import os
import random
import struct
import sys
import time
import urllib.error
import urllib.request
import zlib
from array import array
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Sequence
 
from oc_diagnostics import _dbg
 
_dbg("OC_CHAT.MAIN", "module", "loaded", 4)
 
#: The three verdicts, which are also the three forecast classes.
LABELS: tuple[str, ...] = ("allow", "confirm", "deny")
 
#: The verdict an undecided vote resolves to. Fail closed: the least
#: permission, never the most.
TIE_BREAK = "deny"
 
#: Plan 289 P3's neighbour counts. Chosen on the validation split, never on a
#: case set. 51 is odd so a two-class vote cannot tie on count alone.
DEFAULT_K_GRID: tuple[int, ...] = (1, 5, 15, 51)
 
#: The temperature grid swept beside ``k``. P3 pins the ``k`` values and leaves
#: the temperature to the implementation, so it is pinned here: eight values
#: over more than two orders of magnitude, from near-one-hot to nearly uniform.
#: The top two exist because a grid whose minimum sits on its own boundary is a
#: sweep that was not allowed to finish -- the first run of this module put the
#: minimum at 1.0 for all four k, which is a truncated answer, not an answer.
DEFAULT_TEMPERATURE_GRID: tuple[float, ...] = (
    0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0,
)
 
#: The neighbour count and temperature a caller gets when it does not fit. Only
#: the witnesses use these; every reported number comes from :func:`fit`.
DEFAULT_K = 15
DEFAULT_TEMPERATURE = 0.1
 
#: Plan 289 P4's small-corpus condition: a pinned random subset of the training
#: split, so the report can say how much labelled material the shape needs.
SUBSET_ROWS = 1000
SUBSET_SEED = 289
 
#: The required embedder (P2), from the harness cache, offline.
JINA_MODEL_ID = "jinaai/jina-embeddings-v2-base-code"
 
#: jina-v2's remote code allocates its ALiBi bias for the CONFIGURED maximum
#: length regardless of the input, which is 3.2 GB at its default 8192 and
#: crashed the 2026-09-17 session. Commands are short; cap it.
MAX_SEQ_LENGTH = 512
 
#: Threads for the encoder forward pass. The September run at cpu_count minus
#: two oversubscribed the host, so this is pinned rather than derived.
THREADS = 8
 
#: The optional embedder (P2), only when the operator has LM Studio running.
LMSTUDIO_URL = "http://127.0.0.1:1234/v1/embeddings"
LMSTUDIO_MODEL = "text-embedding-qwen3-embedding-0.6b"
LMSTUDIO_TIMEOUT = 30.0
 
#: Rows per encoder call when embedding a corpus.
BUILD_BATCH = 128
 
#: Query rows per similarity matmul when fitting. 256 keeps the similarity
#: block near 27 MB at corpus scale instead of the 324 MB a single call needs.
FIT_BATCH = 256
 
#: ``.npy`` version 1.0: magic, version, then a uint16 header length. The total
#: of those and the header must be a multiple of this, and the header ends in a
#: newline.
NPY_MAGIC = b"\x93NUMPY"
NPY_ALIGNMENT = 64
NPY_PREAMBLE = 10
 
#: Character n-gram widths :class:`HashEmbedder` projects. The same range the
#: token floor uses, so the stand-in is a recognisable shape rather than noise.
HASH_NGRAM_SIZES: tuple[int, ...] = (1, 2, 3, 4)
 
 
def _numpy() -> Any:
    """Return numpy when it is importable, else ``None``.
 
    The witnesses run in the project venv, which has neither numpy nor torch;
    the index builds and the bench evaluations run in the scratch harness venv,
    which has both. One module serves both, so the import is a question rather
    than a requirement.
    """
 
    try:
        import numpy
    except ImportError:          # pragma: no cover - exercised by the project venv
        return None
    return numpy
 
 
#: numpy when this interpreter has it, else ``None``. See :func:`_numpy`.
NUMPY: Any = _numpy()
 
 
def read_jsonl(path: Path) -> list[dict[str, str]]:
    """Read a JSON-lines corpus split into a list of rows."""
 
    text = Path(path).read_text(encoding="utf-8")
    return [json.loads(line) for line in text.splitlines() if line.strip()]
 
 
def read_cases(path: Path) -> list[dict[str, str]]:
    """Read one of the bench's hand-authored case sets."""
 
    return json.loads(Path(path).read_text(encoding="utf-8"))
 
 
class HashEmbedder:
    """A deterministic stand-in encoder: signed character n-grams, unit-norm.
 
    Not a model and not a baseline -- it is what lets every witness in
    ``tests/test_289_knn.py`` run in the project venv with no torch, no
    download and no network, and what makes ``--embedder hash`` a smoke run of
    the whole pipeline in under a second. CRC-32 rather than :func:`hash`,
    which is salted per process: an index whose columns move between runs
    cannot be built once and queried later.
    """
 
    def __init__(self, dim: int = 64) -> None:
        self.dim = dim
        self.name = "hash"
 
    def encode(self, texts: Sequence[str]) -> list[list[float]]:
        """Project each text onto :attr:`dim` unit-norm dimensions."""
 
        return [self._one(text) for text in texts]
 
    def _one(self, text: str) -> list[float]:
        """One text to one unit-norm vector."""
 
        vector = [0.0] * self.dim
        padded = f"\x02{text}\x03"
        for size in HASH_NGRAM_SIZES:
            for start in range(len(padded) - size + 1):
                key = f"{size}\x1f{padded[start:start + size]}".encode("utf-8")
                digest = zlib.crc32(key)
                sign = 1.0 if digest & 1 else -1.0
                vector[digest % self.dim] += sign
        norm = math.sqrt(sum(value * value for value in vector))
        if norm == 0.0:
            return vector
        return [value / norm for value in vector]
 
 
class SentenceTransformerEmbedder:
    """The required route (P2): a sentence-transformers encoder on CPU.
 
    Imported lazily so this module stays importable in a venv with no torch.
    ``HF_HUB_OFFLINE`` is set here rather than left to the caller's shell,
    because a cache miss that silently reaches the network is the one failure
    mode a reproducible bench cannot tolerate.
    """
 
    def __init__(
        self,
        model_id: str = JINA_MODEL_ID,
        max_seq_length: int = MAX_SEQ_LENGTH,
        threads: int = THREADS,
        name: str | None = None,
    ) -> None:
        func = "SentenceTransformerEmbedder.__init__"
        os.environ.setdefault("HF_HUB_OFFLINE", "1")
        import torch
        from sentence_transformers import SentenceTransformer
 
        torch.set_num_threads(threads)
        started = time.perf_counter()
        self.model = SentenceTransformer(
            model_id,
            trust_remote_code=True,
            config_kwargs={"max_position_embeddings": max_seq_length},
        )
        self.model.max_seq_length = max_seq_length
        self.load_seconds = time.perf_counter() - started
        self.model_id = model_id
        self.threads = threads
        self.max_seq_length = max_seq_length
        self.name = name or model_id.rsplit("/", 1)[-1]
        self.dim = int(self.model.get_sentence_embedding_dimension())
        _dbg("OC_CHAT.MAIN", func,
             f"{self.name} dim {self.dim} in {self.load_seconds:.1f}s", 2)
 
    def encode(self, texts: Sequence[str]) -> Any:
        """Embed *texts* as a numpy float array, one row each."""
 
        return self.model.encode(
            list(texts), batch_size=min(BUILD_BATCH, max(1, len(texts))),
            show_progress_bar=False,
        )
 
 
class LMStudioEmbedder:
    """The optional route (P2): an OpenAI-shaped embeddings endpoint.
 
    Needs no torch, so it runs in either venv. It is optional because it needs
    the operator's LM Studio to be running: :meth:`probe` answers that in one
    short call so a worker can skip the leg instead of waiting on it.
    """
 
    def __init__(
        self,
        url: str = LMSTUDIO_URL,
        model_id: str = LMSTUDIO_MODEL,
        timeout: float = LMSTUDIO_TIMEOUT,
        name: str | None = None,
    ) -> None:
        func = "LMStudioEmbedder.__init__"
        self.url = url
        self.model_id = model_id
        self.timeout = timeout
        self.name = name or f"lmstudio_{model_id.rsplit('/', 1)[-1]}"
        self.dim = len(self.encode(["pip list"])[0])
        _dbg("OC_CHAT.MAIN", func, f"{self.name} dim {self.dim}", 2)
 
    @staticmethod
    def probe(url: str = LMSTUDIO_URL, timeout: float = 3.0) -> bool:
        """Is the endpoint answering? One call, short timeout, never waits."""
 
        models_url = url.rsplit("/", 1)[0] + "/models"
        try:
            with urllib.request.urlopen(models_url, timeout=timeout) as response:
                json.loads(response.read().decode("utf-8"))
        except (urllib.error.URLError, OSError, ValueError):
            return False
        return True
 
    def encode(self, texts: Sequence[str]) -> list[list[float]]:
        """Embed *texts* through one POST, preserving the input order."""
 
        body = json.dumps({"model": self.model_id, "input": list(texts)})
        request = urllib.request.Request(
            self.url, data=body.encode("utf-8"),
            headers={"Content-Type": "application/json"}, method="POST",
        )
        with urllib.request.urlopen(request, timeout=self.timeout) as response:
            payload = json.loads(response.read().decode("utf-8"))
        rows = sorted(payload["data"], key=lambda item: item.get("index", 0))
        return [[float(value) for value in row["embedding"]] for row in rows]
 
 
def _npy_header(rows: int, dim: int) -> bytes:
    """The padded version-1.0 header describing a C-order float32 matrix."""
 
    body = f"{{'descr': '<f4', 'fortran_order': False, 'shape': ({rows}, {dim}), }}"
    padding = -(NPY_PREAMBLE + len(body) + 1) % NPY_ALIGNMENT
    return body.encode("latin-1") + b" " * padding + b"\n"
 
 
def _flat_float32(matrix: Any, rows: int, dim: int) -> bytes:
    """Flatten *matrix* to little-endian float32 bytes.
 
    Two paths for one format: numpy when the caller handed over an array, and
    :mod:`array` when it handed over lists. The bytes are identical.
    """
 
    if NUMPY is not None and hasattr(matrix, "astype"):
        return NUMPY.ascontiguousarray(matrix, dtype="<f4").tobytes()
    flat = array("f")
    for row in matrix:
        if len(row) != dim:
            raise ValueError(f"row of width {len(row)} in a {dim}-wide matrix")
        flat.extend(float(value) for value in row)
    if len(flat) != rows * dim:
        raise ValueError(f"{len(flat)} values for a {rows}x{dim} matrix")
    if sys.byteorder != "little":  # pragma: no cover - x86/ARM are little-endian
        flat.byteswap()
    return flat.tobytes()
 
 
def write_npy(path: Path, matrix: Any) -> Path:
    """Write *matrix* as a C-order float32 ``.npy`` file.
 
    Hand-rolled because the project venv has no numpy and the witnesses must
    still be able to round-trip an index; the harness venv reads the very same
    bytes with :func:`numpy.load`.
    """
 
    rows = len(matrix)
    dim = len(matrix[0]) if rows else 0
    data = _flat_float32(matrix, rows, dim)
    header = _npy_header(rows, dim)
    with Path(path).open("wb") as handle:
        handle.write(NPY_MAGIC)
        handle.write(bytes([1, 0]))
        handle.write(struct.pack("<H", len(header)))
        handle.write(header)
        handle.write(data)
    return Path(path)
 
 
def read_npy(path: Path) -> tuple[int, int, array]:
    """Read a float32 ``.npy`` file into ``(rows, dim, flat values)``."""
 
    with Path(path).open("rb") as handle:
        if handle.read(6) != NPY_MAGIC:
            raise ValueError(f"{path}: not a .npy file")
        major = handle.read(2)[0]
        width = 2 if major == 1 else 4
        length = int.from_bytes(handle.read(width), "little")
        header = ast.literal_eval(handle.read(length).decode("latin-1").strip())
        if header["descr"] != "<f4" or header["fortran_order"]:
            raise ValueError(f"{path}: expected C-order '<f4', got {header}")
        rows, dim = (int(value) for value in header["shape"])
        flat = array("f")
        flat.fromfile(handle, rows * dim)
    if sys.byteorder != "little":  # pragma: no cover - x86/ARM are little-endian
        flat.byteswap()
    return rows, dim, flat
 
 
@dataclass(frozen=True)
class Index:
    """An embedded corpus on disk, loaded and ready to answer queries."""
 
    path: Path
    embedder_name: str
    rows: int
    dim: int
    commands: list[str]
    labels: list[str]
    vectors: Any
    meta: dict[str, Any]
    embedder: Any = None
 
    @property
    def bytes_on_disk(self) -> int:
        """Total size of the index directory, the number P5 asks for."""
 
        return sum(item.stat().st_size for item in self.path.iterdir()
                   if item.is_file())
 
 
def _normalise_rows(matrix: Any) -> Any:
    """Scale every row to unit length, leaving an all-zero row alone.
 
    Cosine similarity is a dot product only when both sides are unit-norm, so
    this is done once at build time and never per query. A zero row cannot be
    normalised and is kept as it is; it produces a similarity of 0.0 against
    everything, which is the degenerate case the tie rule handles.
    """
 
    if NUMPY is not None and hasattr(matrix, "astype"):
        array_2d = NUMPY.asarray(matrix, dtype="float32")
        norms = NUMPY.linalg.norm(array_2d, axis=1, keepdims=True)
        return array_2d / NUMPY.where(norms == 0.0, 1.0, norms)
    scaled: list[list[float]] = []
    for row in matrix:
        norm = math.sqrt(sum(float(value) * float(value) for value in row))
        divisor = norm if norm else 1.0
        scaled.append([float(value) / divisor for value in row])
    return scaled
 
 
def build_index(
    corpus_path: Path,
    embedder: Any,
    index_dir: Path,
    subset: int | None = None,
    seed: int = SUBSET_SEED,
) -> Path:
    """Embed a corpus split once and write the index to *index_dir*.
 
    Three files: ``vectors.npy`` (row-normalised float32), ``index.json`` (the
    commands, the labels and how the index was built) and a ``.gitignore`` of
    ``*``, because the array is megabytes of rebuildable floats and must not be
    committable by accident.
    """
 
    func = "build_index"
    rows = read_jsonl(corpus_path)
    if subset is not None and subset < len(rows):
        rows = random.Random(seed).sample(rows, subset)
    commands = [row["command"] for row in rows]
    labels = [row["truth"] for row in rows]
    _dbg("OC_CHAT.MAIN", func,
         f"embedding {len(commands)} commands with {embedder.name}", 2)
 
    started = time.perf_counter()
    batches: list[Any] = []
    for start in range(0, len(commands), BUILD_BATCH):
        chunk = commands[start:start + BUILD_BATCH]
        batches.append(embedder.encode(chunk))
        _dbg("OC_CHAT.MAIN", func, f"embedded {start + len(chunk)} of {len(commands)}", 3)
    if NUMPY is not None and batches and hasattr(batches[0], "astype"):
        matrix = NUMPY.concatenate(batches, axis=0)
    else:
        matrix = [row for batch in batches for row in batch]
    matrix = _normalise_rows(matrix)
    build_seconds = time.perf_counter() - started
 
    target = Path(index_dir)
    target.mkdir(parents=True, exist_ok=True)
    (target / ".gitignore").write_text(
        "# Plan 289: an embedded corpus is megabytes of float32 that any\n"
        "# checkout can rebuild from the committed corpus and the cached\n"
        "# encoder, so nothing in this directory is ever committed.\n"
        "*\n",
        encoding="utf-8",
    )
    write_npy(target / "vectors.npy", matrix)
    (target / "index.json").write_text(
        json.dumps({
            "embedder": embedder.name,
            "dim": int(embedder.dim),
            "rows": len(commands),
            "corpus_file": Path(corpus_path).name,
            "subset_rows": subset,
            "subset_seed": seed if subset is not None else None,
            "build_seconds": build_seconds,
            "commands": commands,
            "labels": labels,
        }, indent=1, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    _dbg("OC_CHAT.MAIN", func,
         f"wrote {target} in {build_seconds:.1f}s", 2)
    return target
 
 
def load_index(index_dir: Path, embedder: Any = None) -> Index:
    """Load an index from disk, optionally with a live embedder for queries."""
 
    func = "load_index"
    target = Path(index_dir)
    meta = json.loads((target / "index.json").read_text(encoding="utf-8"))
    if embedder is not None and embedder.name != meta["embedder"]:
        raise ValueError(
            f"{target}: index built with {meta['embedder']!r}, "
            f"queried with {embedder.name!r}"
        )
    if NUMPY is not None:
        vectors = NUMPY.load(target / "vectors.npy")
        rows, dim = int(vectors.shape[0]), int(vectors.shape[1])
    else:
        rows, dim, vectors = read_npy(target / "vectors.npy")
    if rows != meta["rows"] or dim != meta["dim"]:
        raise ValueError(
            f"{target}: array is {rows}x{dim}, metadata says "
            f"{meta['rows']}x{meta['dim']}"
        )
    _dbg("OC_CHAT.MAIN", func, f"{target}: {rows}x{dim} {meta['embedder']}", 2)
    return Index(
        path=target, embedder_name=str(meta["embedder"]), rows=rows, dim=dim,
        commands=list(meta["commands"]), labels=list(meta["labels"]),
        vectors=vectors, meta=meta, embedder=embedder,
    )
 
 
def index_row(index: Index, position: int) -> list[float]:
    """One row of the index as a plain list, whichever backend holds it."""
 
    if NUMPY is not None and hasattr(index.vectors, "shape"):
        return [float(value) for value in index.vectors[position]]
    start = position * index.dim
    return list(index.vectors[start:start + index.dim])
 
 
def _similarities(index: Index, query: Sequence[float]) -> Any:
    """Cosine similarity of *query* against every indexed row.
 
    A dot product, because both sides are unit-norm: the index rows were
    normalised at build time and the query is normalised by :func:`classify`.
    """
 
    if NUMPY is not None and hasattr(index.vectors, "shape"):
        return index.vectors @ NUMPY.asarray(query, dtype="float32")
    out: list[float] = []
    vectors = index.vectors
    dim = index.dim
    for position in range(index.rows):
        start = position * dim
        out.append(sum(vectors[start + offset] * query[offset]
                       for offset in range(dim)))
    return out
 
 
def top_neighbours(
    index: Index, query: Sequence[float], k: int
) -> list[tuple[float, str]]:
    """The *k* most similar indexed rows as ``(similarity, label)`` pairs."""
 
    similarities = _similarities(index, query)
    count = min(k, index.rows)
    if NUMPY is not None and hasattr(similarities, "shape"):
        if count < index.rows:
            candidates = NUMPY.argpartition(-similarities, count - 1)[:count]
        else:
            candidates = NUMPY.arange(index.rows)
        order = candidates[NUMPY.argsort(-similarities[candidates], kind="stable")]
        return [(float(similarities[position]), index.labels[int(position)])
                for position in order]
    ranked = sorted(range(index.rows), key=lambda position: -similarities[position])
    return [(float(similarities[position]), index.labels[position])
            for position in ranked[:count]]
 
 
def class_sums(neighbours: Iterable[tuple[float, str]]) -> dict[str, float]:
    """Sum the similarities of the neighbours belonging to each class.
 
    A class absent from the result is a class no neighbour voted for, which is
    not the same thing as a class whose votes summed to zero. Membership, not
    the sign of the sum, is what :func:`probabilities` masks on.
    """
 
    sums: dict[str, float] = {}
    for similarity, label in neighbours:
        sums[label] = sums.get(label, 0.0) + similarity
    return sums
 
 
def probabilities(
    sums: dict[str, float], temperature: float, mask_absent: bool = True
) -> dict[str, float]:
    """Softmax the per-class similarity sums over a temperature.
 
    With ``mask_absent`` (the default, and the reading in the module docstring)
    a class no neighbour voted for receives no mass, so a 1-NN forecast is
    exactly one-hot. Without it, every class enters the softmax with its sum --
    zero when absent -- which is the other reading of P1's formula and the arm
    :func:`fit` sweeps beside the first.
    """
 
    if temperature <= 0.0:
        raise ValueError(f"temperature must be positive, got {temperature}")
    present = [label for label in LABELS
               if not mask_absent or label in sums]
    if not present:                     # pragma: no cover - k >= 1 always votes
        return {label: 0.0 for label in LABELS}
    scaled = {label: sums.get(label, 0.0) / temperature for label in present}
    top = max(scaled.values())
    exponentials = {label: math.exp(value - top) for label, value in scaled.items()}
    total = sum(exponentials.values())
    return {label: exponentials.get(label, 0.0) / total for label in LABELS}
 
 
def verdict_of(probs: dict[str, float]) -> tuple[str, bool]:
    """The verdict a forecast implies, and whether it came from a tie.
 
    Ties break toward :data:`TIE_BREAK`. The caller gets the flag as well as
    the verdict because plan 289 P3 requires the report to say how often the
    rule fired: a tie rule nobody counts is a rule nobody can audit.
    """
 
    best = max(probs.get(label, 0.0) for label in LABELS)
    tied = [label for label in LABELS if probs.get(label, 0.0) == best]
    if len(tied) > 1:
        return TIE_BREAK, True
    return tied[0], False
 
 
def embed_query(index: Index, command: str) -> list[float]:
    """Embed one command with the index's own encoder, normalised."""
 
    if index.embedder is None:
        raise ValueError(f"{index.path}: loaded without an embedder, cannot query")
    vector = index.embedder.encode([command])[0]
    values = [float(value) for value in vector]
    norm = math.sqrt(sum(value * value for value in values))
    if norm == 0.0:
        return values
    return [value / norm for value in values]
 
 
def classify(
    command: str,
    index: Index,
    k: int = DEFAULT_K,
    temperature: float = DEFAULT_TEMPERATURE,
    mask_absent: bool = True,
) -> dict[str, float]:
    """The forecast for one command: embed, search, vote.
 
    This is the whole decision, and it is what the latency column measures:
    one encoder forward pass at batch size one plus one matrix-vector product,
    with no file read and no model load inside it.
    """
 
    neighbours = top_neighbours(index, embed_query(index, command), k)
    return probabilities(class_sums(neighbours), temperature, mask_absent)
 
 
@dataclass(frozen=True)
class FitCell:
    """One ``(k, temperature, masking)`` cell of the validation sweep."""
 
    k: int
    temperature: float
    mask_absent: bool
    brier: float
    correct: int
    total: int
    ties: int
 
    @property
    def agreement(self) -> str:
        """``correct/total`` with its percentage."""
 
        share = 100.0 * self.correct / self.total if self.total else 0.0
        return f"{self.correct}/{self.total} ({share:.1f}%)"
 
 
@dataclass(frozen=True)
class Fit:
    """The chosen hyperparameters and the whole grid they were chosen from."""
 
    k: int
    temperature: float
    mask_absent: bool
    brier: float
    correct: int
    total: int
    rows: int
    source: str
    grid: list[FitCell]
    seconds: float
 
 
def _neighbour_batches(
    index: Index, vectors: Any, k: int
) -> list[list[tuple[float, str]]]:
    """Top-*k* neighbours for a batch of query vectors.
 
    With numpy this is one matmul per batch rather than one per query, which is
    the difference between a fit that takes seconds and one that takes an hour.
    """
 
    if NUMPY is not None and hasattr(index.vectors, "shape"):
        queries = NUMPY.asarray(vectors, dtype="float32")
        similarities = queries @ index.vectors.T
        count = min(k, index.rows)
        out: list[list[tuple[float, str]]] = []
        for position in range(similarities.shape[0]):
            row = similarities[position]
            if count < index.rows:
                candidates = NUMPY.argpartition(-row, count - 1)[:count]
            else:
                candidates = NUMPY.arange(index.rows)
            order = candidates[NUMPY.argsort(-row[candidates], kind="stable")]
            out.append([(float(row[item]), index.labels[int(item)]) for item in order])
        return out
    return [top_neighbours(index, vector, k) for vector in vectors]
 
 
def _normalise_one(vector: Sequence[float]) -> list[float]:
    """One vector scaled to unit length, or left alone when it is all zero."""
 
    values = [float(value) for value in vector]
    norm = math.sqrt(sum(value * value for value in values))
    if norm == 0.0:
        return values
    return [value / norm for value in values]
 
 
def fit(
    validation_path: Path,
    index: Index,
    k_grid: Sequence[int] = DEFAULT_K_GRID,
    temperature_grid: Sequence[float] = DEFAULT_TEMPERATURE_GRID,
    maskings: Sequence[bool] = (True, False),
) -> Fit:
    """Choose ``k`` and the temperature on the validation split, by Brier.
 
    P3: never on a case set. The neighbours are searched ONCE at the largest
    ``k`` in the grid and every cell reads a prefix of that list, so the sweep
    costs one pass over the validation split rather than one per cell.
    """
 
    func = "fit"
    rows = read_jsonl(validation_path)
    truths = [row["truth"] for row in rows]
    widest = max(k_grid)
    started = time.perf_counter()
    neighbours: list[list[tuple[float, str]]] = []
    for start in range(0, len(rows), FIT_BATCH):
        chunk = [row["command"] for row in rows[start:start + FIT_BATCH]]
        vectors = [_normalise_one(vector) for vector in index.embedder.encode(chunk)]
        neighbours.extend(_neighbour_batches(index, vectors, widest))
        _dbg("OC_CHAT.MAIN", func, f"searched {len(neighbours)} of {len(rows)}", 3)
 
    grid: list[FitCell] = []
    for k in k_grid:
        sums = [class_sums(row[:k]) for row in neighbours]
        for mask_absent in maskings:
            for temperature in temperature_grid:
                total_brier = 0.0
                correct = 0
                ties = 0
                for position, truth in enumerate(truths):
                    probs = probabilities(sums[position], temperature, mask_absent)
                    total_brier += sum(
                        (probs.get(label, 0.0) - (1.0 if label == truth else 0.0)) ** 2
                        for label in LABELS
                    )
                    verdict, tied = verdict_of(probs)
                    correct += int(verdict == truth)
                    ties += int(tied)
                cell = FitCell(
                    k=k, temperature=temperature, mask_absent=mask_absent,
                    brier=total_brier / len(rows), correct=correct,
                    total=len(rows), ties=ties,
                )
                grid.append(cell)
                _dbg("OC_CHAT.MAIN", func,
                     f"k={k} T={temperature} mask={mask_absent}: "
                     f"Brier {cell.brier:.4f} {cell.agreement}", 3)
 
    best = min(grid, key=lambda cell: (cell.brier, cell.k, cell.temperature))
    seconds = time.perf_counter() - started
    _dbg("OC_CHAT.MAIN", func,
         f"chose k={best.k} T={best.temperature} mask={best.mask_absent} "
         f"Brier {best.brier:.4f} in {seconds:.1f}s", 2)
    return Fit(
        k=best.k, temperature=best.temperature, mask_absent=best.mask_absent,
        brier=best.brier, correct=best.correct, total=best.total, rows=len(rows),
        source=Path(validation_path).name, grid=grid, seconds=seconds,
    )
 
 
def evaluate(
    cases_path: Path,
    index: Index,
    k: int = DEFAULT_K,
    temperature: float = DEFAULT_TEMPERATURE,
    mask_absent: bool = True,
    model_name: str = "knn",
    fitted: Fit | None = None,
) -> dict[str, Any]:
    """Decide every case in *cases_path* and build a plan-288 result payload.
 
    P1 asks for a result FILE; producing the payload and writing it are split
    so the witnesses can validate the schema without touching disk and so one
    evaluation can be scored twice. :func:`write_result` is the other half.
 
    ``condition`` is ``"corpus"`` for every file this module writes (P2): this
    classifier is never given the policy text, so the policy/nopolicy axis the
    prompted legs use does not exist for it. The case set therefore lives in
    *model_name* and in the provenance block, not in the condition column.
    """
 
    func = "evaluate"
    cases = read_cases(cases_path)
    rows: list[dict[str, Any]] = []
    ties = 0
    for case in cases:
        started = time.perf_counter()
        probs = classify(case["command"], index, k, temperature, mask_absent)
        verdict, tied = verdict_of(probs)
        elapsed_ms = (time.perf_counter() - started) * 1000.0
        ties += int(tied)
        rows.append({
            "command": case["command"],
            "truth": case["truth"],
            "slice": case["slice"],
            "pred": verdict,
            "probs": probs,
            "ms": elapsed_ms,
            "reason": case["reason"],
            "tie": tied,
        })
    hits = sum(row["pred"] == row["truth"] for row in rows)
    _dbg("OC_CHAT.MAIN", func,
         f"{model_name}: {hits}/{len(rows)} agreement, {ties} ties", 3)
    provenance: dict[str, Any] = {
        "route": (f"CPU, {index.embedder_name} embeddings, "
                  f"cosine kNN over the gate-labelled corpus"),
        "embedder": index.embedder_name,
        "embedder_dim": index.dim,
        "case_file": Path(cases_path).name,
        "corpus_file": index.meta.get("corpus_file"),
        "index_rows": index.rows,
        "index_subset_rows": index.meta.get("subset_rows"),
        "index_subset_seed": index.meta.get("subset_seed"),
        "index_build_seconds": index.meta.get("build_seconds"),
        "index_bytes": index.bytes_on_disk,
        "k": k,
        "temperature": temperature,
        "mask_absent": mask_absent,
        "tie_break": TIE_BREAK,
        "ties_broken": ties,
        "threads": THREADS,
        "max_seq_length": MAX_SEQ_LENGTH,
    }
    if fitted is not None:
        provenance["fit"] = {
            "source": fitted.source,
            "rows": fitted.rows,
            "k_grid": list(DEFAULT_K_GRID),
            "temperature_grid": list(DEFAULT_TEMPERATURE_GRID),
            "validation_brier": fitted.brier,
            "validation_agreement": f"{fitted.correct}/{fitted.total}",
            "seconds": fitted.seconds,
        }
    return {
        "model": model_name,
        "condition": "corpus",
        "results": rows,
        "provenance": provenance,
    }
 
 
def write_result(payload: dict[str, Any], out_path: Path) -> Path:
    """Write one result payload where :mod:`score` can find it."""
 
    func = "write_result"
    target = Path(out_path)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(
        json.dumps(payload, indent=1, sort_keys=True) + "\n", encoding="utf-8"
    )
    _dbg("OC_CHAT.MAIN", func, f"wrote {target}", 2)
    return target
 
 
def render_fit_table(fitted: Fit, mask_absent: bool | None = None) -> str:
    """The validation-Brier table W2 requires in the report, as markdown.
 
    Pass *mask_absent* to render one arm of the sweep; omit it for the whole
    grid. :func:`render_masking_table` is the compact comparison of the two.
    """
 
    lines = [
        "| k | temperature | masking | validation Brier | validation agreement | ties |",
        "|---:|---:|---|---:|---:|---:|",
    ]
    cells = [cell for cell in fitted.grid
             if mask_absent is None or cell.mask_absent == mask_absent]
    for cell in sorted(cells,
                       key=lambda item: (item.k, not item.mask_absent, item.temperature)):
        masking = "masked" if cell.mask_absent else "unmasked"
        lines.append(
            f"| {cell.k} | {cell.temperature} | {masking} | {cell.brier:.4f} "
            f"| {cell.agreement} | {cell.ties} |"
        )
    return "\n".join(lines) + "\n"
 
 
def render_masking_table(fitted: Fit) -> str:
    """Best cell per ``k`` under each reading of P1's softmax, as markdown.
 
    This is the table that settles the masking question empirically rather
    than by argument: the same fit rule applied to both arms, side by side.
    """
 
    lines = [
        "| k | masked: best T | masked Brier | unmasked: best T | unmasked Brier |",
        "|---:|---:|---:|---:|---:|",
    ]
    for k in sorted({cell.k for cell in fitted.grid}):
        row = [str(k)]
        for mask_absent in (True, False):
            arm = [cell for cell in fitted.grid
                   if cell.k == k and cell.mask_absent == mask_absent]
            if not arm:
                row.extend(["-", "-"])
                continue
            best = min(arm, key=lambda cell: (cell.brier, cell.temperature))
            row.extend([str(best.temperature), f"{best.brier:.4f}"])
        lines.append("| " + " | ".join(row) + " |")
    return "\n".join(lines) + "\n"
 
 
def _make_embedder(kind: str) -> Any:
    """Build the embedder the CLI was asked for."""
 
    if kind == "hash":
        return HashEmbedder()
    if kind == "jina":
        return SentenceTransformerEmbedder()
    if kind == "lmstudio":
        return LMStudioEmbedder()
    raise ValueError(f"unknown embedder {kind!r}")
 
 
def _build_parser() -> argparse.ArgumentParser:
    """Construct the CLI argument parser."""
 
    parser = argparse.ArgumentParser(
        prog="python bench/argv_policy/baselines/knn_classifier.py",
        description="Embed the gate-labelled corpus and vote with its neighbours.",
    )
    parser.add_argument("--stage", choices=["build", "fit", "evaluate", "all"],
                        default="all", help="Which part to run (default: all).")
    parser.add_argument("--embedder", choices=["jina", "hash", "lmstudio"],
                        default="jina", help="Embedding route (default: jina).")
    parser.add_argument("--subset", type=int, default=None,
                        help=f"Index a pinned random subset of N rows "
                             f"(P4 uses {SUBSET_ROWS}).")
    parser.add_argument("--seed", type=int, default=SUBSET_SEED,
                        help=f"Subset seed (default: {SUBSET_SEED}).")
    parser.add_argument("--index-dir", default=None,
                        # OCT-LINT: disable=no_hardcoded_secrets reason="plan 289: argparse help text naming the default index path; no credential (operator ruling 2026-09-17)"
                        help="Index directory (default: ../results/knn_index_<tag>).")
    parser.add_argument("--out-dir", default=None,
                        help="Where result files go (default: ../results).")
    parser.add_argument("--tag", default=None,
                        help="Name for this leg (default: <embedder>_<size>).")
    parser.add_argument("--k", type=int, default=None,
                        help="Neighbour count; omit to use the fitted value.")
    parser.add_argument("--temperature", type=float, default=None,
                        help="Softmax temperature; omit to use the fitted value.")
    parser.add_argument("--unmasked", action="store_true",
                        help="Give absent classes exp(0) mass (the other reading of P1).")
    parser.add_argument("--fit-json", default=None,
                        help="Also write the whole fit grid as JSON to this path.")
    return parser
 
 
def _leg_tag(args: argparse.Namespace) -> str:
    """The short name identifying one embedder-and-corpus-size leg."""
 
    if args.tag:
        return str(args.tag)
    embedder = "jina" if args.embedder == "jina" else args.embedder
    size = "full" if args.subset is None else str(args.subset)
    return f"{embedder}_{size}"
 
 
def main(argv: list[str] | None = None) -> int:
    """CLI entry point: build the index, fit on validation, score both sets."""
 
    func = "main"
    args = _build_parser().parse_args(argv)
    bench_dir = Path(__file__).resolve().parent.parent
    tag = _leg_tag(args)
    index_dir = Path(args.index_dir) if args.index_dir else (
        bench_dir / "results" / f"knn_index_{tag}"
    )
    out_dir = Path(args.out_dir) if args.out_dir else bench_dir / "results"
 
    if args.stage in ("build", "all"):
        if args.stage == "build" or not (index_dir / "index.json").is_file():
            embedder = _make_embedder(args.embedder)
            print(f"knn: embedding corpus.jsonl with {embedder.name} "
                  f"(subset {args.subset}) -> {index_dir}")
            build_index(bench_dir / "corpus.jsonl", embedder, index_dir,
                        args.subset, args.seed)
        else:
            print(f"knn: reusing the index at {index_dir}")
    if args.stage == "build":
        _dbg("OC_CHAT.MAIN", func, "exit 0", 2)
        return 0
 
    index = load_index(index_dir, _make_embedder(args.embedder))
    print(f"knn: index {index.rows}x{index.dim} {index.embedder_name}, "
          f"{index.bytes_on_disk / (1 << 20):.1f} MiB, "
          f"built in {float(index.meta.get('build_seconds') or 0.0):.1f}s")
 
    fitted: Fit | None = None
    if args.k is None or args.temperature is None or args.stage == "fit":
        fitted = fit(bench_dir / "corpus_val.jsonl", index)
        print(f"knn: fit on {fitted.source} ({fitted.rows} rows) in "
              f"{fitted.seconds:.1f}s -> k={fitted.k} T={fitted.temperature} "
              f"masked={fitted.mask_absent} Brier {fitted.brier:.4f} "
              f"agreement {fitted.correct}/{fitted.total}")
        print()
        print(render_fit_table(fitted, mask_absent=True))
        print(render_masking_table(fitted))
        if args.fit_json:
            Path(args.fit_json).write_text(
                json.dumps({
                    "leg": tag, "source": fitted.source, "rows": fitted.rows,
                    "chosen": {"k": fitted.k, "temperature": fitted.temperature,
                               "mask_absent": fitted.mask_absent,
                               "brier": fitted.brier,
                               "agreement": f"{fitted.correct}/{fitted.total}"},
                    "grid": [{"k": cell.k, "temperature": cell.temperature,
                              "mask_absent": cell.mask_absent, "brier": cell.brier,
                              "correct": cell.correct, "total": cell.total,
                              "ties": cell.ties} for cell in fitted.grid],
                }, indent=1, sort_keys=True) + "\n",
                encoding="utf-8",
            )
            print(f"knn: wrote {args.fit_json}")
    if args.stage == "fit":
        _dbg("OC_CHAT.MAIN", func, "exit 0", 2)
        return 0
 
    k = args.k if args.k is not None else (fitted.k if fitted else DEFAULT_K)
    temperature = args.temperature if args.temperature is not None else (
        fitted.temperature if fitted else DEFAULT_TEMPERATURE
    )
    mask_absent = False if args.unmasked else (
        fitted.mask_absent if fitted else True
    )
    for condition, case_file in (("hand", "cases.json"), ("hard", "cases_hard.json")):
        payload = evaluate(
            bench_dir / case_file, index, k, temperature, mask_absent,
            model_name=f"knn_{tag}_{condition}", fitted=fitted,
        )
        hits = sum(row["pred"] == row["truth"] for row in payload["results"])
        target = write_result(payload, out_dir / f"knn_{tag}_{condition}.json")
        print(f"  {condition:5} {hits}/{len(payload['results'])} agreement, "
              f"{payload['provenance']['ties_broken']} ties -> {target.name}")
    _dbg("OC_CHAT.MAIN", func, "exit 0", 2)
    return 0
 
 
if __name__ == "__main__":
    sys.exit(main())