encoder_classifier.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# bench/argv_policy/baselines/encoder_classifier.py
 
"""
A small DeBERTa-v3 encoder fine-tuned on the gate-labelled corpus, calibrated.
 
Purpose
-------
The 2026-09-07 classifier note proposed discriminative encoder classifiers as
gates: tens to a few hundred million parameters, CPU milliseconds, a typed
verdict with a probability, trained on our own corpus. The 2026-09-17 session
tested the hosted end of that idea and the zero-shot end of it; this module
tests the proposal itself, which is a small encoder FINE-TUNED on the
gate-labelled commands plan 288 generated. It trains, it calibrates its logits
with one scalar, and it writes result files in the bench schema so the same
scorer reads it beside Jev, DeepSeek and the token floor.
 
Four commands, in the order a run uses them::
 
    train      --backbone <hf id> --out <dir>      # fine-tune, early-stopped
    calibrate  --model <dir>                       # fit one temperature
    evaluate   --model <dir> --cases <file>        # write a result file
    relabel    --out <dir>                         # the policy-change probe
 
``train``, ``calibrate`` and ``evaluate`` need torch and transformers and run
under a SCRATCH HARNESS venv; ``relabel`` needs the argv gate and runs under
the PROJECT venv. That split is why torch and transformers are imported inside
the functions that use them rather than at module scope: the tests in
``tests/test_290_encoder.py`` and ``tests/test_294_speed.py`` exercise the
loader, the temperature fit, the vocabulary pruning, the tensor cache, the
relabel filter and the result writer in the project venv, where no torch
exists, and an import at module scope would put every one of them out of reach.
 
Two scratch venvs, and which one to use
---------------------------------------
Plan 290 ran on the CPU venv at ``...\\harness\\.venv``. Plan 294 added a second
one at ``...\\harness_cuda`` carrying ``torch==2.14.*+cu126``, because the host's
GTX 1060 is a Pascal card (compute capability 6.1) and **the CUDA 12.6 wheel of
PyTorch 2.14 is the last prebuilt build that still carries ``sm_61`` kernels**:
PyTorch dropped Pascal from its CUDA 12.8 and 12.9 wheels at release 2.8 and
CUDA 13 cannot build them. Both venvs stay; the CPU one is what keeps every
number plan 290 produced reproducible, and ``--device cpu`` reproduces a CPU run
inside the CUDA venv as well.
 
Run the trainer from the subproject root with the harness interpreter, with
``HF_HOME`` and ``HF_HUB_OFFLINE=1`` set in the shell::
 
    <harness>\\python.exe bench/argv_policy/baselines/encoder_classifier.py train
        --backbone microsoft/deberta-v3-xsmall --out <scratch>/bench_models/xsmall-7
 
Model directories, their checkpoints and the raw training logs live under
:data:`DEFAULT_MODEL_ROOT` on ``G:`` and are never committed; only the result
JSON and a copy of the training log go into the worktree.
 
The tokenizer
-------------
``AutoTokenizer.from_pretrained("microsoft/deberta-v3-xsmall")`` cannot be
loaded in the harness venv: the snapshot ships only ``spm.model``, and
converting that to a fast tokenizer needs protobuf and sentencepiece, neither
of which is installed (and installing is not this plan's to do). The
``MoritzLaurer/deberta-v3-base-zeroshot-v2.0`` snapshot in the same cache
ships an already-converted ``tokenizer.json``, and its ``spm.model`` is
BYTE-IDENTICAL to both microsoft backbones' -- SHA-256
``c679fbf93643d19aab7ee10c0b99e460bdbc02fedf34b92b05af343b4af586fd`` for all
three. So the tokenizer is loaded from that directory and
:func:`verify_tokenizer_vocab` re-checks the digest equality at every training
run rather than trusting this paragraph.
 
Pricing the direction of an error (plan 298)
--------------------------------------------
Cross-entropy prices an ``allow`` where the gate says ``deny`` exactly as it
prices a ``deny`` where the gate says ``allow``. The bench does not: the column
it is read by counts the errors that GRANT MORE PERMISSION than the gate, and
the other kind costs a human one unnecessary confirmation. Plan 298 puts that
asymmetry into the arithmetic as one cost matrix and two levers over it.
 
``--alpha A`` sets the matrix: zero on the diagonal, ``A`` where the prediction
is more permissive than the truth, ``1`` where it is less. Then:
 
- ``evaluate --decision min-cost`` reads the verdict minimising expected cost
  out of the calibrated probabilities. It retrains nothing and rewrites no
  probability -- only the ``pred`` field moves -- so it costs one flag.
- ``train --cost-weight B`` adds ``B`` times the expected cost under that same
  matrix to the training loss.
 
Both default to today's behaviour exactly, not approximately: ``--alpha 1.0``
makes every off-diagonal cell ``1``, under which the min-cost verdict and the
argmax verdict agree on every input including ties, and ``--cost-weight 0.0``
removes the loss term altogether.
 
Watching the same column the objective prices (plan 300)
--------------------------------------------------------
Plan 298 held the stopping rule fixed on validation cross-entropy so that the
objective was the only variable, and measured what that cost: at
``--cost-weight 3`` the run stopped at epoch 2 and checkpointed an undertrained
model, because an objective that prices the direction moves AWAY from
cross-entropy by construction and a rule watching cross-entropy reads that
movement as a run that has stopped improving. The two pieces are supplied here.
 
- ``train --stop-metric cost`` makes early stopping and checkpoint selection
  watch the subset validation EXPECTED COST under the same ``--alpha`` matrix
  the loss term prices with, rather than the subset cross-entropy. Both figures
  are computed and logged every epoch whichever one stops the run, so an
  after-the-fact reader can see the epoch the other metric would have kept --
  and :func:`select_best_epoch` states it in the training record rather than
  leaving it to be inferred from the log.
- ``evaluate`` counts the PERMISSION-GRANTING errors a result file makes
  confidently, at 0.9 and at 0.99 top probability, into its provenance block.
  Plan 298 found that class by hand and it is the class the escalation
  threshold cannot reach: an error the model is sure about is not escalated and
  no decision rule over a calibrated distribution moves it. Counting it in the
  file makes it a column of the table rather than an afternoon's grep.
 
``--stop-metric ce`` is the default and is plan 290's rule unchanged, so an
unflagged run is still the run every earlier tag on this bench was produced by;
the counts are additive provenance and leave the row shape ``score.py`` reads
untouched.
 
Responsibilities
----------------
- :func:`held_out_commands` / :func:`load_split` / :func:`label_counts` --
  the corpus loader, with the hand-case guard.
- :func:`epoch_batches` -- one epoch of length-bucketed minibatches.
- :func:`resolve_device` -- plan 294 P1's device selection.
- :func:`prune_vocab` / :func:`token_coverage` -- plan 294 P2's vocabulary
  pruning, in pure Python.
- :func:`remap_lookup` / :func:`remap_table` -- the pruned id space.
- :func:`cache_key` / :func:`load_cached_ids` / :func:`store_cached_ids` /
  :func:`validation_subset` / :func:`resolve_max_length` -- plan 294 P3's
  tokenised-id cache and subset validation.
- :func:`fit_temperature` / :func:`mean_negative_log_likelihood` /
  :func:`probabilities` / :func:`brier` -- calibration, in pure Python.
- :func:`cost_matrix` / :func:`expected_cost_of` / :func:`min_cost_verdict` /
  :func:`decide` / :func:`expected_cost` -- plan 298's direction-weighted cost
  matrix, its decision rule and its loss term, in pure Python.
- :func:`cost_tensor` / :func:`expected_cost_term` -- the same loss term on
  tensors, which is what the training loop actually calls.
- :func:`stopping_value` / :func:`select_best_epoch` -- plan 300 P1's
  direction-aware stopping rule, in pure Python.
- :func:`grants_more_permission` / :func:`confident_error_counts` -- plan 300
  P2's confident permission-granting error count, in pure Python.
- :func:`strip_recursive_option` / :func:`relabel_recursive_allow` /
  :func:`gate_classifier` -- the policy-change probe's documented relabel.
- :func:`result_row` / :func:`write_result_file` -- the plan-288 result schema.
- :func:`train` / :func:`calibrate` / :func:`evaluate` -- the torch commands.
- :func:`main` -- the CLI.
 
Diagnostics
-----------
Domain: OC_CHAT.MAIN
Levels:
    L2 - command lifecycle (training start and finish, files written)
    L3 - per-epoch validation loss, per-case-set agreement, fitted temperature
    L4 - module load
 
Contracts
---------
- Read-only with respect to the committed corpus, the case sets and the gate:
  the policy-change probe writes a RELABELLED COPY under the scratch root and
  never edits a committed file.
- The relabelled labels come from the unmodified gate applied to a rewritten
  command, never from a hand edit of a verdict.
- Bitwise reproduction is not claimed (CPU thread scheduling reorders float
  reductions); a run is IDENTIFIABLE instead -- every result file carries the
  backbone, the seed, the epoch count, the training wall time, the parameter
  count and the SHA-256 of the saved weights.
- Nothing from either case set touches training or calibration.
- Console output is ASCII only; commands are printed through :func:`ascii`.
- No dependency enters any ``pyproject.toml``; torch and transformers are
  imported lazily and only exist in the scratch harness venv.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
- torch (LAZY, harness venv only: tensors, AdamW, the forward and backward pass)
- transformers (LAZY, harness venv only: AutoTokenizer and AutoModelForSequenceClassification)
- oc_chat.tools._argv_policy (LAZY, project venv only: the gate the relabel asks)
"""
from __future__ import annotations
 
import argparse
import hashlib
import json
import math
import random
import sys
import time
from dataclasses import dataclass
from pathlib import Path, PurePath
 
from oc_diagnostics import _dbg
 
_dbg("OC_CHAT.MAIN", "module", "loaded", 4)
 
#: The three verdicts, in the column order of the classification head.
LABELS: tuple[str, ...] = ("allow", "confirm", "deny")
 
#: Plan 298 P1. The same three verdicts ordered by how much permission they
#: grant, LEAST first. Pinned as a tuple rather than derived from
#: :data:`LABELS`, because every figure this plan produces is read off this
#: order and an order inferred from the head's column order would follow a
#: reordering of the head without saying so. It states what ``score.py``'s
#: ``PERMISSIVENESS`` states, and ``tests/test_298_direction.py`` asserts the
#: two agree rather than letting the bench carry two opinions about which way
#: is dangerous.
PERMISSIVENESS: tuple[str, ...] = ("deny", "confirm", "allow")
 
#: Plan 298 P1. How much more an error granting MORE permission costs than one
#: granting less. ``1.0`` prices them the same, which is the matrix plain
#: cross-entropy already implies, so it is the default and every number this
#: bench carried before plan 298 reproduces under it.
DEFAULT_ALPHA = 1.0
 
#: Plan 298 P3. The weight on the expected-cost term added to the training
#: loss. ``0.0`` removes the term entirely, so an unflagged training run is
#: plan 295's training run.
DEFAULT_COST_WEIGHT = 0.0
 
#: Plan 298 P2. How a verdict is read out of a calibrated distribution.
#: ``argmax`` is the rule every result file before this plan was written under;
#: ``min-cost`` minimises expected cost under the cost matrix.
DECISION_CHOICES: tuple[str, ...] = ("argmax", "min-cost")
DEFAULT_DECISION = "argmax"
 
#: Plan 300 P1. Which validation figure early stopping and checkpoint selection
#: watch, in the order :func:`select_best_epoch` reads an ``(ce, cost)`` pair:
#: ``ce`` is the subset validation cross-entropy plan 290's rule watches and is
#: the default, so an unflagged run is the run every earlier tag was produced
#: by; ``cost`` is the subset validation expected cost under the ``--alpha``
#: matrix, which is the column a direction-weighted objective is actually
#: trying to move.
STOP_METRIC_CHOICES: tuple[str, ...] = ("ce", "cost")
DEFAULT_STOP_METRIC = "ce"
 
#: Plan 300 P2. The top-probability thresholds the permission-granting errors
#: are counted at. 0.9 is ``score.py``'s own escalation threshold, so the count
#: at 0.9 is exactly the set of dangerous errors escalation would NOT catch;
#: 0.99 separates the merely confident from the ones plan 298 found at 0.999,
#: which no decision rule over a calibrated distribution can reach.
CONFIDENT_THRESHOLDS: tuple[float, ...] = (0.9, 0.99)
 
#: The pinned seed. It fixes the classifier head's initialisation and the batch
#: order, and it names the model directory, so two runs are distinguishable.
DEFAULT_SEED = 7
 
#: Plan 290 P2. Commands are short -- the corpus's median is 10 tokens and its
#: longest is 44 -- so 64 is a cap that never truncates rather than a shape.
DEFAULT_MAX_LENGTH = 64
DEFAULT_BATCH = 32
DEFAULT_LEARNING_RATE = 3e-5
DEFAULT_MAX_EPOCHS = 4
DEFAULT_WARMUP_SHARE = 0.1
 
#: Threads are pinned so the reported CPU latency means something. Measured on
#: this host: 8 threads is no slower than the 14 available (1.22 s against
#: 1.28 s per training step), so pinning costs nothing.
DEFAULT_THREADS = 8
 
#: Plan 294 P1. ``auto`` takes the GPU when torch reports one and the CPU
#: otherwise; ``cuda`` REFUSES rather than falling back, because a silent
#: fallback turns a 90-second epoch into an 18-minute one and the only evidence
#: is the wall clock.
DEVICE_CHOICES: tuple[str, ...] = ("auto", "cpu", "cuda")
DEFAULT_DEVICE = "auto"
 
#: Plan 294 P2. Written beside the weights when and only when the embedding
#: table was sliced, so ``calibrate`` and ``evaluate`` can tell which id space
#: a saved model is indexed by without being told on the command line.
VOCAB_PRUNE_FILE = "vocab_prune.json"
 
#: Plan 294 P2. How many unseen vocabulary rows the pruning keeps anyway, so an
#: unseen but plausible piece still has a row of its own rather than collapsing
#: onto the unknown token. A SentencePiece vocabulary is emitted in descending
#: unigram-score order after its control symbols, so "the most frequent ids that
#: do not occur" is "the lowest-numbered ids that do not occur" and no frequency
#: table has to be carried alongside the model.
DEFAULT_VOCAB_RESERVE = 2000
 
#: Plan 294 P3. Where the tokenised corpus is cached. Under the scratch root
#: with the models, never in the worktree: it is rebuildable, it is megabytes,
#: and ``R:`` refuses writes under host memory pressure.
DEFAULT_CACHE_ROOT = Path(r"G:\Temp\claude\g--PythonApps-option-c\harness\bench_cache")
 
#: Plan 294 P3. The per-epoch validation subset: how many rows, and the seed
#: that fixes WHICH rows, so two runs validate on the same commands.
DEFAULT_SUBSET_ROWS = 1000
DEFAULT_SUBSET_SEED = 294
 
#: Plan 294 P3. The corpus's longest command is 44 tokens, so 48 is a cap that
#: still never truncates. :func:`resolve_max_length` re-measures rather than
#: trusting this number, and falls back to :data:`DEFAULT_MAX_LENGTH` if a
#: regenerated corpus ever exceeds it.
FAST_MAX_LENGTH = 48
 
#: Plan 290 P2's wall-time budget per backbone. Training stops at the end of
#: the epoch after which another epoch would not fit, and the record says the
#: stop was budget-bound rather than convergence-bound.
DEFAULT_BUDGET_SECONDS = 3600.0
 
#: How much better a validation cross-entropy has to be to count as better.
EARLY_STOP_MIN_DELTA = 1e-4
 
#: Batches are formed inside chunks of this many, sorted by length, so padding
#: follows the commands in the chunk instead of the longest in the epoch.
BUCKET_BATCHES = 50
 
#: The scratch root every model directory, checkpoint and raw log lives under.
#: ``R:`` refuses writes under host memory pressure and nothing here is
#: committed, so neither the worktree nor the RAM disk is the place for it.
DEFAULT_MODEL_ROOT = Path(
    r"G:\Temp\claude\g--PythonApps-option-c\harness\bench_models"
)
 
#: The snapshot holding the converted DeBERTa-v3 fast tokenizer (see "The
#: tokenizer" above). Its vocabulary is verified against the backbone's own
#: ``spm.model`` by digest at every training run.
DEFAULT_TOKENIZER_DIR = Path(
    # OCT-LINT: disable=no_hardcoded_secrets reason="local HF snapshot path; the entropy is the 40-hex model revision id, not a credential (operator ruling 2026-09-17)"
    r"G:\Temp\claude\g--PythonApps-option-c\harness\hf_home\hub"
    r"\models--MoritzLaurer--deberta-v3-base-zeroshot-v2.0\snapshots"
    r"\8e7e5af5983a0ddb1a5b45a38b129ab69e2258e8"
)
 
#: The HF cache the backbones were downloaded into, used to locate a backbone's
#: ``spm.model`` for the digest check.
DEFAULT_HF_HOME = Path(r"G:\Temp\claude\g--PythonApps-option-c\harness\hf_home")
 
#: The programs whose option list the policy-change probe edits. The gate
#: normalises a program the same way: basename, lowercased, ``.exe`` stripped.
LISTING_PROGRAMS: frozenset[str] = frozenset({"ls", "dir"})
 
#: The option the probe turns from unknown into allowed.
RECURSIVE_OPTION = "R"
 
 
# --------------------------------------------------------------------------
# The corpus loader
# --------------------------------------------------------------------------
 
 
def read_jsonl(path: Path) -> list[dict[str, str]]:
    """Read a JSON-lines file 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 held_out_commands(bench_dir: Path) -> frozenset[str]:
    """Every command in the two case sets, which no training row may repeat."""
 
    func = "held_out_commands"
    commands: set[str] = set()
    for name in ("cases.json", "cases_hard.json"):
        path = Path(bench_dir) / name
        for case in json.loads(path.read_text(encoding="utf-8")):
            commands.add(case["command"])
    _dbg("OC_CHAT.MAIN", func, f"{len(commands)} held-out commands", 3)
    return frozenset(commands)
 
 
def load_split(bench_dir: Path, name: str, held_out: frozenset[str]) -> list[dict[str, str]]:
    """Load one corpus split, refusing any row that is a held-out case.
 
    ``gen_corpus.py`` already drops the case sets at generation time, so this
    guard should never fire -- which is exactly why it is worth asserting here:
    a hand case that reached training would raise every agreement figure on
    this bench and nothing downstream would look wrong.
    """
 
    func = "load_split"
    rows = read_jsonl(Path(bench_dir) / name)
    leaked = [row["command"] for row in rows if row["command"] in held_out]
    if leaked:
        raise ValueError(
            f"{name}: {len(leaked)} held-out case(s) present in the corpus, "
            f"first is {ascii(leaked[0])}"
        )
    _dbg("OC_CHAT.MAIN", func, f"{name}: {len(rows)} rows", 3)
    return rows
 
 
def label_counts(rows: list[dict[str, str]]) -> dict[str, int]:
    """How many rows carry each verdict."""
 
    return {label: sum(1 for row in rows if row["truth"] == label) for label in LABELS}
 
 
def epoch_batches(
    rows: list[dict[str, str]], batch: int, rng: random.Random
) -> list[list[dict[str, str]]]:
    """One epoch of minibatches: a permutation of *rows*, length-bucketed.
 
    An epoch is a permutation, so the class balance a batch sequence sees is
    the corpus's own and no class weighting is needed. Within a chunk of
    :data:`BUCKET_BATCHES` batches the rows are sorted by command length before
    being cut into batches, and the batch order is then shuffled: padding
    follows the chunk rather than the longest command in the epoch, which on
    this corpus is the difference between 20.6 and 17.1 minutes per epoch, and
    the shuffle keeps the gradient sequence stochastic.
    """
 
    order = list(range(len(rows)))
    rng.shuffle(order)
    chunk_size = batch * BUCKET_BATCHES
    batches: list[list[dict[str, str]]] = []
    for start in range(0, len(order), chunk_size):
        chunk = sorted(order[start:start + chunk_size], key=lambda index: len(rows[index]["command"]))
        for offset in range(0, len(chunk), batch):
            batches.append([rows[index] for index in chunk[offset:offset + batch]])
    rng.shuffle(batches)
    return batches
 
 
# --------------------------------------------------------------------------
# Plan 294 P2 and P3: pruning, caching and subset validation, in pure Python
# --------------------------------------------------------------------------
 
 
def tokenise(commands: list[str], tokenizer: object, max_length: int) -> list[list[int]]:
    """The token ids of every command, unpadded, truncated at *max_length*."""
 
    encoded = tokenizer(list(commands), truncation=True,        # type: ignore[operator]
                        max_length=max_length)
    return [list(ids) for ids in encoded["input_ids"]]
 
 
def tokenise_keys(tokenizer: object) -> tuple[str, ...]:
    """Which tensors this tokenizer emits, asked of it rather than assumed."""
 
    return tuple(tokenizer(["ls"]).keys())                   # type: ignore[operator]
 
 
def full_vocab_size(tokenizer: object) -> int:
    """How many ids the tokenizer can emit, counting any added tokens."""
 
    try:
        return len(tokenizer)                                   # type: ignore[arg-type]
    except TypeError:                                           # pragma: no cover - stubs
        return int(tokenizer.vocab_size)                        # type: ignore[attr-defined]
 
 
def prune_vocab(
    corpus: list[str],
    tokenizer: object,
    reserve: int = DEFAULT_VOCAB_RESERVE,
    max_length: int = DEFAULT_MAX_LENGTH,
) -> tuple[list[int], dict[int, int]]:
    """Plan 294 P2: the vocabulary rows worth keeping, and the map onto them.
 
    DeBERTa-v3-xsmall carries 70.8M parameters of which the 128,100-row
    embedding table at width 384 is about 49M -- two thirds of the model
    devoted to pieces this corpus never emits. The kept set is every id the
    corpus actually produces, plus every special token, plus a reserve of the
    most frequent ids that do NOT occur, so an unseen but plausible piece still
    has a row of its own rather than collapsing onto the unknown token.
 
    Returns the kept ids ASCENDING and a map from old id to its position in
    that list. Ascending order is not cosmetic: it keeps the control symbols at
    the front in their original relative order, so the padding id survives the
    remap as itself and a saved config's ``pad_token_id`` stays true.
    """
 
    func = "prune_vocab"
    occurring: set[int] = set()
    for ids in tokenise(corpus, tokenizer, max_length):
        occurring.update(int(value) for value in ids)
    specials = {int(value) for value in tokenizer.all_special_ids}   # type: ignore[attr-defined]
    seen = occurring | specials
    full = full_vocab_size(tokenizer)
    reserved: list[int] = []
    for token_id in range(full):
        if len(reserved) >= reserve:
            break
        if token_id not in seen:
            reserved.append(token_id)
    kept = sorted(seen | set(reserved))
    remap = {old: new for new, old in enumerate(kept)}
    _dbg("OC_CHAT.MAIN", func,
         f"{len(kept)} of {full} rows kept ({len(occurring)} seen, "
         f"{len(reserved)} reserved)", 2)
    return kept, remap
 
 
def token_coverage(commands: list[str], tokenizer: object, kept: set[int] | list[int],
                   max_length: int = DEFAULT_MAX_LENGTH) -> float:
    """The share of token OCCURRENCES in *commands* that a kept id covers.
 
    Occurrences rather than distinct ids: one rare piece appearing once in one
    hard case is not the same loss as a common piece appearing in a hundred,
    and the figure that says whether pruning hurt the case sets is the one
    weighted the way the case sets actually use the vocabulary.
    """
 
    kept_set = {int(value) for value in kept}
    total = 0
    covered = 0
    for ids in tokenise(commands, tokenizer, max_length):
        for value in ids:
            total += 1
            covered += int(int(value) in kept_set)
    return covered / total if total else 1.0
 
 
def resolve_max_length(commands: list[str], tokenizer: object, requested: int,
                       fallback: int = DEFAULT_MAX_LENGTH) -> int:
    """*requested* if it never truncates *commands*, else *fallback*.
 
    Plan 294 P3 shortens the sequence from 64 to 48 because the corpus's
    longest command is 44 tokens. That is a measured fact about a committed
    file rather than a property of the task, so it is re-measured here: a
    regenerated corpus carrying a longer command silently loses its tail
    otherwise, and a truncated command is a mislabelled training row.
    """
 
    func = "resolve_max_length"
    longest = max((len(ids) for ids in tokenise(commands, tokenizer, fallback)), default=0)
    if longest <= requested:
        return requested
    _dbg("OC_CHAT.MAIN", func,
         f"longest command is {longest} tokens, over {requested}; using {fallback}", 2)
    return fallback
 
 
def cache_key(corpus_path: Path, tokenizer_digest: str, max_length: int) -> str:
    """The cache identity of one tokenised split.
 
    The corpus file's own SHA-256 is in the key, so a regenerated corpus
    invalidates the cache instead of training the next run on the previous
    corpus's ids -- which would be a silent wrong answer, not a crash.
    """
 
    material = f"{file_digest(corpus_path)}:{tokenizer_digest}:{max_length}"
    return hashlib.sha256(material.encode("utf-8")).hexdigest()[:32]
 
 
def load_cached_ids(cache_root: Path, key: str) -> list[list[int]] | None:
    """The tokenised split cached under *key*, or ``None`` when it is absent."""
 
    func = "load_cached_ids"
    path = Path(cache_root) / f"ids_{key}.json"
    if not path.is_file():
        return None
    rows = json.loads(path.read_text(encoding="utf-8"))
    _dbg("OC_CHAT.MAIN", func, f"{len(rows)} rows from {path.name}", 3)
    return [list(row) for row in rows]
 
 
def store_cached_ids(cache_root: Path, key: str, rows: list[list[int]]) -> Path:
    """Write a tokenised split into the cache and return where it landed."""
 
    func = "store_cached_ids"
    Path(cache_root).mkdir(parents=True, exist_ok=True)
    path = Path(cache_root) / f"ids_{key}.json"
    path.write_text(json.dumps(rows), encoding="utf-8")
    _dbg("OC_CHAT.MAIN", func, f"{len(rows)} rows to {path.name}", 3)
    return path
 
 
def remap_lookup(remap: dict[int, int], unknown_id: int, size: int) -> list[int]:
    """The old-id to new-id lookup as a plain list, unknown-filled.
 
    Every id the pruning dropped points at the unknown token, which is the
    pruning's whole contract for a piece the corpus never showed. Built here in
    pure Python rather than inside the tensor constructor so the contract is
    assertable in the project venv, where there is no torch.
    """
 
    lookup = [unknown_id] * size
    for old, new in remap.items():
        lookup[old] = new
    return lookup
 
 
def validation_subset(rows: list[dict[str, str]], size: int,
                      seed: int = DEFAULT_SUBSET_SEED) -> list[dict[str, str]]:
    """A fixed *size*-row sample of *rows*, the same one for every run.
 
    The sample is drawn from a pinned seed and then restored to corpus order,
    so the subset a run validates on depends on the seed alone and not on the
    order the shuffle happened to produce.
    """
 
    if size <= 0 or size >= len(rows):
        return list(rows)
    order = list(range(len(rows)))
    random.Random(seed).shuffle(order)
    return [rows[index] for index in sorted(order[:size])]
 
 
# --------------------------------------------------------------------------
# Calibration, in pure Python
# --------------------------------------------------------------------------
 
 
def probabilities(logits: list[float], temperature: float = 1.0) -> dict[str, float]:
    """The temperature-scaled softmax of one logit row, keyed by verdict."""
 
    scaled = [value / temperature for value in logits]
    top = max(scaled)
    exponentials = [math.exp(value - top) for value in scaled]
    total = sum(exponentials)
    return {label: value / total for label, value in zip(LABELS, exponentials)}
 
 
def brier(probs: dict[str, float], truth: str) -> float:
    """The multi-class Brier score of one forecast (``score.brier_row``).
 
    Kept here rather than imported because ``score.py`` sits one directory up
    and this module is run as a script from the harness venv; the arithmetic is
    two lines and ``tests/test_290_encoder.py`` asserts the two agree.
    """
 
    return sum((probs.get(label, 0.0) - (1.0 if label == truth else 0.0)) ** 2
               for label in LABELS)
 
 
def mean_negative_log_likelihood(
    logits: list[list[float]],
    labels: list[int],
    temperature: float,
    weights: list[float] | None = None,
) -> float:
    """Mean NLL of *logits* divided by *temperature* against *labels*.
 
    *weights* exist for the witness, not for the run: weighting each logit row
    by the true probability of each class turns this into the EXPECTED NLL,
    whose minimiser over the scalar is exactly the temperature the logits were
    scaled by, which is what lets the test assert a known answer to two
    decimals instead of a sampling interval.
    """
 
    total = 0.0
    mass = 0.0
    for index, row in enumerate(logits):
        weight = 1.0 if weights is None else weights[index]
        if weight == 0.0:
            continue
        scaled = [value / temperature for value in row]
        top = max(scaled)
        normaliser = top + math.log(sum(math.exp(value - top) for value in scaled))
        total += weight * (normaliser - scaled[labels[index]])
        mass += weight
    return total / mass if mass else 0.0
 
 
def fit_temperature(
    logits: list[list[float]],
    labels: list[int],
    weights: list[float] | None = None,
    low: float = 0.05,
    high: float = 20.0,
) -> float:
    """Fit one temperature by minimising NLL, by golden-section search.
 
    Plan 290 P3 asks for the negative log-likelihood to be minimised over one
    scalar with the optimisation written out rather than called from a fitting
    library, and W1 asks for a fit that runs in the project venv, where there
    is no torch. A golden-section search over log-temperature satisfies both:
    the objective is unimodal in log T, the search needs no derivative, and
    eighty iterations close the bracket to about 1e-7.
    """
 
    func = "fit_temperature"
    ratio = (math.sqrt(5.0) - 1.0) / 2.0
    left, right = math.log(low), math.log(high)
    probe_a, probe_b = right - ratio * (right - left), left + ratio * (right - left)
    value_a = mean_negative_log_likelihood(logits, labels, math.exp(probe_a), weights)
    value_b = mean_negative_log_likelihood(logits, labels, math.exp(probe_b), weights)
    for _ in range(80):
        if value_a <= value_b:
            right, probe_b, value_b = probe_b, probe_a, value_a
            probe_a = right - ratio * (right - left)
            value_a = mean_negative_log_likelihood(logits, labels, math.exp(probe_a), weights)
        else:
            left, probe_a, value_a = probe_a, probe_b, value_b
            probe_b = left + ratio * (right - left)
            value_b = mean_negative_log_likelihood(logits, labels, math.exp(probe_b), weights)
    fitted = math.exp((left + right) / 2.0)
    _dbg("OC_CHAT.MAIN", func, f"temperature {fitted:.4f} over {len(logits)} rows", 3)
    return fitted
 
 
def calibration_report(
    logits: list[list[float]], rows: list[dict[str, str]], temperature: float
) -> dict[str, float | int]:
    """NLL, Brier and agreement on the validation split, before and after."""
 
    labels = [LABELS.index(row["truth"]) for row in rows]
    before = [probabilities(row, 1.0) for row in logits]
    after = [probabilities(row, temperature) for row in logits]
    hits = sum(max(probs, key=lambda label: probs[label]) == row["truth"]
               for probs, row in zip(before, rows))
    return {
        "validation_rows": len(rows),
        "temperature": temperature,
        "nll_before": mean_negative_log_likelihood(logits, labels, 1.0),
        "nll_after": mean_negative_log_likelihood(logits, labels, temperature),
        "brier_before": sum(brier(probs, row["truth"])
                            for probs, row in zip(before, rows)) / len(rows),
        "brier_after": sum(brier(probs, row["truth"])
                           for probs, row in zip(after, rows)) / len(rows),
        "validation_accuracy": hits / len(rows),
    }
 
 
# --------------------------------------------------------------------------
# Plan 298: pricing the direction of an error, in pure Python
# --------------------------------------------------------------------------
 
 
def cost_matrix(alpha: float = DEFAULT_ALPHA) -> list[list[float]]:
    """Plan 298 P1: ``COST[true][pred]``, both axes in :data:`LABELS` order.
 
    Zero on the diagonal; *alpha* in every cell where the prediction grants MORE
    permission than the gate's verdict; ``1.0`` in every cell where it grants
    less. At *alpha* ``1.0`` every off-diagonal cell is ``1.0``, which is the
    matrix a plain cross-entropy objective and an argmax decision rule already
    imply -- which is why that is the default and why it reproduces every number
    on this bench rather than approximating it.
 
    The asymmetry is the whole of this plan. An ``allow`` where the gate says
    ``deny`` runs something the gate refused; a ``deny`` where the gate says
    ``allow`` costs a human one unnecessary confirmation. Cross-entropy prices
    those two mistakes identically. The gate does not, and neither does the
    column this bench is read by.
    """
 
    func = "cost_matrix"
    ranks = {label: PERMISSIVENESS.index(label) for label in LABELS}
    matrix = [
        [0.0 if pred == true else (alpha if ranks[pred] > ranks[true] else 1.0)
         for pred in LABELS]
        for true in LABELS
    ]
    _dbg("OC_CHAT.MAIN", func, f"alpha {alpha}", 4)
    return matrix
 
 
def expected_cost_of(probs: dict[str, float], pred: str,
                     costs: list[list[float]]) -> float:
    """What predicting *pred* costs in expectation under the forecast *probs*.
 
    ``sum_true p(true) * COST[true][pred]`` -- the quantity plan 298 P2's rule
    minimises, named so a test can assert the arithmetic directly rather than
    only through the verdict it happens to produce.
    """
 
    pred_index = LABELS.index(pred)
    return sum(probs.get(LABELS[true_index], 0.0) * costs[true_index][pred_index]
               for true_index in range(len(LABELS)))
 
 
def min_cost_verdict(probs: dict[str, float], costs: list[list[float]]) -> str:
    """Plan 298 P2: the verdict with the lowest expected cost under *probs*.
 
    ``argmin_pred sum_true p(true) * COST[true][pred]``. This is the cheap
    lever: it changes no weights and touches no probability, only which verdict
    is read out of one, so a model already trained and calibrated can be made
    more cautious for the price of a second pass over its own output.
 
    A tie breaks towards the earlier :data:`LABELS` column, which is precisely
    how ``max`` breaks an argmax tie over the same dict. That matters more than
    it looks: it is what makes this rule agree with the argmax rule on EVERY
    input at *alpha* ``1.0``, rather than on every input that happens not to be
    tied, so the default really is the old behaviour and not merely close to it.
    """
 
    return min(LABELS, key=lambda pred: expected_cost_of(probs, pred, costs))
 
 
def decide(probs: dict[str, float], decision: str = DEFAULT_DECISION,
           costs: list[list[float]] | None = None) -> str:
    """The verdict *decision* reads out of one calibrated distribution.
 
    The two rules are kept behind one function because ``result_row`` must not
    grow a branch: a result file differs under the two rules in its ``pred``
    field and in nothing else, and that stays true more easily when the choice
    is made in one place.
    """
 
    func = "decide"
    if decision not in DECISION_CHOICES:
        raise ValueError(
            f"{func}: decision must be one of {DECISION_CHOICES}, not {decision!r}"
        )
    if decision == "argmax":
        return max(probs, key=lambda label: probs[label])
    return min_cost_verdict(probs, costs if costs is not None else cost_matrix())
 
 
def expected_cost(prob_rows: list[list[float]], labels: list[int],
                  costs: list[list[float]]) -> float:
    """Plan 298 P3's training term, written out in plain Python.
 
    Per row ``sum_pred p[pred] * COST[true][pred]``, averaged over the batch so
    the term lands on the same scale as the mean cross-entropy it is added to
    and ``--cost-weight`` therefore means the same thing at any batch size.
 
    This function is the DEFINITION and :func:`expected_cost_term` is the fast
    copy of it; ``tests/test_298_direction.py`` asserts the two agree on a toy
    batch, which is the only reason the tensor version can be trusted inside a
    run whose per-step output nobody checks by hand.
    """
 
    if not prob_rows:
        return 0.0
    total = 0.0
    for row, label in zip(prob_rows, labels):
        total += sum(row[index] * costs[label][index] for index in range(len(LABELS)))
    return total / len(prob_rows)
 
 
# --------------------------------------------------------------------------
# Plan 300: stopping on the column the objective prices, in pure Python
# --------------------------------------------------------------------------
 
 
def stopping_value(record: dict[str, float | int], metric: str = DEFAULT_STOP_METRIC) -> float:
    """Plan 300 P1: the figure a run watching *metric* compares for one epoch.
 
    One function rather than a conditional expression in two places, because the
    training loop and :func:`select_best_epoch` have to agree about which number
    is being watched or the record's "the other metric would have chosen epoch
    N" line is measuring something no run would have done.
    """
 
    func = "stopping_value"
    if metric not in STOP_METRIC_CHOICES:
        raise ValueError(
            f"{func}: stop metric must be one of {STOP_METRIC_CHOICES}, not {metric!r}"
        )
    key = "validation_loss" if metric == "ce" else "validation_cost"
    return float(record[key])
 
 
def select_best_epoch(
    pairs: list[tuple[float, float]],
    metric: str = DEFAULT_STOP_METRIC,
    patience: int = 1,
    min_delta: float = EARLY_STOP_MIN_DELTA,
) -> dict[str, float | int | bool]:
    """Plan 300 P1: which epoch a run watching *metric* keeps, and when it stops.
 
    *pairs* is one ``(cross-entropy, expected cost)`` per epoch in epoch order,
    in :data:`STOP_METRIC_CHOICES` order, which is how *metric* selects the
    element to compare. Returns the 1-based best epoch (``0`` when no epoch ever
    improved), the value it scored, how many epochs the run would have got
    through, and whether patience ended it.
 
    This is the loop's rule written out as a function, and it earns its place
    twice. It is what W1 asserts on hand-written sequences, where the answer is
    known and a real run's is not. And it is what the training record calls on
    at the end with the OTHER metric, so the plan-300 table can say which epoch
    a run would have checkpointed had it watched the other column -- the one
    figure that separates "the stopping rule moved the model" from "the run
    happened to differ". Run over the epochs one run actually produced: a run
    watching the other metric might have stopped earlier and never seen the
    later pairs, so the alternate epoch is a reading of THIS run's curve, not a
    prediction of a run nobody made.
    """
 
    func = "select_best_epoch"
    if metric not in STOP_METRIC_CHOICES:
        raise ValueError(
            f"{func}: stop metric must be one of {STOP_METRIC_CHOICES}, not {metric!r}"
        )
    index = STOP_METRIC_CHOICES.index(metric)
    best_value = float("inf")
    best_epoch = 0
    epochs_run = 0
    stale = 0
    stopped_early = False
    for epoch, pair in enumerate(pairs, 1):
        epochs_run = epoch
        value = float(pair[index])
        if value < best_value - min_delta:
            stale = 0
            best_value, best_epoch = value, epoch
            continue
        stale += 1
        if stale >= patience:
            stopped_early = True
            break
    _dbg("OC_CHAT.MAIN", func,
         f"{metric}: best epoch {best_epoch} of {epochs_run} run", 4)
    return {
        "best_epoch": best_epoch,
        "best_value": best_value,
        "epochs_run": epochs_run,
        "stopped_early": stopped_early,
    }
 
 
def other_stop_metric(metric: str) -> str:
    """The stopping metric *metric* is not, which is what the record compares."""
 
    func = "other_stop_metric"
    if metric not in STOP_METRIC_CHOICES:
        raise ValueError(
            f"{func}: stop metric must be one of {STOP_METRIC_CHOICES}, not {metric!r}"
        )
    return next(choice for choice in STOP_METRIC_CHOICES if choice != metric)
 
 
def grants_more_permission(truth: str, pred: str) -> bool:
    """Whether predicting *pred* where the gate said *truth* grants MORE permission.
 
    ``score.py``'s :func:`error_direction` in a boolean, and asserted equal to
    it in ``tests/test_300_stopping.py`` rather than left as two opinions about
    which way is dangerous. An unrecognised prediction counts as granting more,
    for the reason the scorer gives: understating that column is the one
    direction this bench must never round in.
    """
 
    if pred == truth:
        return False
    ranks = {label: PERMISSIVENESS.index(label) for label in PERMISSIVENESS}
    if pred not in ranks:
        return True
    return ranks[pred] > ranks.get(truth, 0)
 
 
def confident_error_counts(
    rows: list[dict[str, object]], thresholds: tuple[float, ...] = CONFIDENT_THRESHOLDS
) -> dict[str, int]:
    """Plan 300 P2: how many permission-granting errors were made CONFIDENTLY.
 
    Returns the total, and the count at or above each threshold's top
    probability. The class matters because it is the one nothing downstream can
    catch: escalation fires BELOW its threshold, so an error at 0.99 is not
    escalated; and a min-cost reading moves a verdict only where the
    distribution leaves enough mass elsewhere to move, which at 0.999 it does
    not. Plan 298 found five of ``small``'s thirteen at 0.9 or better and three
    at 0.999 or better by hand; this makes it a number every result file
    carries.
 
    A row with no usable forecast counts towards the total and towards no
    threshold, which is the same treatment ``score.py`` gives it in the Brier
    column: a missing probability is not evidence of a confident mistake.
    """
 
    func = "confident_error_counts"
    counts = {"granted_more": 0}
    for threshold in thresholds:
        counts[f"granted_more_at_{threshold:g}"] = 0
    for row in rows:
        if not grants_more_permission(str(row.get("truth", "")), str(row.get("pred", ""))):
            continue
        counts["granted_more"] += 1
        probs = row.get("probs")
        if not isinstance(probs, dict) or not probs:
            continue
        top = max(float(value) for value in probs.values())
        for threshold in thresholds:
            if top >= threshold:
                counts[f"granted_more_at_{threshold:g}"] += 1
    _dbg("OC_CHAT.MAIN", func,
         f"{counts['granted_more']} permission-granting error(s) over {len(rows)} rows", 3)
    return counts
 
 
# --------------------------------------------------------------------------
# The policy-change probe's relabel
# --------------------------------------------------------------------------
 
 
def _normalised_program(token: str) -> str:
    """The program name the gate would read out of *token*."""
 
    name = PurePath(token).name.lower()
    return name[:-4] if name.endswith(".exe") else name
 
 
def _is_short_cluster(token: str) -> bool:
    """Whether *token* is a short-option cluster such as ``-la``."""
 
    return len(token) > 1 and token[0] == "-" and token[1] != "-" and token[1:].isalnum()
 
 
def strip_recursive_option(command: str) -> str | None:
    """*command* with the recursive listing option removed, or ``None``.
 
    This is the whole of the policy change the probe measures: ``-R`` becomes
    an allowed option of the Tier-1 listing commands, and nothing else moves.
    Because the option carries no other meaning to the gate, the verdict of a
    command under the changed policy is the verdict the UNMODIFIED gate gives
    the same command with ``-R`` removed -- which is how
    :func:`relabel_recursive_allow` derives every new label without editing
    ``_argv_policy.py`` and without hand-writing a verdict.
 
    ``None`` when the command is not a listing command carrying an unquoted
    ``-R``: a quoted ``'-R'`` is data, ``--reverse`` is a different option, and
    ``python -R`` is a real option of a different program in a different tier.
    """
 
    tokens = command.split()
    if not tokens or _normalised_program(tokens[0]) not in LISTING_PROGRAMS:
        return None
    kept = [tokens[0]]
    matched = False
    for token in tokens[1:]:
        if _is_short_cluster(token) and RECURSIVE_OPTION in token[1:]:
            matched = True
            remainder = token[1:].replace(RECURSIVE_OPTION, "")
            if remainder:
                kept.append(f"-{remainder}")
            continue
        kept.append(token)
    return " ".join(kept) if matched else None
 
 
def relabel_recursive_allow(
    rows: list[dict[str, str]], classify: object
) -> tuple[list[dict[str, str]], int]:
    """Relabel every matched row by asking the gate about the stripped command.
 
    Returns the new rows and how many verdicts actually moved. A matched row
    whose verdict does not move is the interesting case and is kept, not
    dropped: ``ls -laR /etc`` is still a deny after ``-R`` becomes allowed,
    because the path is still outside the sandbox, and a blanket
    deny-to-allow rewrite would have taught the model the opposite.
    """
 
    func = "relabel_recursive_allow"
    relabelled: list[dict[str, str]] = []
    changed = 0
    for row in rows:
        stripped = strip_recursive_option(row["command"])
        if stripped is None:
            relabelled.append(dict(row))
            continue
        verdict, reason = classify(stripped)            # type: ignore[operator]
        if verdict != row["truth"]:
            changed += 1
        updated = dict(row)
        updated["truth"] = verdict
        updated["reason"] = reason
        relabelled.append(updated)
    _dbg("OC_CHAT.MAIN", func, f"{changed} of {len(rows)} verdicts moved", 3)
    return relabelled, changed
 
 
def gate_classifier(sandbox_root: Path) -> object:
    """A ``command -> (verdict, reason)`` callable backed by the real gate.
 
    Imported lazily: the argv gate lives in the ``oc_chat`` package, which the
    scratch harness venv cannot import, and only ``relabel`` needs it.
    """
 
    from oc_chat.tools._argv_policy import classify_command
 
    def classify(command: str) -> tuple[str, str]:
        """Classify one command against the unmodified gate."""
 
        decision = classify_command(command, sandbox_root)
        return decision.decision, decision.reason
 
    return classify
 
 
# --------------------------------------------------------------------------
# The result schema
# --------------------------------------------------------------------------
 
 
def result_row(
    case: dict[str, str], logits: list[float], temperature: float, elapsed_ms: float,
    decision: str = DEFAULT_DECISION, costs: list[list[float]] | None = None,
) -> dict[str, object]:
    """One result row in the plan-288 schema, calibrated and timed.
 
    *decision* is plan 298 P2's rule, defaulting to the argmax every result file
    written before that plan carries. Note what it does NOT touch: ``probs`` and
    ``logits`` are the model's own output either way, so two result files
    differing only in the rule differ only in ``pred``, and the Brier score and
    the reliability bands the scorer computes are identical across the pair.
    That is the property that lets the table read the rule's effect off the
    agreement and direction columns alone.
    """
 
    probs = probabilities(logits, temperature)
    return {
        "command": case["command"],
        "truth": case["truth"],
        "slice": case.get("slice", ""),
        "pred": decide(probs, decision, costs),
        "probs": probs,
        "ms": elapsed_ms,
        "reason": case.get("reason", ""),
        "logits": [round(value, 6) for value in logits],
    }
 
 
def write_result_file(
    path: Path,
    model: str,
    condition: str,
    rows: list[dict[str, object]],
    provenance: dict[str, object],
) -> None:
    """Write one result file the shared scorer can read."""
 
    func = "write_result_file"
    payload = {
        "model": model,
        "condition": condition,
        "results": rows,
        "provenance": provenance,
    }
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(json.dumps(payload, indent=1, sort_keys=True) + "\n", encoding="utf-8")
    _dbg("OC_CHAT.MAIN", func, f"wrote {Path(path).name} ({len(rows)} rows)", 2)
 
 
def file_digest(path: Path) -> str:
    """The SHA-256 of one file, which is how a trained run is identified."""
 
    digest = hashlib.sha256()
    with Path(path).open("rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            digest.update(block)
    return digest.hexdigest()
 
 
def weights_digest(model_dir: Path) -> str:
    """The SHA-256 of the saved weights, whichever format they were saved in."""
 
    for name in ("model.safetensors", "pytorch_model.bin"):
        candidate = Path(model_dir) / name
        if candidate.is_file():
            return file_digest(candidate)
    return ""
 
 
def verify_tokenizer_vocab(
    tokenizer_dir: Path, backbone: str, hf_home: Path
) -> tuple[str, bool]:
    """Check the tokenizer's vocabulary is the backbone's own, by digest.
 
    Returns the tokenizer's ``spm.model`` digest and whether a ``spm.model``
    found in the backbone's cache directory matches it. A mismatch means the
    tokenizer is the wrong one for the weights and training would be nonsense,
    so the caller refuses rather than warning.
    """
 
    func = "verify_tokenizer_vocab"
    tokenizer_spm = Path(tokenizer_dir) / "spm.model"
    digest = file_digest(tokenizer_spm) if tokenizer_spm.is_file() else ""
    folder = "models--" + backbone.replace("/", "--")
    candidates = sorted((Path(hf_home) / "hub" / folder).glob("snapshots/*/spm.model"))
    matched = any(file_digest(candidate) == digest for candidate in candidates) if digest else False
    _dbg("OC_CHAT.MAIN", func,
         f"{backbone}: {len(candidates)} spm candidate(s), match {matched}", 3)
    return digest, matched
 
 
# --------------------------------------------------------------------------
# The torch commands
# --------------------------------------------------------------------------
 
 
@dataclass(frozen=True)
class TrainConfig:
    """Everything one training run is pinned by."""
 
    backbone: str
    corpus_dir: Path
    out_dir: Path
    tokenizer_dir: Path = DEFAULT_TOKENIZER_DIR
    hf_home: Path = DEFAULT_HF_HOME
    seed: int = DEFAULT_SEED
    max_epochs: int = DEFAULT_MAX_EPOCHS
    batch: int = DEFAULT_BATCH
    learning_rate: float = DEFAULT_LEARNING_RATE
    warmup_share: float = DEFAULT_WARMUP_SHARE
    max_length: int = DEFAULT_MAX_LENGTH
    threads: int = DEFAULT_THREADS
    budget_seconds: float = DEFAULT_BUDGET_SECONDS
    device: str = DEFAULT_DEVICE
    prune_vocab: bool = False
    vocab_reserve: int = DEFAULT_VOCAB_RESERVE
    cache_ids: bool = False
    cache_root: Path = DEFAULT_CACHE_ROOT
    val_subset: int = 0
    patience: int = 1
    warmup_steps: int = 0
    alpha: float = DEFAULT_ALPHA
    cost_weight: float = DEFAULT_COST_WEIGHT
    stop_metric: str = DEFAULT_STOP_METRIC
 
 
def _require_torch() -> tuple[object, object, object]:
    """Import torch and transformers, or say which interpreter is needed."""
 
    try:
        import torch
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
    except ImportError as exc:                              # pragma: no cover - env guard
        raise RuntimeError(
            "encoder_classifier: train/calibrate/evaluate need torch and transformers; "
            "run this script with the scratch harness interpreter (see the module docstring)"
        ) from exc
    return torch, AutoTokenizer, AutoModelForSequenceClassification
 
 
def resolve_device(requested: str = DEFAULT_DEVICE) -> str:
    """The device a run will actually use, or a refusal saying why it cannot.
 
    ``auto`` is the default and reads the machine: CUDA when torch reports a
    usable device, the CPU otherwise. ``cuda`` is an ASSERTION rather than a
    preference -- if torch cannot see the card it raises, because the failure
    mode this guards against is a run that silently takes the CPU path and is
    read as a GPU measurement eighteen minutes later. ``cpu`` always succeeds
    and is how the CUDA venv reproduces a plan-290 CPU number.
    """
 
    func = "resolve_device"
    if requested not in DEVICE_CHOICES:
        raise ValueError(f"{func}: device must be one of {DEVICE_CHOICES}, not {requested!r}")
    if requested == "cpu":
        return "cpu"
    torch, _, _ = _require_torch()
    available = bool(torch.cuda.is_available())             # type: ignore[attr-defined]
    if requested == "cuda" and not available:
        raise RuntimeError(
            # OCT-LINT: disable=no_hardcoded_secrets reason="error-message prose; the entropy is the dotted torch identifiers, not a credential (operator ruling 2026-09-18)"
            "resolve_device: --device cuda was asked for and torch.cuda.is_available() "
            "is False; this interpreter is the CPU venv, or the driver does not carry "
            "sm_61 (see the module docstring on the CUDA 12.6 deadline)"
        )
    device = "cuda" if available else "cpu"
    _dbg("OC_CHAT.MAIN", func, f"{requested} -> {device}", 2)
    return device
 
 
def peak_device_memory(device: str) -> int:
    """Peak bytes torch allocated on *device* since the last reset, or ``0``.
 
    ``max_memory_allocated`` rather than ``nvidia-smi``: the card also carries
    the desktop and every browser on this host, so the only figure that means
    "what this training run needed" is torch's own allocator high-water mark.
    """
 
    if device != "cuda":
        return 0
    torch, _, _ = _require_torch()
    return int(torch.cuda.max_memory_allocated())           # type: ignore[attr-defined]
 
 
def describe_device(device: str) -> str:
    """A one-line description of *device* for a log line and a result file."""
 
    if device != "cuda":
        return f"CPU, {DEFAULT_THREADS} threads"
    torch, _, _ = _require_torch()
    name = torch.cuda.get_device_name(0)                    # type: ignore[attr-defined]
    return f"CUDA, {name}, torch {getattr(torch, '__version__', '?')}"
 
 
def _to_device(encoded: object, device: str) -> dict[str, object]:
    """Every tensor of a tokenizer's output, moved onto *device*."""
 
    return {key: value.to(device) for key, value in encoded.items()}   # type: ignore[attr-defined]
 
 
def remap_table(remap: dict[int, int], unknown_id: int, size: int, device: str) -> object:
    """:func:`remap_lookup` as a tensor on *device*, built once per run.
 
    A dict lookup per token would cost more than the forward pass it feeds, so
    plan 294 P2's map is materialised as a tensor indexed by the old id.
    """
 
    torch, _, _ = _require_torch()
    return torch.tensor(remap_lookup(remap, unknown_id, size),   # type: ignore[attr-defined]
                        dtype=torch.long, device=device)
 
 
def _remap_input_ids(input_ids: object, table: object) -> object:
    """Carry a batch of full-vocabulary ids through a :func:`remap_table`."""
 
    return table[input_ids]                                 # type: ignore[index]
 
 
def cost_tensor(costs: list[list[float]], device: str) -> object:
    """A :func:`cost_matrix` as a float tensor on *device*, built once per run."""
 
    torch, _, _ = _require_torch()
    return torch.tensor(costs, dtype=torch.float32, device=device)   # type: ignore[attr-defined]
 
 
def expected_cost_term(logits: object, targets: object, table: object) -> object:
    """Plan 298 P3's training term on tensors: the batch's mean expected cost.
 
    ``softmax(logits)`` is the model's own forecast, so the term is
    differentiable in the logits and pushes probability mass off the verdicts
    the matrix prices highly -- not merely off the wrong ones, which is what
    cross-entropy alone does. ``table[targets]`` selects each row's true-verdict
    cost row, and that gather is the whole of the indexing.
 
    :func:`expected_cost` is the definition this reproduces; the witness in
    ``tests/test_298_direction.py`` asserts the two agree on a toy batch.
    """
 
    torch, _, _ = _require_torch()
    probs = torch.softmax(logits.float(), dim=-1)           # type: ignore[attr-defined]
    return (probs * table[targets]).sum(dim=-1).mean()      # type: ignore[index]
 
 
def slice_embeddings(model: object, kept: list[int], pad_row: int) -> int:
    """Replace the model's word embeddings with the *kept* rows only.
 
    Returns how many rows the table held before. The classification head is a
    pooler and a linear layer, so nothing is tied to this table and slicing it
    is the whole of the change; ``config.vocab_size`` moves with it so that a
    ``save_pretrained`` round trip rebuilds the same shape.
    """
 
    func = "slice_embeddings"
    torch, _, _ = _require_torch()
    embeddings = model.get_input_embeddings()               # type: ignore[attr-defined]
    weight = embeddings.weight.data
    before = int(weight.shape[0])
    index = torch.tensor(kept, dtype=torch.long, device=weight.device)   # type: ignore[attr-defined]
    sliced = torch.nn.Embedding(len(kept), int(weight.shape[1]), padding_idx=pad_row)  # type: ignore[attr-defined]
    sliced.weight.data = weight.index_select(0, index).clone()
    model.set_input_embeddings(sliced)                      # type: ignore[attr-defined]
    model.config.vocab_size = len(kept)                     # type: ignore[attr-defined]
    _dbg("OC_CHAT.MAIN", func, f"{before} rows sliced to {len(kept)}", 2)
    return before
 
 
def _collate(rows: list[list[int]], pad_id: int, extra_keys: tuple[str, ...],
             device: str) -> dict[str, object]:
    """Pad cached ids into the tensors the model's forward pass expects.
 
    *extra_keys* mirrors whatever the tokenizer itself emits beside the ids and
    the mask -- ``token_type_ids`` for this tokenizer -- as zeros, so a cached
    batch and a freshly tokenised one are the same call.
    """
 
    torch, _, _ = _require_torch()
    width = max(len(row) for row in rows)
    input_ids = torch.full((len(rows), width), pad_id, dtype=torch.long)   # type: ignore[attr-defined]
    mask = torch.zeros((len(rows), width), dtype=torch.long)               # type: ignore[attr-defined]
    for position, row in enumerate(rows):
        input_ids[position, :len(row)] = torch.tensor(row, dtype=torch.long)  # type: ignore[attr-defined]
        mask[position, :len(row)] = 1
    encoded = {"input_ids": input_ids, "attention_mask": mask}
    for key in extra_keys:
        encoded[key] = torch.zeros((len(rows), width), dtype=torch.long)   # type: ignore[attr-defined]
    return _to_device(encoded, device)
 
 
def _saved_remap_table(model_dir: Path, device: str) -> object | None:
    """The lookup table for a model trained with a pruned vocabulary, or ``None``.
 
    ``train`` writes :data:`VOCAB_PRUNE_FILE` beside the weights when and only
    when it pruned, so its presence is what tells ``calibrate`` and ``evaluate``
    that the ids a tokenizer hands them are not the ids this model's embedding
    table is indexed by. A model trained without pruning has no such file and
    every id goes through untouched.
    """
 
    func = "_saved_remap_table"
    path = Path(model_dir) / VOCAB_PRUNE_FILE
    if not path.is_file():
        return None
    saved = json.loads(path.read_text(encoding="utf-8"))
    remap = {int(old): int(new) for old, new in saved["remap"].items()}
    _dbg("OC_CHAT.MAIN", func,
         f"{len(remap)} kept of {saved['full_vocab']} vocabulary rows", 3)
    return remap_table(remap, int(saved["unknown_id"]), int(saved["full_vocab"]), device)
 
 
def _load_tokenizer(tokenizer_dir: Path) -> object:
    """Load the converted DeBERTa-v3 fast tokenizer from a local directory."""
 
    _, auto_tokenizer, _ = _require_torch()
    return auto_tokenizer.from_pretrained(str(tokenizer_dir))    # type: ignore[attr-defined]
 
 
def _forward_logits(
    model: object,
    tokenizer: object,
    commands: list[str],
    max_length: int,
    batch: int,
    device: str = "cpu",
    table: object | None = None,
) -> list[list[float]]:
    """Logits for every command, in batches, in evaluation mode.
 
    *table* is plan 294 P2's pruned-vocabulary lookup: when the model's
    embedding table has been sliced, the ids the tokenizer produces are the FULL
    vocabulary's and have to be carried through the map before they index the
    sliced table, or every id past the kept count is out of range.
    """
 
    torch, _, _ = _require_torch()
    rows: list[list[float]] = []
    with torch.no_grad():                                   # type: ignore[attr-defined]
        for start in range(0, len(commands), batch):
            chunk = commands[start:start + batch]
            encoded = tokenizer(chunk, padding=True, truncation=True,    # type: ignore[operator]
                                max_length=max_length, return_tensors="pt")
            encoded = _to_device(encoded, device)
            if table is not None:
                encoded["input_ids"] = _remap_input_ids(encoded["input_ids"], table)
            output = model(**encoded)                       # type: ignore[operator]
            rows.extend(output.logits.float().tolist())
    return rows
 
 
def score_full_validation(
    model_dir: Path, rows: list[dict[str, str]], tokenizer: object,
    max_length: int, batch: int, device: str, table: object | None,
) -> float:
    """The WHOLE validation split's cross-entropy at the saved best checkpoint.
 
    Plan 294 P3 validates each epoch on a 1,000-row subset, which is a
    different number from plan 290's 3,000-row one and cannot be compared with
    it. The best checkpoint is therefore reloaded once at the end and scored on
    the full split, so the run still produces the figure the previous plan's
    log can be read against.
    """
 
    func = "score_full_validation"
    _, _, auto_model = _require_torch()
    model = auto_model.from_pretrained(str(model_dir))      # type: ignore[attr-defined]
    model.to(device)
    model.eval()
    logits = _forward_logits(model, tokenizer, [row["command"] for row in rows],
                             max_length, batch, device, table)
    indexes = [LABELS.index(row["truth"]) for row in rows]
    loss = mean_negative_log_likelihood(logits, indexes, 1.0)
    _dbg("OC_CHAT.MAIN", func, f"full validation loss {loss:.4f} over {len(rows)} rows", 3)
    return loss
 
 
def train(config: TrainConfig) -> dict[str, object]:
    """Fine-tune one backbone on the corpus, early-stopped on a watched metric.
 
    Writes the best epoch's weights, the tokenizer and a training record into
    ``config.out_dir``, and returns the record. The stopping rule is plan 290
    P2's -- validation cross-entropy stops improving -- plus the budget guard
    P2's one-hour wall-time limit implies: when another epoch would not fit,
    the run stops at the end of the current one and the record says so, which
    is the opposite of quietly overrunning the budget.
 
    Plan 298 P3 adds ``config.cost_weight`` times :func:`expected_cost_term` to
    the per-batch loss. Two deliberate limits on it. It is off at ``0.0``, so an
    unflagged run is still plan 295's run. And the STOPPING rule was left on
    plain validation cross-entropy even when the term was on: the recipe was
    being held constant so that the objective was the only variable, and moving
    the early-stopping criterion at the same time would have changed two things
    and measured their sum.
 
    Plan 300 P1 supplies the second thing, as its own flag.
    ``config.stop_metric`` chooses which validation figure early stopping and
    checkpoint selection watch: ``"ce"`` is the rule above, unchanged and still
    the default, and ``"cost"`` watches the subset validation expected cost
    under the same ``config.alpha`` matrix -- the column a direction-weighted
    objective is trying to move, and the one plan 298 measured a run being
    checkpointed against. Both figures are computed every epoch whichever one
    stops the run, the epoch log carries both, and the record says which epoch
    the other metric would have kept. Nothing else moves: the full-split
    cross-entropy is still scored on the reloaded best checkpoint, and the
    temperature is still fitted by NLL afterwards, so the probabilities a
    reading consumes are calibrated the way every earlier tag's were.
    """
 
    func = "train"
    torch, _, auto_model = _require_torch()
    torch.set_num_threads(config.threads)                   # type: ignore[attr-defined]
    device = resolve_device(config.device)
    digest, matched = verify_tokenizer_vocab(
        config.tokenizer_dir, config.backbone, config.hf_home
    )
    if not matched:
        raise RuntimeError(
            f"train: tokenizer at {config.tokenizer_dir} does not carry the "
            f"vocabulary of {config.backbone} (spm.model digest {digest[:16]})"
        )
    held_out = held_out_commands(config.corpus_dir)
    training = load_split(config.corpus_dir, "corpus.jsonl", held_out)
    validation = load_split(config.corpus_dir, "corpus_val.jsonl", held_out)
    print(f"train: {config.backbone} on {len(training)} rows, "
          f"{len(validation)} validation, batch {config.batch}, "
          f"lr {config.learning_rate}, up to {config.max_epochs} epochs, "
          f"{config.threads} threads", flush=True)
    print(f"train: device {device} ({ascii(describe_device(device))})", flush=True)
    print(f"train: label counts {label_counts(training)}", flush=True)
 
    random.seed(config.seed)
    torch.manual_seed(config.seed)                          # type: ignore[attr-defined]
    if device == "cuda":
        torch.cuda.manual_seed_all(config.seed)             # type: ignore[attr-defined]
    tokenizer = _load_tokenizer(config.tokenizer_dir)
    model = auto_model.from_pretrained(                     # type: ignore[attr-defined]
        config.backbone, num_labels=len(LABELS),
        id2label={index: label for index, label in enumerate(LABELS)},
        label2id={label: index for index, label in enumerate(LABELS)},
    )
    Path(config.out_dir).mkdir(parents=True, exist_ok=True)
    commands = [row["command"] for row in training] + [row["command"] for row in validation]
    max_length = resolve_max_length(commands, tokenizer, config.max_length)
    if max_length != config.max_length:
        print(f"train: sequence length {config.max_length} would truncate; "
              f"using {max_length}", flush=True)
 
    full_vocab = int(model.get_input_embeddings().weight.shape[0])
    kept: list[int] = []
    table: object | None = None
    if config.prune_vocab:
        kept, remap = prune_vocab(commands, tokenizer, config.vocab_reserve, max_length)
        unknown_row = remap[int(tokenizer.unk_token_id)]     # type: ignore[attr-defined]
        pad_row = remap[int(tokenizer.pad_token_id)]         # type: ignore[attr-defined]
        before = slice_embeddings(model, kept, pad_row)
        model.config.pad_token_id = pad_row                 # type: ignore[attr-defined]
        table = remap_table(remap, unknown_row, full_vocab, device)
        Path(config.out_dir).joinpath(VOCAB_PRUNE_FILE).write_text(
            json.dumps({"kept": kept, "remap": {str(old): new for old, new in remap.items()},
                        "unknown_id": unknown_row, "pad_id": pad_row,
                        "full_vocab": full_vocab, "reserve": config.vocab_reserve},
                       sort_keys=True),
            encoding="utf-8",
        )
        print(f"train: vocabulary pruned {before} -> {len(kept)} rows "
              f"(reserve {config.vocab_reserve})", flush=True)
 
    model.to(device)
    if device == "cuda":
        torch.cuda.reset_peak_memory_stats()                # type: ignore[attr-defined]
    parameters = sum(parameter.numel() for parameter in model.parameters())
    optimiser = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)  # type: ignore[attr-defined]
    steps_per_epoch = math.ceil(len(training) / config.batch)
    warmup_steps = config.warmup_steps or max(
        1, int(config.warmup_share * steps_per_epoch * config.max_epochs)
    )
    print(f"train: {parameters / 1e6:.1f}M parameters, {steps_per_epoch} steps per epoch, "
          f"{warmup_steps} warm-up steps", flush=True)
    _dbg("OC_CHAT.MAIN", func, f"{config.backbone}: {parameters} parameters", 2)
 
    costs = cost_matrix(config.alpha)
    cost_weights = cost_tensor(costs, device) if config.cost_weight else None
    if cost_weights is None:
        print("train: objective cross-entropy (plan 295's; --cost-weight is 0)", flush=True)
    else:
        print(f"train: objective cross-entropy + {config.cost_weight} * expected cost "
              f"at alpha {config.alpha}; granting more permission costs "
              f"{config.alpha}x granting less", flush=True)
    if config.stop_metric not in STOP_METRIC_CHOICES:
        raise ValueError(
            f"{func}: stop metric must be one of {STOP_METRIC_CHOICES}, "
            f"not {config.stop_metric!r}"
        )
    if config.stop_metric == "ce":
        print("train: stopping on validation cross-entropy (plan 290's rule; "
              "--stop-metric is ce)", flush=True)
    else:
        print(f"train: stopping on validation expected cost at alpha {config.alpha} "
              f"(plan 300 P1); both figures are logged every epoch", flush=True)
 
    pad_id = int(tokenizer.pad_token_id)                    # type: ignore[attr-defined]
    extra_keys = tuple(key for key in tokenise_keys(tokenizer)
                       if key not in ("input_ids", "attention_mask"))
    cached_ids: dict[str, list[int]] | None = None
    if config.cache_ids:
        key = cache_key(Path(config.corpus_dir) / "corpus.jsonl", digest, max_length)
        rows_ids = load_cached_ids(config.cache_root, key)
        origin = "cache"
        if rows_ids is None:
            origin = "tokenizer"
            rows_ids = tokenise([row["command"] for row in training], tokenizer, max_length)
            store_cached_ids(config.cache_root, key, rows_ids)
        cached_ids = {row["command"]: ids for row, ids in zip(training, rows_ids)}
        print(f"train: {len(rows_ids)} training rows tokenised once from the {origin} "
              f"(key {key[:12]})", flush=True)
 
    subset = validation_subset(validation, config.val_subset)
    if len(subset) != len(validation):
        print(f"train: validating each epoch on {len(subset)} of {len(validation)} rows "
              f"(seed {DEFAULT_SUBSET_SEED}); the full split is scored at the end",
              flush=True)
 
    log_lines: list[str] = []
    epoch_records: list[dict[str, float | int]] = []
    rng = random.Random(config.seed)
    started = time.perf_counter()
    step = 0
    best_value = float("inf")
    best_loss = float("inf")
    best_cost = float("inf")
    best_epoch = 0
    stale = 0
    full_validation_loss = float("nan")
    stop_reason = "max-epochs"
 
    def write_record(reason: str) -> dict[str, object]:
        """Write the training record and log as they stand, and return it.
 
        Called at every improving checkpoint as well as at the end, because a
        run on this host can die mid-epoch: the first plan-290 run lost a
        54-minute record to a CPU allocation failure in epoch 3 while the
        epoch-2 weights sat complete on disk beside it. A record written with
        the weights it describes is the difference between a usable checkpoint
        and an unidentifiable one.
        """
 
        record: dict[str, object] = {
            "backbone": config.backbone,
            "seed": config.seed,
            "epochs": best_epoch,
            "epochs_run": len(epoch_records),
            "max_epochs": config.max_epochs,
            "stop_reason": reason,
            "train_seconds": time.perf_counter() - started,
            "train_rows": len(training),
            "validation_rows": len(validation),
            "batch": config.batch,
            "learning_rate": config.learning_rate,
            "warmup_steps": warmup_steps,
            "max_length": max_length,
            "threads": config.threads,
            "device": device,
            "device_name": describe_device(device),
            "peak_device_bytes": peak_device_memory(device),
            "patience": config.patience,
            "alpha": config.alpha,
            "cost_weight": config.cost_weight,
            "stop_metric": config.stop_metric,
            "alternate_stop_metric": other_stop_metric(config.stop_metric),
            "alternate_best_epoch": select_best_epoch(
                [(float(entry["validation_loss"]), float(entry["validation_cost"]))
                 for entry in epoch_records],
                other_stop_metric(config.stop_metric), config.patience,
            )["best_epoch"],
            "vocab_full": full_vocab,
            "vocab_kept": len(kept) if kept else full_vocab,
            "cached_ids": bool(config.cache_ids),
            "validation_subset": len(subset),
            "full_validation_loss": full_validation_loss,
            "parameters": parameters,
            "budget_seconds": config.budget_seconds,
            "corpus_dir": str(config.corpus_dir),
            "tokenizer_dir": str(config.tokenizer_dir),
            "tokenizer_spm_sha256": digest,
            "epoch_log": epoch_records,
            "best_validation_loss": best_loss,
            "best_validation_cost": best_cost,
            "weights_sha256": weights_digest(config.out_dir),
            "torch": getattr(torch, "__version__", "?"),
        }
        Path(config.out_dir).joinpath("training.json").write_text(
            json.dumps(record, indent=1, sort_keys=True) + "\n", encoding="utf-8"
        )
        header = (f"{config.backbone} seed {config.seed} batch {config.batch} "
                  f"lr {config.learning_rate} threads {config.threads} "
                  f"device {device} {parameters} parameters")
        Path(config.out_dir).joinpath("train.log").write_text(
            header + "\n" + "\n".join(log_lines)
            + f"\nstopped after epoch {len(epoch_records)} ({reason}); "
              f"best epoch {best_epoch}; "
              f"{record['train_seconds']:.0f}s total\n",           # type: ignore[str-format]
            encoding="utf-8",
        )
        return record
 
    try:
        for epoch in range(1, config.max_epochs + 1):
            model.train()
            epoch_started = time.perf_counter()
            running = 0.0
            for position, batch_rows in enumerate(epoch_batches(training, config.batch, rng), 1):
                if cached_ids is None:
                    encoded = tokenizer([row["command"] for row in batch_rows],   # type: ignore[operator]
                                        padding=True, truncation=True,
                                        max_length=max_length, return_tensors="pt")
                    encoded = _to_device(encoded, device)
                else:
                    encoded = _collate([cached_ids[row["command"]] for row in batch_rows],
                                       pad_id, extra_keys, device)
                if table is not None:
                    encoded["input_ids"] = _remap_input_ids(encoded["input_ids"], table)
                targets = torch.tensor(                        # type: ignore[attr-defined]
                    [LABELS.index(row["truth"]) for row in batch_rows], dtype=torch.long,  # type: ignore[attr-defined]
                    device=device,
                )
                step += 1
                for group in optimiser.param_groups:
                    group["lr"] = config.learning_rate * min(1.0, step / warmup_steps)
                output = model(**encoded, labels=targets)
                loss = output.loss
                if cost_weights is not None:
                    loss = loss + config.cost_weight * expected_cost_term(   # type: ignore[operator]
                        output.logits, targets, cost_weights
                    )
                loss.backward()
                optimiser.step()
                optimiser.zero_grad(set_to_none=True)
                running += float(loss.item())
                if position % 50 == 0:
                    elapsed = time.perf_counter() - epoch_started
                    print(f"  epoch {epoch} step {position}/{steps_per_epoch} "
                          f"loss {running / position:.4f} "
                          f"{elapsed:.0f}s elapsed, {elapsed / position * steps_per_epoch:.0f}s projected",
                          flush=True)
            train_loss = running / steps_per_epoch
            model.eval()
            logits = _forward_logits(model, tokenizer, [row["command"] for row in subset],
                                     max_length, config.batch, device, table)
            indexes = [LABELS.index(row["truth"]) for row in subset]
            validation_loss = mean_negative_log_likelihood(logits, indexes, 1.0)
            forecasts = [probabilities(row) for row in logits]
            hits = sum(max(probs, key=lambda label: probs[label])
                       == subset[position]["truth"] for position, probs in enumerate(forecasts))
            validation_cost = expected_cost(
                [[probs[label] for label in LABELS] for probs in forecasts], indexes, costs
            )
            epoch_seconds = time.perf_counter() - epoch_started
            line = (f"epoch {epoch}: train loss {train_loss:.4f} "
                    f"validation loss {validation_loss:.4f} "
                    f"validation cost {validation_cost:.4f} "
                    f"validation accuracy {hits / len(subset):.4f} "
                    f"{epoch_seconds:.0f}s")
            print(f"train: {line}", flush=True)
            log_lines.append(line)
            epoch_records.append({
                "epoch": epoch, "train_loss": train_loss, "validation_loss": validation_loss,
                "validation_cost": validation_cost,
                "validation_accuracy": hits / len(subset), "seconds": epoch_seconds,
            })
            _dbg("OC_CHAT.MAIN", func, line, 3)
 
            watched = stopping_value(epoch_records[-1], config.stop_metric)
            improved = watched < best_value - EARLY_STOP_MIN_DELTA
            if improved:
                stale = 0
                best_value, best_epoch = watched, epoch
                best_loss, best_cost = validation_loss, validation_cost
                model.save_pretrained(str(config.out_dir))
                tokenizer.save_pretrained(str(config.out_dir))   # type: ignore[operator]
                log_lines.append(f"epoch {epoch}: checkpoint saved (best so far)")
                write_record("in-progress")
            else:
                stale += 1
                log_lines.append(f"epoch {epoch}: no improvement in {config.stop_metric} "
                                 f"({stale} of {config.patience} allowed)")
                if stale >= config.patience:
                    stop_reason = "early-stop"
                    log_lines.append(f"epoch {epoch}: stopping; best was epoch {best_epoch} "
                                     f"by {config.stop_metric}")
                    break
            elapsed = time.perf_counter() - started
            if epoch < config.max_epochs and elapsed + epoch_seconds > config.budget_seconds:
                stop_reason = "budget"
                log_lines.append(
                    f"epoch {epoch}: {elapsed:.0f}s elapsed and another epoch needs "
                    f"about {epoch_seconds:.0f}s, which would pass the "
                    f"{config.budget_seconds:.0f}s budget; stopping here"
                )
                break
 
    except RuntimeError as exc:
        allocator = any(mark in str(exc) for mark in
                        ("not enough memory", "DefaultCPUAllocator", "CUDA out of memory"))
        stop_reason = "memory" if allocator else "error"
        line = (f"epoch {len(epoch_records) + 1}: aborted after "
                f"{time.perf_counter() - started:.0f}s: {exc}")
        print(f"train: {line}", flush=True)
        log_lines.append(line)
        _dbg("OC_CHAT.MAIN", func, line, 2)
        if best_epoch == 0:
            write_record(stop_reason)
            raise
 
    if best_epoch and len(subset) != len(validation):
        full_validation_loss = score_full_validation(
            config.out_dir, validation, tokenizer, max_length, config.batch, device, table
        )
        line = (f"full validation loss over {len(validation)} rows at the best "
                f"checkpoint: {full_validation_loss:.4f}")
        print(f"train: {line}", flush=True)
        log_lines.append(line)
 
    record = write_record(stop_reason)
    train_seconds = float(record["train_seconds"])          # type: ignore[arg-type]
    print(f"train: best epoch {best_epoch} by {config.stop_metric} ({stop_reason}) "
          f"in {train_seconds:.0f}s -> {config.out_dir}", flush=True)
    print(f"train: validation cross-entropy {best_loss:.4f} and expected cost "
          f"{best_cost:.4f} at that epoch; watching "
          f"{record['alternate_stop_metric']} would have kept epoch "
          f"{record['alternate_best_epoch']}", flush=True)
    _dbg("OC_CHAT.MAIN", func, f"trained in {train_seconds:.0f}s, best epoch {best_epoch}", 2)
    return record
 
 
def calibrate(
    model_dir: Path, corpus_dir: Path, batch: int = DEFAULT_BATCH,
    max_length: int = DEFAULT_MAX_LENGTH, threads: int = DEFAULT_THREADS,
    device: str = DEFAULT_DEVICE,
) -> dict[str, float | int]:
    """Fit one temperature on the validation split and record what it moved.
 
    Plan 294 P3 validates on a 1,000-row subset DURING training, for speed; the
    temperature is still fitted on the whole 3,000-row split here, because a
    calibration scalar fitted on a third of the rows is a third as well
    determined and costs one forward pass to get right.
    """
 
    func = "calibrate"
    torch, _, auto_model = _require_torch()
    torch.set_num_threads(threads)                          # type: ignore[attr-defined]
    device = resolve_device(device)
    held_out = held_out_commands(corpus_dir)
    validation = load_split(corpus_dir, "corpus_val.jsonl", held_out)
    tokenizer = _load_tokenizer(model_dir)
    model = auto_model.from_pretrained(str(model_dir))      # type: ignore[attr-defined]
    model.to(device)
    model.eval()
    table = _saved_remap_table(model_dir, device)
    started = time.perf_counter()
    logits = _forward_logits(model, tokenizer, [row["command"] for row in validation],
                             max_length, batch, device, table)
    labels = [LABELS.index(row["truth"]) for row in validation]
    temperature = fit_temperature(logits, labels)
    report = calibration_report(logits, validation, temperature)
    report["forward_seconds"] = time.perf_counter() - started
    Path(model_dir).joinpath("calibration.json").write_text(
        json.dumps(report, indent=1, sort_keys=True) + "\n", encoding="utf-8"
    )
    Path(model_dir).joinpath("val_logits.json").write_text(
        json.dumps({"commands": [row["command"] for row in validation],
                    "truth": [row["truth"] for row in validation],
                    "logits": [[round(value, 6) for value in row] for row in logits]}),
        encoding="utf-8",
    )
    print(f"calibrate: temperature {temperature:.4f}; "
          f"validation NLL {report['nll_before']:.4f} -> {report['nll_after']:.4f}; "
          f"Brier {report['brier_before']:.4f} -> {report['brier_after']:.4f}; "
          f"accuracy {report['validation_accuracy']:.4f}")
    _dbg("OC_CHAT.MAIN", func, f"temperature {temperature:.4f}", 2)
    return report
 
 
def evaluate(
    model_dir: Path, cases_path: Path, condition: str, out_path: Path,
    model_name: str, batch: int = 1, max_length: int = DEFAULT_MAX_LENGTH,
    threads: int = DEFAULT_THREADS, device: str = DEFAULT_DEVICE,
    decision: str = DEFAULT_DECISION, alpha: float = DEFAULT_ALPHA,
) -> dict[str, object]:
    """Predict a case set and write its result file, timing every decision.
 
    At ``batch=1`` the clock runs around tokenise-plus-forward for one command,
    which is the per-decision latency a gate would pay. At ``batch>1`` the
    clock runs around the whole batch and the per-row figure is the amortised
    one, which is the throughput reading rather than the latency reading.
 
    *decision* and *alpha* are plan 298 P2's rule. They are a property of the
    READING and not of the model, so both go into the provenance block: two
    result files from one checkpoint under two rules are otherwise
    indistinguishable except by their verdicts, and the scorer reports no field
    that would say which rule produced which.
 
    Plan 300 P2 adds :func:`confident_error_counts` to that same block. It goes
    in the PROVENANCE and not into the rows on purpose: ``score.py`` reads the
    row shape and a new per-row field would change what every consumer of a
    result file sees, whereas a provenance key is additive and the scorer never
    looks at it. The counts are recomputable from the rows by anyone who wants
    them at another threshold, which is why the file carries the summary rather
    than a flag per row.
    """
 
    func = "evaluate"
    torch, _, auto_model = _require_torch()
    torch.set_num_threads(threads)                          # type: ignore[attr-defined]
    device = resolve_device(device)
    training_record = json.loads(
        Path(model_dir).joinpath("training.json").read_text(encoding="utf-8")
    )
    calibration = json.loads(
        Path(model_dir).joinpath("calibration.json").read_text(encoding="utf-8")
    )
    temperature = float(calibration["temperature"])
    cases = json.loads(Path(cases_path).read_text(encoding="utf-8"))
    tokenizer = _load_tokenizer(model_dir)
    model = auto_model.from_pretrained(str(model_dir))     # type: ignore[attr-defined]
    model.to(device)
    model.eval()
    table = _saved_remap_table(model_dir, device)
 
    costs = cost_matrix(alpha)
    rows: list[dict[str, object]] = []
    with torch.no_grad():                                   # type: ignore[attr-defined]
        for start in range(0, len(cases), batch):
            chunk = cases[start:start + batch]
            clock = time.perf_counter()
            encoded = tokenizer([case["command"] for case in chunk],      # type: ignore[operator]
                                padding=True, truncation=True,
                                max_length=max_length, return_tensors="pt")
            encoded = _to_device(encoded, device)
            if table is not None:
                encoded["input_ids"] = _remap_input_ids(encoded["input_ids"], table)
            logits = model(**encoded).logits.float()       # type: ignore[operator]
            if device == "cuda":
                torch.cuda.synchronize()                    # type: ignore[attr-defined]
            elapsed_ms = (time.perf_counter() - clock) * 1000.0 / len(chunk)
            for case, row in zip(chunk, logits.tolist()):
                rows.append(result_row(case, row, temperature, elapsed_ms,
                                       decision, costs))
    hits = sum(row["pred"] == row["truth"] for row in rows)
    confident = confident_error_counts(rows)
    provenance = {
        "date": time.strftime("%Y-%m-%d"),
        "route": (f"{describe_device(device)}, fine-tuned encoder, batch {batch}"
                  if device == "cuda"
                  else f"CPU, {threads} threads, fine-tuned encoder, batch {batch}"),
        "case_file": Path(cases_path).name,
        "model_dir": str(model_dir),
        "batch": batch,
        "max_length": max_length,
        "threads": threads,
        "device": device,
        "temperature": temperature,
        "decision": decision,
        "alpha": alpha,
        "confident_errors": confident,
        "validation_brier_before": calibration["brier_before"],
        "validation_brier_after": calibration["brier_after"],
        "validation_accuracy": calibration["validation_accuracy"],
    }
    for key in ("backbone", "seed", "epochs", "epochs_run", "stop_reason", "train_seconds",
                "train_rows", "learning_rate", "parameters", "weights_sha256",
                "tokenizer_spm_sha256", "corpus_dir", "torch"):
        provenance[key] = training_record[key]
    for key in ("device", "vocab_kept", "vocab_full", "cached_ids", "validation_subset",
                "alpha", "cost_weight", "stop_metric", "alternate_stop_metric",
                "alternate_best_epoch", "best_validation_loss", "best_validation_cost"):
        if key in training_record:
            provenance[f"train_{key}"] = training_record[key]
    write_result_file(out_path, model_name, condition, rows, provenance)
    print(f"evaluate: {condition:5} {hits}/{len(rows)} agreement "
          f"({decision}, alpha {alpha}) -> {Path(out_path).name}")
    print("evaluate: {} permission-granting error(s), {}".format(
        confident["granted_more"],
        ", ".join(f"{confident[f'granted_more_at_{threshold:g}']} at {threshold:g} or better"
                  for threshold in CONFIDENT_THRESHOLDS)))
    _dbg("OC_CHAT.MAIN", func, f"{condition}: {hits}/{len(rows)} agreement", 3)
    return {"agreement": hits / len(rows), "rows": len(rows)}
 
 
def relabel(bench_dir: Path, out_dir: Path) -> dict[str, object]:
    """Write the policy-change probe's relabelled copy of the corpus.
 
    The copy goes under *out_dir* -- never over a committed file -- and is a
    self-consistent bench corpus under the changed policy: both corpus splits
    and both case sets, each under its canonical name, so ``train
    --corpus-dir <out_dir>`` finds the held-out guard's case files and
    ``evaluate --cases <out_dir>/cases.json`` measures against the labels the
    CHANGED policy implies rather than the committed ones. A report beside
    them records the filter, the counts and the verdict movement.
    """
 
    func = "relabel"
    classify = gate_classifier(Path(bench_dir).parent.parent)
    Path(out_dir).mkdir(parents=True, exist_ok=True)
    held_out = held_out_commands(bench_dir)
    report: dict[str, object] = {
        "filter": ("listing command (ls/dir, basename lowercased, .exe stripped) "
                   "carrying an unquoted short-option cluster containing R; the new "
                   "label is the unmodified gate's verdict for the same command with "
                   "the R removed"),
        "bench_dir": str(bench_dir),
        "splits": {},
    }
    for name in ("corpus.jsonl", "corpus_val.jsonl"):
        rows = load_split(bench_dir, name, held_out)
        matched = sum(1 for row in rows if strip_recursive_option(row["command"]) is not None)
        relabelled, changed = relabel_recursive_allow(rows, classify)
        Path(out_dir).joinpath(name).write_text(
            "".join(json.dumps(row, sort_keys=True) + "\n" for row in relabelled),
            encoding="utf-8",
        )
        report["splits"][name] = {          # type: ignore[index]
            "rows": len(rows), "matched": matched, "verdicts_moved": changed,
            "labels_before": label_counts(rows), "labels_after": label_counts(relabelled),
        }
        print(f"relabel: {name}: {matched} matched, {changed} verdicts moved, "
              f"labels {label_counts(relabelled)}", flush=True)
    report["cases"] = {}
    for name in ("cases.json", "cases_hard.json"):
        cases = json.loads(Path(bench_dir).joinpath(name).read_text(encoding="utf-8"))
        relabelled_cases, changed_cases = relabel_recursive_allow(cases, classify)
        target = Path(out_dir) / name
        target.write_text(json.dumps(relabelled_cases, indent=1) + "\n", encoding="utf-8")
        moved = [case["command"] for before, case in zip(cases, relabelled_cases)
                 if before["truth"] != case["truth"]]
        report["cases"][name] = {         # type: ignore[index]
            "rows": len(cases),
            "matched": sum(1 for case in cases
                           if strip_recursive_option(case["command"]) is not None),
            "verdicts_moved": changed_cases,
            "moved_commands": moved,
        }
        print(f"relabel: {name}: {len(moved)} of {len(cases)} verdicts moved "
              f"-> {target.name}", flush=True)
    Path(out_dir).joinpath("relabel_report.json").write_text(
        json.dumps(report, indent=1, sort_keys=True) + "\n", encoding="utf-8"
    )
    _dbg("OC_CHAT.MAIN", func, f"relabelled corpus written to {out_dir}", 2)
    return report
 
 
# --------------------------------------------------------------------------
# The CLI
# --------------------------------------------------------------------------
 
 
def _bench_dir() -> Path:
    """``bench/argv_policy`` -- the directory holding the corpus and the cases."""
 
    return Path(__file__).resolve().parent.parent
 
 
def _build_parser() -> argparse.ArgumentParser:
    """Construct the CLI argument parser."""
 
    parser = argparse.ArgumentParser(
        prog="python bench/argv_policy/baselines/encoder_classifier.py",
        description="Fine-tune, calibrate and score a small encoder on the argv-policy bench.",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)
 
    trainer = subparsers.add_parser("train", help="Fine-tune one backbone.")
    trainer.add_argument("--backbone", required=True, help="Hugging Face model id.")
    trainer.add_argument("--out", required=True, help="Model directory to write.")
    trainer.add_argument("--corpus-dir", default=None,
                         help="Where corpus.jsonl lives (default: the bench directory).")
    trainer.add_argument("--tokenizer-dir", default=str(DEFAULT_TOKENIZER_DIR),
                         help="Directory holding the converted fast tokenizer.")
    trainer.add_argument("--seed", type=int, default=DEFAULT_SEED)
    trainer.add_argument("--epochs", type=int, default=DEFAULT_MAX_EPOCHS,
                         help=f"Maximum epochs (default: {DEFAULT_MAX_EPOCHS}).")
    trainer.add_argument("--batch", type=int, default=DEFAULT_BATCH)
    trainer.add_argument("--learning-rate", type=float, default=DEFAULT_LEARNING_RATE)
    trainer.add_argument("--max-length", type=int, default=DEFAULT_MAX_LENGTH)
    trainer.add_argument("--threads", type=int, default=DEFAULT_THREADS)
    trainer.add_argument("--budget-seconds", type=float, default=DEFAULT_BUDGET_SECONDS)
    trainer.add_argument("--device", choices=DEVICE_CHOICES, default=DEFAULT_DEVICE,
                         help="cuda refuses rather than falling back to the CPU.")
    trainer.add_argument("--prune-vocab", action="store_true",
                         help="Plan 294 P2: slice the embedding table to the ids in use.")
    trainer.add_argument("--vocab-reserve", type=int, default=DEFAULT_VOCAB_RESERVE,
                         help="Unseen vocabulary rows kept anyway (P2).")
    trainer.add_argument("--cache-ids", action="store_true",
                         help="Plan 294 P3: tokenise the training split once, then reuse it.")
    trainer.add_argument("--cache-root", default=str(DEFAULT_CACHE_ROOT),
                         help="Where the tokenised split is cached (never the worktree).")
    trainer.add_argument("--val-subset", type=int, default=0,
                         help="Plan 294 P3: rows to validate on each epoch; 0 is the full split.")
    trainer.add_argument("--patience", type=int, default=1,
                         help="Epochs without improvement before stopping (P4 uses 2).")
    trainer.add_argument("--warmup-steps", type=int, default=0,
                         help="Pin the warm-up ramp; 0 derives it from --epochs as before.")
    trainer.add_argument("--alpha", type=float, default=DEFAULT_ALPHA,
                         # OCT-LINT: disable=no_hardcoded_secrets reason="argparse help prose; the entropy is ordinary English, not a credential (operator ruling 2026-09-18)"
                         help="Plan 298 P1: what an error granting MORE permission costs "
                              "relative to one granting less (default: 1.0, symmetric).")
    trainer.add_argument("--cost-weight", type=float, default=DEFAULT_COST_WEIGHT,
                         help="Plan 298 P3: weight on the expected-cost term added to the "
                              "loss; 0.0 is plain cross-entropy (default).")
    trainer.add_argument("--stop-metric", choices=STOP_METRIC_CHOICES,
                         default=DEFAULT_STOP_METRIC,
                         help="Plan 300 P1: which validation figure early stopping and "
                              "checkpoint selection watch; ce is plan 290's rule and the "
                              "default, cost watches the expected cost under --alpha.")
 
    calibrator = subparsers.add_parser("calibrate", help="Fit the temperature.")
    calibrator.add_argument("--model", required=True, help="Model directory from train.")
    calibrator.add_argument("--corpus-dir", default=None)
    calibrator.add_argument("--batch", type=int, default=DEFAULT_BATCH)
    calibrator.add_argument("--max-length", type=int, default=DEFAULT_MAX_LENGTH)
    calibrator.add_argument("--threads", type=int, default=DEFAULT_THREADS)
    calibrator.add_argument("--device", choices=DEVICE_CHOICES, default=DEFAULT_DEVICE)
 
    evaluator = subparsers.add_parser("evaluate", help="Score a case set.")
    evaluator.add_argument("--model", required=True, help="Model directory from train.")
    evaluator.add_argument("--cases", required=True, help="Case file to predict.")
    evaluator.add_argument("--condition", required=True, help="Condition name for the file.")
    evaluator.add_argument("--model-name", required=True, help="Model name for the file.")
    evaluator.add_argument("--out", required=True, help="Result file to write.")
    evaluator.add_argument("--batch", type=int, default=1,
                           help="1 for per-decision latency, 32 for throughput.")
    evaluator.add_argument("--max-length", type=int, default=DEFAULT_MAX_LENGTH)
    evaluator.add_argument("--threads", type=int, default=DEFAULT_THREADS)
    evaluator.add_argument("--device", choices=DEVICE_CHOICES, default=DEFAULT_DEVICE,
                           help="cpu is the deployment latency reading.")
    evaluator.add_argument("--decision", choices=DECISION_CHOICES, default=DEFAULT_DECISION,
                           help="Plan 298 P2: argmax is the rule every earlier result file "
                                "carries; min-cost minimises expected cost under --alpha.")
    evaluator.add_argument("--alpha", type=float, default=DEFAULT_ALPHA,
                           # OCT-LINT: disable=no_hardcoded_secrets reason="argparse help prose; the entropy is ordinary English, not a credential (operator ruling 2026-09-18)"
                           help="Plan 298 P1: what an error granting MORE permission costs "
                                "relative to one granting less (default: 1.0, symmetric, "
                                "under which min-cost and argmax agree everywhere).")
 
    relabeller = subparsers.add_parser(
        "relabel", help="Write the policy-change probe's relabelled corpus copy."
    )
    relabeller.add_argument("--out", required=True, help="Directory to write the copy into.")
    relabeller.add_argument("--bench-dir", default=None)
    return parser
 
 
def main(argv: list[str] | None = None) -> int:
    """CLI entry point: train, calibrate, evaluate, or relabel."""
 
    func = "main"
    args = _build_parser().parse_args(argv)
    bench_dir = Path(getattr(args, "bench_dir", None) or _bench_dir())
    corpus_dir = Path(getattr(args, "corpus_dir", None) or _bench_dir())
    if args.command == "train":
        train(TrainConfig(
            backbone=args.backbone, corpus_dir=corpus_dir, out_dir=Path(args.out),
            tokenizer_dir=Path(args.tokenizer_dir), seed=args.seed, max_epochs=args.epochs,
            batch=args.batch, learning_rate=args.learning_rate, max_length=args.max_length,
            threads=args.threads, budget_seconds=args.budget_seconds, device=args.device,
            prune_vocab=args.prune_vocab, vocab_reserve=args.vocab_reserve,
            cache_ids=args.cache_ids, cache_root=Path(args.cache_root),
            val_subset=args.val_subset, patience=args.patience,
            warmup_steps=args.warmup_steps, alpha=args.alpha,
            cost_weight=args.cost_weight, stop_metric=args.stop_metric,
        ))
    elif args.command == "calibrate":
        calibrate(Path(args.model), corpus_dir, args.batch, args.max_length, args.threads,
                  args.device)
    elif args.command == "evaluate":
        evaluate(Path(args.model), Path(args.cases), args.condition, Path(args.out),
                 args.model_name, args.batch, args.max_length, args.threads, args.device,
                 args.decision, args.alpha)
    else:
        relabel(bench_dir, Path(args.out))
    _dbg("OC_CHAT.MAIN", func, f"{args.command}: exit 0", 2)
    return 0
 
 
if __name__ == "__main__":
    sys.exit(main())