token_classifier.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# bench/argv_policy/baselines/token_classifier.py
 
"""
The token floor: hashed character n-grams into a multinomial logistic regression.
 
Purpose
-------
This is the number every later result on this bench is read against. The argv
gate is, at bottom, a set of rules over the characters of a command string, so
before anyone reports that an embedding index or a fine-tuned encoder "learned
the policy", the bench needs to know what a bag of character n-grams and plain
gradient descent already achieve. If the floor sits near the gate, the task is a
token rule and plan 290's encoder has to be read against the floor, not against
chance.
 
Features are character n-grams for n in 1 to 4 over the command padded with
start and end sentinels, hashed into a fixed space with CRC-32 (Python's own
``hash`` is salted per process, so a model trained in one run would be
meaningless in the next), counted, and L2-normalised so a long command does not
simply outvote a short one. The model is a three-class softmax trained by
minibatch gradient descent with L2 regularisation. No fitting library is used:
the gradient is written out.
 
Run it under the PROJECT venv from the subproject root::
 
    python bench/argv_policy/baselines/token_classifier.py
 
Implementation note (a deliberate departure from plan 288 P4, which reads
"numpy or torch"): this is standard-library only. The project venv has neither
numpy nor torch, installing one is the orchestrator's hard stop, and the scratch
harness venv that does have them cannot import ``oc_diagnostics``, so a module
run there could not be Option C compliant. Pure Python is also the better fit
for the shape of the problem: the features are sparse -- a thirty-character
command touches roughly 120 of 65536 columns -- so a dense design matrix would
be about three gigabytes of mostly zeros, and the sparse inner loop this module
runs is the same arithmetic a sparse numpy implementation would perform.
 
Responsibilities
----------------
- :func:`feature_index` / :func:`features` -- the hashed n-gram representation.
- :class:`Model` -- the weights and the softmax forecast.
- :func:`train` -- minibatch gradient descent over a corpus.
- :func:`evaluate` -- a case set to result rows in the bench schema.
- :func:`main` -- train, evaluate both case sets, write both result files.
 
Diagnostics
-----------
Domain: OC_CHAT.MAIN
Levels:
    L2 - training lifecycle (start, per-epoch loss, finish, files written)
    L3 - per-epoch accuracy and per-case-set agreement
    L4 - module load
 
Contracts
---------
- Deterministic: a corpus, a seed and the hyperparameters fix the weights, so
  the committed result files are reproducible from the committed corpus.
- Read-only with respect to the corpus and the case sets.
- Console output is ASCII only; commands are printed through :func:`ascii`.
- Standard library only, by the reasoning in the implementation note above.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
"""
from __future__ import annotations
 
import argparse
import json
import math
import random
import sys
import time
import zlib
from dataclasses import dataclass
from pathlib import Path
 
from oc_diagnostics import _dbg
 
_dbg("OC_CHAT.MAIN", "module", "loaded", 4)
 
#: The three verdicts, in the column order of the weight matrix.
LABELS: tuple[str, ...] = ("allow", "confirm", "deny")
 
#: Character n-gram lengths. One character carries the metacharacter signal,
#: four carries option and subcommand stems like ``-la`` or ``inst``.
NGRAM_SIZES: tuple[int, ...] = (1, 2, 3, 4)
 
#: The fixed hash space. Large enough that collisions are rare at this corpus
#: size, small enough that the weight matrix is a few hundred thousand floats.
HASH_SPACE = 1 << 16
 
#: Sentinels marking the start and end of a command, so a leading dash or a
#: trailing quote is distinguishable from the same character in the middle.
START_SENTINEL = "\x02"
END_SENTINEL = "\x03"
 
#: Hyperparameters chosen by a sweep on the validation split, not by taste: the
#: rate and epoch count below are where held-out agreement stops improving
#: (lr 40 / 50 epochs reaches 98.7 percent validation; lr 80 and 100 epochs do
#: no better). A floor tuned too low would understate the baseline that plans
#: 289 and 290 are read against, which is the one way this module could mislead.
DEFAULT_EPOCHS = 50
DEFAULT_SEED = 7
DEFAULT_LEARNING_RATE = 40.0
DEFAULT_L2 = 1e-6
DEFAULT_BATCH = 32
 
 
def feature_index(size: int, gram: str) -> int:
    """Hash one n-gram of length *size* into the feature space.
 
    CRC-32 rather than :func:`hash`, which is salted per process: a model whose
    columns move between runs cannot be trained once and reported on later. The
    length is folded into the key so a two-character gram and the first two
    characters of a four-character gram do not systematically share a column.
    """
 
    return zlib.crc32(f"{size}\x1f{gram}".encode("utf-8")) % HASH_SPACE
 
 
def features(command: str) -> dict[int, float]:
    """The L2-normalised hashed n-gram counts of *command*, plus a bias column."""
 
    padded = f"{START_SENTINEL}{command}{END_SENTINEL}"
    counts: dict[int, float] = {}
    for size in NGRAM_SIZES:
        for start in range(len(padded) - size + 1):
            index = feature_index(size, padded[start:start + size])
            counts[index] = counts.get(index, 0.0) + 1.0
    norm = math.sqrt(sum(value * value for value in counts.values())) or 1.0
    scaled = {index: value / norm for index, value in counts.items()}
    scaled[0] = scaled.get(0, 0.0) + 1.0        # bias column
    return scaled
 
 
@dataclass
class Model:
    """A three-class softmax over the hashed n-gram space."""
 
    weights: list[list[float]]
    biases: list[float]
 
    def scores(self, vector: dict[int, float]) -> list[float]:
        """The unnormalised class scores of one feature vector."""
 
        return [
            self.biases[label] + sum(value * self.weights[label][index]
                                     for index, value in vector.items())
            for label in range(len(LABELS))
        ]
 
    def probabilities(self, vector: dict[int, float]) -> list[float]:
        """The softmax of :meth:`scores`, computed shift-stably."""
 
        raw = self.scores(vector)
        top = max(raw)
        exponentials = [math.exp(value - top) for value in raw]
        total = sum(exponentials)
        return [value / total for value in exponentials]
 
    def predict(self, command: str) -> tuple[str, dict[str, float]]:
        """The predicted verdict for *command* and the full forecast."""
 
        probabilities = self.probabilities(features(command))
        best = max(range(len(LABELS)), key=probabilities.__getitem__)
        return LABELS[best], dict(zip(LABELS, probabilities))
 
 
def train(
    rows: list[dict[str, str]],
    epochs: int = DEFAULT_EPOCHS,
    seed: int = DEFAULT_SEED,
    learning_rate: float = DEFAULT_LEARNING_RATE,
    l2: float = DEFAULT_L2,
    batch: int = DEFAULT_BATCH,
) -> Model:
    """Fit the softmax by minibatch gradient descent over *rows*.
 
    Features are extracted once up front rather than per epoch, which is the
    difference between a run of seconds and a run of minutes.
    """
 
    func = "train"
    _dbg("OC_CHAT.MAIN", func, f"{len(rows)} rows, {epochs} epochs, batch {batch}", 2)
    label_of = {name: index for index, name in enumerate(LABELS)}
    examples = [(features(row["command"]), label_of[row["truth"]]) for row in rows]
    model = Model(
        weights=[[0.0] * HASH_SPACE for _ in LABELS],
        biases=[0.0] * len(LABELS),
    )
    rng = random.Random(seed)
    order = list(range(len(examples)))
    for epoch in range(epochs):
        rng.shuffle(order)
        rate = learning_rate / (1.0 + epoch)
        loss = 0.0
        hits = 0
        for start in range(0, len(order), batch):
            chunk = order[start:start + batch]
            gradients: list[dict[int, float]] = [{} for _ in LABELS]
            bias_gradients = [0.0] * len(LABELS)
            for position in chunk:
                vector, truth = examples[position]
                probabilities = model.probabilities(vector)
                loss -= math.log(max(probabilities[truth], 1e-12))
                hits += int(max(range(len(LABELS)),
                                key=probabilities.__getitem__) == truth)
                for label in range(len(LABELS)):
                    error = probabilities[label] - (1.0 if label == truth else 0.0)
                    if error == 0.0:
                        continue
                    bias_gradients[label] += error
                    target = gradients[label]
                    for index, value in vector.items():
                        target[index] = target.get(index, 0.0) + error * value
            step = rate / len(chunk)
            for label in range(len(LABELS)):
                row = model.weights[label]
                for index, gradient in gradients[label].items():
                    row[index] -= step * (gradient + l2 * row[index])
                model.biases[label] -= step * bias_gradients[label]
        _dbg("OC_CHAT.MAIN", func,
             f"epoch {epoch + 1}: loss {loss / len(examples):.4f} "
             f"train accuracy {hits / len(examples):.4f}", 3)
    _dbg("OC_CHAT.MAIN", func, "training complete", 2)
    return model
 
 
def evaluate(model: Model, cases: list[dict[str, str]]) -> list[dict[str, object]]:
    """Predict every case, timing each decision in process."""
 
    func = "evaluate"
    rows: list[dict[str, object]] = []
    for case in cases:
        started = time.perf_counter()
        prediction, probabilities = model.predict(case["command"])
        elapsed_ms = (time.perf_counter() - started) * 1000.0
        rows.append({
            "command": case["command"],
            "truth": case["truth"],
            "slice": case["slice"],
            "pred": prediction,
            "probs": probabilities,
            "ms": elapsed_ms,
            "reason": case["reason"],
        })
    hits = sum(row["pred"] == row["truth"] for row in rows)
    _dbg("OC_CHAT.MAIN", func, f"{hits}/{len(rows)} agreement", 3)
    return rows
 
 
def _accuracy(model: Model, rows: list[dict[str, str]]) -> float:
    """Share of *rows* the model gets right."""
 
    if not rows:
        return 0.0
    return sum(model.predict(row["command"])[0] == row["truth"] for row in rows) / len(rows)
 
 
def _build_parser() -> argparse.ArgumentParser:
    """Construct the CLI argument parser."""
 
    parser = argparse.ArgumentParser(
        prog="python bench/argv_policy/baselines/token_classifier.py",
        description="Train the hashed n-gram floor and score both case sets.",
    )
    parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS,
                        help=f"Passes over the corpus (default: {DEFAULT_EPOCHS}).")
    parser.add_argument("--seed", type=int, default=DEFAULT_SEED,
                        help=f"Shuffle seed (default: {DEFAULT_SEED}).")
    parser.add_argument("--learning-rate", type=float, default=DEFAULT_LEARNING_RATE,
                        help=f"Initial step size (default: {DEFAULT_LEARNING_RATE}).")
    parser.add_argument("--l2", type=float, default=DEFAULT_L2,
                        help=f"L2 penalty (default: {DEFAULT_L2}).")
    parser.add_argument("--batch", type=int, default=DEFAULT_BATCH,
                        help=f"Minibatch size (default: {DEFAULT_BATCH}).")
    parser.add_argument("--out-dir", default=None,
                        help="Where to write result files (default: ../results).")
    return parser
 
 
def main(argv: list[str] | None = None) -> int:
    """CLI entry point: train the floor and write both result files."""
 
    func = "main"
    args = _build_parser().parse_args(argv)
    bench_dir = Path(__file__).resolve().parent.parent
    out_dir = Path(args.out_dir) if args.out_dir else bench_dir / "results"
    out_dir.mkdir(parents=True, exist_ok=True)
 
    corpus = [json.loads(line) for line
              in (bench_dir / "corpus.jsonl").read_text(encoding="utf-8").splitlines() if line]
    validation = [json.loads(line) for line
                  in (bench_dir / "corpus_val.jsonl").read_text(encoding="utf-8").splitlines()
                  if line]
    print(f"token_classifier: training on {len(corpus)} commands, "
          f"{args.epochs} epochs, hash space {HASH_SPACE}")
    started = time.perf_counter()
    model = train(corpus, args.epochs, args.seed, args.learning_rate, args.l2, args.batch)
    train_seconds = time.perf_counter() - started
    train_accuracy = _accuracy(model, corpus)
    validation_accuracy = _accuracy(model, validation)
    print(f"token_classifier: trained in {train_seconds:.1f}s; "
          f"train {train_accuracy:.4f}, validation {validation_accuracy:.4f}")
 
    for condition, case_file in (("hand", "cases.json"), ("hard", "cases_hard.json")):
        cases = json.loads((bench_dir / case_file).read_text(encoding="utf-8"))
        rows = evaluate(model, cases)
        hits = sum(row["pred"] == row["truth"] for row in rows)
        payload = {
            "model": "token_classifier",
            "condition": condition,
            "results": rows,
            "provenance": {
                "route": "CPU, standard library, hashed character n-grams",
                "case_file": case_file,
                "train_rows": len(corpus),
                "validation_rows": len(validation),
                "epochs": args.epochs,
                "seed": args.seed,
                "learning_rate": args.learning_rate,
                "l2": args.l2,
                "batch": args.batch,
                "hash_space": HASH_SPACE,
                "ngram_sizes": list(NGRAM_SIZES),
                "train_seconds": train_seconds,
                "train_accuracy": train_accuracy,
                "validation_accuracy": validation_accuracy,
            },
        }
        target = out_dir / f"token_classifier_{condition}.json"
        target.write_text(json.dumps(payload, indent=1, sort_keys=True) + "\n", encoding="utf-8")
        print(f"  {condition:5} {hits}/{len(rows)} agreement -> {target.name}")
    _dbg("OC_CHAT.MAIN", func, "exit 0", 2)
    return 0
 
 
if __name__ == "__main__":
    sys.exit(main())