app/fp/parser2.py

"""Train and score the production head design on the contract v2.1 corpus.

The architecture question this answers is narrow and specific: can a small
encoder emit *typed spans* over layouts it never saw, with every value resolved
by deterministic contract code? There is no value head here on purpose. The
pilot needed one, and its hard-split season accuracy was the only field that
survived pretraining; if this run reproduces field-level reliability without it,
the head was compensating for narrow data rather than doing necessary work.

Predictions are scored through `fp.normalize.record`, the same function that
derived the gold. A prediction is therefore a parse under the ratified contract,
not a string diff: abstention, ranges, lists and case folding are all evaluated
as the contract defines them, which is what makes these numbers comparable with
the real-data numbers we cannot produce yet.

Usage:
    python -m fp.parser2 --arm B-broad --mode characters --seed 0 \
        --out runs/parser-v2/broad-characters-0
"""

import argparse
import importlib.metadata
import json
import math
import os
import random
import statistics
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path

# Required by deterministic CUDA matrix multiplication. This has to be present
# before the first CUDA context is initialized; setting it here keeps remote runs
# reproducible without relying on one shell's environment.
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")

import torch
from torch import nn
from transformers import BertConfig, BertModel

from fp import corpus2, normalize
from fp.corpus2 import ARM_SPLITS, ROOT as CORPUS, augment, augment_training, digest, rows
from fp.corpus2 import KINDS as corpus_KINDS
from fp.normalize import record as contract_record

SNAP = Path(os.environ.get(
    "ANIFILEBERT_SNAPSHOT",
    "/home/lain/.cache/huggingface/hub/models--ModerRAS--AniFileBERT/snapshots/"
    "d8ddb0b54dad4d65a60ab7eba5acb06a2a1fab02"))
KINDS = list(corpus_KINDS)
LABELS = ["O"] + [prefix + kind for kind in KINDS for prefix in ("B-", "I-")]
# A schedule is a named, frozen list, not the meaning of a flag. Round 4 shipped
# `balanced` as four views; growing that same name to six views in round 5 would
# have made `scripts/parser-v2-round4.sh` silently produce a different experiment
# under a published name, so the older list stays pinned under `balanced4`.
SCHEDULES = {"balanced4": ("none", "shuffle", "shuffle-values", "shuffle-paths"),
             "balanced": ("none", "shuffle", "shuffle-values", "shuffle-paths",
                          "shuffle-markers", "shuffle-boundary"),
             # Round 5b's dose control. Uniform rotation makes every view's share
             # 1/|views|, so adding two views silently cut the value view from 3/12 to
             # 2/12 -- which explains the probe-value drop without any interference
             # hypothesis. This list restores the value share to 0.25 (C1's) while
             # keeping the marker share at C4's 2/12, trading only boundary, the view
             # with the weakest measured effect. A 12-slot list visited once per epoch
             # makes each row's counts exact, not merely balanced within one.
             "balanced-dose": ("none", "none", "shuffle", "shuffle",
                               "shuffle-values", "shuffle-values", "shuffle-values",
                               "shuffle-paths", "shuffle-paths",
                               "shuffle-markers", "shuffle-markers",
                               "shuffle-boundary")}
BALANCED_VIEWS = SCHEDULES["balanced"]
SELECTION_VIEWS = ("plain", "shuffle", "markers", "boundary")


def balanced_view(row_index, epoch, seed, views=BALANCED_VIEWS):
    """Exact rotation over a named schedule.

    Unlike random independent gates, this guarantees each row receives each view the
    same number of times over a full schedule: three times per view in round 4's
    twelve epochs and four views, twice per view in round 5's six. `seed` rotates the
    starting cell without changing the balance.
    """
    return views[(row_index + epoch + seed) % len(views)]


def selection_views(validation_rows, seed, forbidden_queries):
    """Fixed derived views of the validation split, for checkpoint selection.

    Round 4 selected on seen-layout validation while the acceptance target is about
    transformed views, so the choice now averages the same transformations it is
    meant to be good at. Only validation rows are read. A derived view that lands on
    any scored query is replaced by its parent: an augmented row that happened to
    reproduce held-out text would leak the evaluation into selection through the same
    back door that round 3 had to close for training rows.
    """
    streams = {"shuffle": (augment, random.Random(seed + 20260919)),
               "markers": (corpus2.augment_markers, random.Random(seed + 20260920)),
               "boundary": (corpus2.augment_boundary, random.Random(seed + 20260921))}
    views = {"plain": list(validation_rows)}
    counts = {name: {"accepted": 0, "fell_back": 0, "leakage_guarded": 0}
              for name in streams}
    for name, (transform, rng) in streams.items():
        derived = []
        for row in validation_rows:
            candidate = transform(row, rng)
            if candidate is None:
                counts[name]["fell_back"] += 1
                derived.append(row)
                continue
            if candidate["query"] in forbidden_queries:
                counts[name]["leakage_guarded"] += 1
                derived.append(row)
                continue
            counts[name]["accepted"] += 1
            derived.append(candidate)
        views[name] = derived
    return views, counts
SCORE_SPLITS = ("validation", "easy", "heldout-layout", "heldout-lexical", "heldout-value",
                "heldout-negative", "probe-base", "probe-layout", "probe-script",
                "probe-value", "probe-episode-title")


def save_json(path, value):
    Path(path).write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n")


def resolve_device(requested):
    if requested == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if requested == "cuda" and not torch.cuda.is_available():
        raise SystemExit("--device cuda requested, but torch.cuda.is_available() is false")
    return torch.device(requested)


def on_device(batch, device):
    return tuple(tensor.to(device, non_blocking=device.type == "cuda") for tensor in batch)


def cpu_state_dict(model):
    """Portable checkpoint: never make a local rescore depend on CUDA existing."""
    return {name: value.detach().cpu() for name, value in model.state_dict().items()}


class Extractor(nn.Module):
    def __init__(self, mode, n_tags):
        super().__init__()
        self.mode, self.tokenizer = mode, None
        if mode != "bytes":
            sys.path.insert(0, str(SNAP))
            from anifilebert.model import load_model
            from anifilebert.tokenizer import load_tokenizer
            self.tokenizer = load_tokenizer(str(SNAP))
            if mode == "pretrained":
                self.encoder = load_model(str(SNAP)).bert
            else:
                config = BertConfig.from_pretrained(str(SNAP), local_files_only=True)
                self.encoder = BertModel(config, add_pooling_layer=False)
        else:
            config = BertConfig(vocab_size=257, hidden_size=256, num_hidden_layers=4,
                                num_attention_heads=8, intermediate_size=1024,
                                max_position_embeddings=512, pad_token_id=256)
            self.encoder = BertModel(config, add_pooling_layer=False)
        self.tags = nn.Linear(256, n_tags)

    def encode(self, row):
        """Tokens plus the [start, end) token window of every character.

        Byte mode widens the window for multi-byte characters, so a boundary can
        only ever fall between complete UTF-8 characters.
        """
        query = row["query"]
        if self.mode == "bytes":
            ids, starts, ends, pos = list(query.encode()), [], [], 0
            for char in query:
                starts.append(pos)
                pos += len(char.encode())
                ends.append(pos)
        else:
            tokens = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(query))
            ids = [self.tokenizer.cls_token_id] + tokens + [self.tokenizer.sep_token_id]
            starts, ends = list(range(1, len(query) + 1)), list(range(2, len(query) + 2))
            assert len(tokens) == len(query), (
                "char-window pooling assumes one token per character; the tokenizer "
                "changed. Refusing to train on misaligned labels.")
        assert len(ids) <= self.encoder.config.max_position_embeddings, "No truncation"
        labels = [0] * len(query)
        for span in row["spans"]:
            start, end = span["start"], span["end"]
            labels[start] = LABELS.index("B-" + span["kind"])
            for i in range(start + 1, end):
                labels[i] = LABELS.index("I-" + span["kind"])
        return ids, starts, ends, labels, row

    def forward(self, batch):
        ids, attention, starts, ends = batch
        hidden = self.encoder(input_ids=ids, attention_mask=attention).last_hidden_state
        cumulative = torch.cat((torch.zeros_like(hidden[:, :1]), hidden.cumsum(1)), 1)
        a = starts.unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
        b = ends.unsqueeze(-1).expand_as(a)
        width = (ends - starts).clamp_min(1).unsqueeze(-1)
        chars = (cumulative.gather(1, b) - cumulative.gather(1, a)) / width
        # Flatten before the tag layer: DynamicQuantizedLinear, which is what the
        # INT8 export installs, only accepts 2-D input, so a rank-3 call would make
        # the shipped artifact crash while the fp32 prototype scored fine.
        tagged = self.tags(chars.reshape(-1, chars.shape[-1]))
        return tagged.reshape(*chars.shape[:2], -1), chars


def collate(items, pad):
    length = max(len(item[0]) for item in items)
    chars = max(len(item[1]) for item in items)
    ids = torch.full((len(items), length), pad, dtype=torch.long)
    attention = torch.zeros_like(ids)
    starts = torch.zeros((len(items), chars), dtype=torch.long)
    ends, labels = starts.clone(), torch.full((len(items), chars), -100, dtype=torch.long)
    for i, (tokens, a, b, y, _) in enumerate(items):
        ids[i, :len(tokens)] = torch.tensor(tokens)
        attention[i, :len(tokens)] = 1
        starts[i, :len(a)], ends[i, :len(b)] = torch.tensor(a), torch.tensor(b)
        labels[i, :len(y)] = torch.tensor(y)
    return (ids, attention, starts, ends), labels


def legality_matrix():
    """Zero where a BIO jump is legal, -inf where it is not.

    This is the strict control for the empirical table below: a constraint says
    which sequences are admissible and lets the emissions decide among them. The
    empirical matrix additionally supplies a learned ordering prior, so its effect
    must be measured separately rather than confused with BIO legality.
    """
    allow = torch.zeros(len(LABELS), len(LABELS))
    for source, from_tag in enumerate(LABELS):
        for target, to_tag in enumerate(LABELS):
            if to_tag.startswith("I-") and from_tag not in (to_tag, "B-" + to_tag[2:]):
                allow[source, target] = -torch.inf
    return allow


def transitions(items):
    """Empirical log-transition matrix over the tag sequence.

    This answers a different question than `legality_matrix`: whether the
    *distribution* of arrangements the model was trained on supplies a useful tag
    order. It is not the default because it can trade transfer for fit; `--transition`
    keeps that factor explicit.

    Greedy per-position argmax is allowed to emit an `I-X` that opens a span, and
    the legalizer silently throws those away. A transition table makes the model
    answer for that: it costs one matrix, needs no retraining, and can be applied
    to an existing checkpoint, which is why this is the cheapest hypothesis here.
    """
    counts = torch.zeros(len(LABELS), len(LABELS))
    for item in items:
        previous = LABELS.index("O")
        for label in item[3]:
            counts[previous, label] += 1
            previous = label
        counts[previous, LABELS.index("O")] += 1
    probabilities = counts / counts.sum(1, keepdim=True).clamp_min(1)
    return torch.log(probabilities + 1e-12)


def viterbi(logits, transition):
    """Exact best tag sequence under a transition matrix. Legal in the sense the
    matrix defines it: an impossible jump is -inf, not a silent drop.

    The reduction is over the *source* axis. `scores` is indexed [from, to], so the
    best predecessor of each target is `max(0)`; an earlier version of this function
    took `max(1)`, which collapsed targets instead of sources and returned an all-`O`
    path on every split - silently scoring 0 spans on rows the same checkpoint got
    311/311 right. The test `test_viterbi_returns_one_state_per_frame_and_repairs_an_orphan`
    exists because a bug that makes a decoder *abstain* looks like a cautious model
    in an aggregate table.
    """
    best = torch.full((len(LABELS),), -torch.inf)
    best[LABELS.index("O")] = 0.0
    back = []
    for frame in logits:
        scores = best.unsqueeze(1) + transition  # [from, to]
        back.append(scores.argmax(0))            # best predecessor for each target
        best = frame + scores.max(0).values      # score of the best path ending at each target
    path = [int(best.argmax())]
    for step in reversed(back):
        path.append(int(step[path[-1]]))
    path.reverse()
    # One state per frame, not n + 1: the recursion's seed state is not a tag. An
    # earlier version returned it, so every frame was scored with its neighbour's
    # label - which in a span task shifts every boundary by one character.
    return path[1:]


def viterbi_batch(logits, lengths, transition):
    """Vectorized equivalent of `viterbi`, with a distinct end for every row."""
    batch, frames, labels = logits.shape
    lengths = torch.as_tensor(lengths, device=logits.device)
    best = torch.full((batch, labels), -torch.inf, device=logits.device,
                      dtype=logits.dtype)
    best[:, LABELS.index("O")] = 0.0
    back = []
    for index in range(frames):
        scores = best.unsqueeze(2) + transition.unsqueeze(0)
        predecessor = scores.argmax(1)
        candidate = logits[:, index] + scores.max(1).values
        active = (index < lengths).unsqueeze(1)
        best = torch.where(active, candidate, best)
        back.append(predecessor)
    current = best.argmax(1)
    paths = torch.zeros((batch, frames), dtype=torch.long, device=logits.device)
    for index in reversed(range(frames)):
        active = index < lengths
        paths[:, index] = torch.where(active, current, paths[:, index])
        predecessor = back[index].gather(1, current.unsqueeze(1)).squeeze(1)
        current = torch.where(active, predecessor, current)
    return [paths[row, :int(length)].tolist() for row, length in enumerate(lengths)]


def decode(tag_ids, limit):
    """BIO to spans. An `I-` that does not continue its own kind ends the span, and
    an orphan `I-` becomes nothing: the decoder never invents a boundary."""
    spans, kind, begin = [], None, 0
    for i, tag in enumerate([LABELS[t] for t in tag_ids[:limit]] + ["O"]):
        if kind is not None and tag != "I-" + kind:
            spans.append((kind, begin, i))
            kind = None
        if tag.startswith("B-"):
            kind, begin = tag[2:], i
    return spans


def predict_spans(query, tag_ids):
    """Decoder output as contract spans, with no text and no value: the record is
    whatever `fp.normalize` can derive from the boundaries alone."""
    return [{"kind": kind, "start": start, "end": end}
            for kind, start, end in decode(tag_ids, len(query))]


def tag_paths(logits, lengths, decode_mode="greedy", transition=None):
    """Decode only real character frames, never another row's batch padding.

    Viterbi backtracking starts at its final frame, so decoding the padded suffix
    can change the state selected for the final real character. Greedy happened to
    be insensitive because `predict_spans` truncates, but both decoders should be
    explicitly batch-composition invariant.
    """
    if decode_mode == "greedy":
        return [frame[:length].argmax(-1).tolist()
                for frame, length in zip(logits, lengths)]
    return viterbi_batch(logits, lengths, transition)


def label_weights(items):
    """Inverse frequency, sqrt-damped, mean 1. DISCARD and TITLE dominate the
    character budget, and without this the rare kinds are unlearnable.

    Keys are label *ids*: returning them indexed by name silently produced all
    zeros, which made the loss NaN and the exported artifact crash.
    """
    counts = Counter(label for item in items for label in item[3] if label >= 0)
    total = sum(counts.values())
    raw = {label_id: 1 / math.sqrt(count / total * len(LABELS))
           for label_id, count in counts.items()}
    mean = sum(raw[label] * counts[label] for label in raw) / total
    return [raw.get(index, 0.0) / mean for index in range(len(LABELS))]


@torch.inference_mode()
def evaluate(model, encoded, batch_size=64, timed=False, decode_mode="greedy", transition=None):
    model.eval()
    device = next(model.parameters()).device
    counts, bylang, bylayout, predictions, latencies = Counter(), {}, {}, [], []
    per_kind = defaultdict(Counter)
    start = time.perf_counter()
    for offset in range(0, len(encoded), batch_size):
        items = encoded[offset:offset + batch_size]
        batch, _ = collate(items, model.encoder.config.pad_token_id)
        batch = on_device(batch, device)
        logits, _ = model(batch)
        # Decoding is tiny and intentionally stays on CPU, where Viterbi's Python
        # loop and the deterministic contract resolver already operate.
        logits = logits.detach().cpu()
        tags = tag_paths(logits, [len(item[1]) for item in items], decode_mode, transition)
        for i, item in enumerate(items):
            row = item[-1]
            spans = predict_spans(row["query"], tags[i])
            actual, issues = normalize.gated_record(row["query"], spans)
            expected = row["record"]
            exact = actual == expected
            counts["n"] += 1
            counts["exact"] += exact
            counts["abstained"] += actual is None
            counts["wrong_abstain"] += actual is None and expected is not None
            counts["false_emit"] += actual is not None and expected is None
            counts["issue_rows"] += bool(issues)
            for span in row["spans"]:
                per_kind[span["kind"]]["gold"] += 1
            for span in spans:
                per_kind[span["kind"]]["pred"] += 1
            gold = {(s["kind"], s["start"], s["end"]) for s in row["spans"]}
            pred = {(s["kind"], s["start"], s["end"]) for s in spans}
            counts["span_tp"] += len(gold & pred)
            counts["span_pred"], counts["span_gold"] = counts["span_pred"] + len(pred), counts["span_gold"] + len(gold)
            for kind, _, _ in gold & pred:
                per_kind[kind]["tp"] += 1
            for key, value in (actual or {}).items():
                counts["emitted_fields"] += 1
                counts["wrong_emitted_fields"] += (expected or {}).get(key) != value
            for key, value in (expected or {}).items():
                counts["gold_fields"] += 1
                # Compare prediction to gold. Writing `(expected or {})` on both
                # sides made every gold field look correct and reported 100%.
                counts["correct_fields"] += (actual or {}).get(key) == value
            counts["correct_abstain"] += actual is None and expected is None
            bylang.setdefault(row["language"], Counter())["n"] += 1
            bylang[row["language"]]["exact"] += exact
            slot = bylayout.setdefault(row["layout"], Counter())
            slot["n"] += 1
            slot["exact"] += exact
            predictions.append({"query": row["query"], "layout": row["layout"],
                                "language": row["language"], "expected": expected,
                                "actual": actual, "spans": spans,
                                "issues": issues, "exact": exact})
            if timed:
                latencies.append(1000 * single_row_ms(model, item))
    elapsed = time.perf_counter() - start
    out = dict(counts)
    out["span_f1"] = f1(counts["span_tp"], counts["span_pred"], counts["span_gold"])
    out["kinds"] = {kind: {"gold": c["gold"], "pred": c["pred"], "tp": c["tp"],
                           "f1": f1(c["tp"], c["pred"], c["gold"])}
                    for kind, c in sorted(per_kind.items())}
    out["by_language"] = {k: dict(v) for k, v in sorted(bylang.items())}
    out["by_layout"] = {k: dict(v) for k, v in sorted(bylayout.items())}
    out["ms_per_row_batched"] = 1000 * elapsed / max(1, len(encoded))
    if latencies:
        out["latency_ms_p50"] = round(statistics.median(latencies), 2)
        out["latency_ms_p95"] = round(sorted(latencies)[int(0.95 * len(latencies))], 2)
    return out, predictions


def single_row_ms(model, item):
    began = time.perf_counter()
    batch, _ = collate([item], model.encoder.config.pad_token_id)
    device = next(model.parameters()).device
    batch = on_device(batch, device)
    logits, _ = model(batch)
    torch.cuda.synchronize(device) if device.type == "cuda" else None
    logits.argmax(-1)
    return time.perf_counter() - began


def f1(tp, pred, gold):
    precision, recall = (tp / pred if pred else 0.0), (tp / gold if gold else 0.0)
    return round(2 * precision * recall / (precision + recall), 4) if precision + recall else 0.0


def train(args):
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=False)
    torch.set_num_threads(args.threads)
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.use_deterministic_algorithms(True)
    device = resolve_device(args.device)
    train_splits = ARM_SPLITS[args.arm]
    train_rows = [row for split in train_splits for row in rows(split)]
    data = {split: rows(split) for split in SCORE_SPLITS}
    scored_queries = set()
    for value in data.values():
        scored_queries |= {row["query"] for row in value}
    leaked = scored_queries.intersection(row["query"] for row in train_rows)
    assert not leaked, f"train rows appear in a scored split: {sorted(leaked)[:3]}"
    if args.smoke:
        train_rows, data = train_rows[:160], {k: v[:40] for k, v in data.items()}
    model = Extractor(args.mode, len(LABELS))
    encoded = {key: [model.encode(row) for row in value] for key, value in
               list(data.items()) + [("train", train_rows)]}
    model.to(device)
    weights = torch.tensor(label_weights(encoded["train"]), dtype=torch.float, device=device)
    transition = (legality_matrix() if args.transition == "legal"
                  else transitions(encoded["train"]))
    versions = {p: importlib.metadata.version(p) for p in ("torch", "transformers", "numpy")}
    unknown = {}
    if model.tokenizer:
        unk = model.tokenizer.unk_token_id
        # The audit number that matters is how much of the training text the
        # vocabulary cannot represent at all, recorded per run.
        tokens = [token for item in encoded["train"] for token in item[0]]
        unknown = {"train": round(sum(t == unk for t in tokens) / len(tokens), 4)}
    # Recorded verbatim so a published run can never be reinterpreted after the
    # view list changes: `--augment balanced` means whatever this file says it meant.
    views = SCHEDULES.get(args.augment)
    selected_views, selection_counts = selection_views(data["validation"], args.seed,
                                                       scored_queries)
    selection_encoded = {name: [model.encode(row) for row in value]
                         for name, value in selected_views.items()}
    save_json(out / "config.json", {
        "args": vars(args), "versions": versions, "corpus_sha256": digest(CORPUS / "manifest.json"),
        "source_sha256s": {"parser2": digest(__file__),
                           "corpus2": digest(Path(corpus2.__file__)),
                           "normalize": digest(Path(normalize.__file__))},
        "corpus_splits": train_splits, "train_rows": len(train_rows),
        "parameters": sum(p.numel() for p in model.parameters()),
        "source_sha256": digest(__file__), "encoder": model.encoder.config.to_dict(),
        "snapshot": str(SNAP) if args.mode != "bytes" else None,
        "pretraining_overlap": "unknown for upstream weights" if args.mode == "pretrained" else None,
        "value_head": False, "unknown_token_fraction": unknown,
        "device": {"requested": args.device, "resolved": str(device),
                   "cuda": torch.version.cuda,
                   "name": torch.cuda.get_device_name(device) if device.type == "cuda" else None},
        "label_weights_min": round(float(weights.min()), 3),
        "label_weights_max": round(float(weights.max()), 3),
        "balanced_views": list(views) if views else None,
        "selection": ("mean exact-record over fixed plain/shuffle/markers/boundary "
                      "views of the validation split; earliest checkpoint breaks ties")
        if args.select == "multi-view" else
        "validation exact-record; earliest checkpoint breaks ties",
        # Recorded before training so a reviewer can see that selection never read a
        # held-out row, and can rebuild the same views from the seed.
        "selection_view_counts": selection_counts,
        "selection_view_queries": {name: [row["query"] for row in value]
                                   for name, value in selected_views.items()}})
    optimizer = torch.optim.AdamW([
        {"params": model.encoder.parameters(), "lr": 5e-5 if args.mode == "pretrained" else 3e-4},
        {"params": model.tags.parameters(), "lr": 3e-4}], weight_decay=0.01)
    best, history, started = -1.0, [], time.perf_counter()
    rng = random.Random(args.seed)
    balanced_position_rng = random.Random(args.seed + 20260918)
    value_rng = random.Random(args.seed + 20260916)
    path_rng = random.Random(args.seed + 20260917)
    marker_rng = random.Random(args.seed + 20260922)
    boundary_rng = random.Random(args.seed + 20260923)
    for epoch in range(args.epochs):
        model.train()
        items, augmentation = [], Counter()
        for row_index, row in enumerate(train_rows):
            view = row
            view_mode = (balanced_view(row_index, epoch, args.seed, views)
                         if args.augment in SCHEDULES else args.augment)
            selected = (view_mode != "none" and
                        (args.augment in SCHEDULES or rng.random() < args.augment_rate))
            augmentation["scheduled_" + view_mode] += 1
            if selected:
                augmentation["selected"] += 1
                # A different arrangement of the same atoms every epoch: the label
                # multiset is unchanged unless this arm explicitly adds value
                # support. Axis-specific RNGs do not perturb which positional view
                # the paired round-2 control would have seen.
                position_rng = balanced_position_rng if args.augment in SCHEDULES else rng
                candidate = augment_training(row, view_mode, position_rng, value_rng, path_rng,
                                             marker_rng, boundary_rng)
                # The augmented view is built at train time, so it never passed the
                # corpus dedupe. A view that lands on a scored query would be test
                # leakage through the back door, and a shuffled negative row is the
                # realistic way to hit one.
                if candidate and candidate["query"] not in scored_queries:
                    view = candidate
                    augmentation["accepted"] += 1
                    for axis in candidate.get("augmented_axes", []):
                        augmentation[axis] += 1
                elif candidate:
                    augmentation["scored_query_collision"] += 1
                else:
                    augmentation["rejected"] += 1
            items.append(model.encode(view))
        rng.shuffle(items)
        losses = []
        for offset in range(0, len(items), args.batch):
            minibatch = items[offset:offset + args.batch]
            batch, labels = collate(minibatch, model.encoder.config.pad_token_id)
            batch, labels = on_device(batch, device), labels.to(device, non_blocking=True)
            logits, _ = model(batch)
            loss = nn.functional.cross_entropy(logits.reshape(-1, len(LABELS)),
                                               labels.reshape(-1), weight=weights)
            if not torch.isfinite(loss):
                raise SystemExit(f"non-finite loss at epoch {epoch + 1} step {offset}: "
                                 f"a zero/nan weight or lr is producing a model that only "
                                 f"looks trained")
            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            losses.append(float(loss.detach()))
        val, _ = evaluate(model, encoded["validation"], decode_mode=args.decode,
                          transition=transition)
        # The selection objective is the mean over the fixed derived views; the
        # plain validation number stays recorded so rounds 3 and 4 remain readable.
        view_scores = {}
        for name in SELECTION_VIEWS:
            if name == "plain":
                view_scores[name] = val["exact"]
                continue
            scored_view, _ = evaluate(model, selection_encoded[name],
                                      decode_mode=args.decode, transition=transition)
            view_scores[name] = scored_view["exact"]
        objective = (sum(view_scores.values()) /
                     (len(SELECTION_VIEWS) * max(1, val["n"])))
        entry = {"epoch": epoch + 1, "loss": sum(losses) / len(losses),
                  "validation_exact": val["exact"], "validation_span_f1": val["span_f1"],
                  "validation_kinds": {k: v["f1"] for k, v in val["kinds"].items()},
                  "selection_objective": round(objective, 4),
                  "selection_views": view_scores,
                  "augmentation": dict(sorted(augmentation.items())),
                  "elapsed_s": round(time.perf_counter() - started, 1)}
        history.append(entry)
        save_json(out / "history.json", history)
        print(f"{args.arm} {args.mode} {args.seed} ep{epoch + 1} "
              f"val {val['exact']}/{val['n']} sel {objective:.3f} f1 {val['span_f1']} "
              f"{entry['elapsed_s']}s", flush=True)
        chosen = objective if args.select == "multi-view" else val["exact"] / max(1, val["n"])
        if chosen > best:
            best = chosen
            torch.save(cpu_state_dict(model), out / "best.pt")
    model.load_state_dict(torch.load(out / "best.pt", map_location="cpu", weights_only=True))
    results = {}
    for split in SCORE_SPLITS:
        scored, preds = evaluate(model, encoded[split], timed=split == "easy",
                                 decode_mode=args.decode, transition=transition)
        results[split] = scored
        (out / f"{split}.jsonl").write_text("".join(
            json.dumps(p, ensure_ascii=False) + "\n" for p in preds))
    # Purely positional probe, derived at evaluation time from `probe-base`: the
    # same atoms in a new arrangement, so unlike `probe-layout` it introduces no
    # unseen marker vocabulary. `heldout-layout` confounds the two, and round 1
    # showed that a zero there is ambiguous between "format" and "lexicon".
    # Gold is the parent record, which `augment` proves equal before returning.
    probe_rng = random.Random(args.seed + 20260915)
    base = data["probe-base"]
    views = [(augment(row, probe_rng) or row) for row in base]
    scored, preds = evaluate(model, [model.encode(row) for row in views])
    scored["derived_rearranged"] = sum(1 for row, view in zip(base, views) if view is not row)
    scored["derived_from"] = "probe-base re-arranged by corpus2.augment at eval time"
    results["probe-shuffle"] = scored
    (out / "probe-shuffle.jsonl").write_text("".join(
        json.dumps(p, ensure_ascii=False) + "\n" for p in preds))
    results["training_seconds"] = round(time.perf_counter() - started, 1)
    results["checkpoint_sha256"] = digest(out / "best.pt")
    # `best` is whichever objective was registered: a count for seen-layout
    # validation, a rate for the multi-view mean. Reporting both explicitly keeps a
    # round-4 and a round-5 history readable without pretending they are the same
    # number.
    results["selected_checkpoint_objective"] = {"mode": args.select,
                                                "value": round(best, 4)}
    results["selection_view_counts"] = selection_counts
    export(model, args, out, encoded, results)
    save_json(out / "results.json", results)
    print(json.dumps({split: {"exact": f"{results[split]['exact']}/{results[split]['n']}",
                               "span_f1": results[split]["span_f1"]}
                      for split in ("validation", "easy", "heldout-layout", "heldout-lexical",
                                    "heldout-value", "heldout-negative", "probe-base",
                                    "probe-layout", "probe-shuffle")}, indent=2), flush=True)


def export(model, args, out, encoded, results):
    """Evaluate the artifact that would actually ship, not the fp32 prototype."""
    # Dynamic INT8 is a CPU deployment format. The checkpoint is deliberately
    # portable and all artifact comparisons happen after the same fp32 model is
    # returned to CPU, regardless of where training ran.
    model.cpu().eval()
    fp32 = out / "model.fp32.pt"
    torch.save(model.state_dict(), fp32)
    quantized = torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)
    qpath = out / "model.int8.pt"
    torch.save(quantized.state_dict(), qpath)
    sizes = {"fp32_mib": round(fp32.stat().st_size / 2 ** 20, 2),
             "int8_mib": round(qpath.stat().st_size / 2 ** 20, 2)}
    rows_sample = encoded["easy"][:64]
    scored, _ = evaluate(quantized, rows_sample, timed=False)
    # Like-for-like: the quantized number only means anything next to the fp32
    # number on the same rows, which is the comparison the pilot had to add.
    baseline, _ = evaluate(model, rows_sample)
    # Quantization is scored on held-out formats too, not only on in-distribution
    # rows: a drop that only appears on unseen layouts would be invisible here
    # otherwise, and the shipped artifact is the one that has to answer for it.
    unseen, _ = evaluate(quantized, encoded["heldout-layout"])
    results["export"] = {**sizes, "int8_subset_n": scored["n"],
                         "int8_subset_exact": scored["exact"],
                         "fp32_subset_exact": baseline["exact"],
                         "int8_span_f1_subset": scored["span_f1"],
                         "fp32_span_f1_subset": baseline["span_f1"],
                         "int8_heldout_layout_exact": f"{unseen['exact']}/{unseen['n']}",
                         "fp32_heldout_layout_exact": f"{results['heldout-layout']['exact']}/"
                                                      f"{results['heldout-layout']['n']}",
                         "int8_latency_p50": None, "int8_latency_p95": None}
    with torch.inference_mode():
        latencies = [1000 * single_row_ms(quantized, item) for item in rows_sample]
    results["export"]["int8_latency_p50"] = round(statistics.median(latencies), 2)
    results["export"]["int8_latency_p95"] = round(sorted(latencies)[int(0.95 * len(latencies))], 2)


def score_artifact(args):
    """Score the *shipped* INT8 file on the full battery, under both decoders.

    Training writes `model.int8.pt`, but until now only a 64-row subset and
    heldout-layout were ever read from that file: the artifact comparisons in the
    documentation were re-quantizations of the checkpoint on the scoring machine,
    not the bytes that would ship. This loads the exported file and states explicitly
    whether it equals a fresh local re-quantization, so a reader knows which object
    produced each number.
    """
    out = Path(args.score_artifact)
    config = json.loads((out / "config.json").read_text())
    target = out / f"artifact-int8-{args.transition}.json"
    if target.exists():
        raise SystemExit(f"refusing to overwrite {target}")
    shipped_path = out / "model.int8.pt"
    if not shipped_path.exists():
        raise SystemExit(f"{out} has no exported artifact to score")
    torch.set_num_threads(args.threads)
    mode = config["args"]["mode"]

    def quantized_model():
        # Dynamic INT8 has no state_dict-compatible loader for an unquantized module:
        # quantize first, then load the shipped weights into the same structure.
        candidate = torch.quantization.quantize_dynamic(Extractor(mode, len(LABELS)),
                                                       {nn.Linear}, dtype=torch.qint8)
        candidate.load_state_dict(torch.load(shipped_path, map_location="cpu",
                                             weights_only=False))
        return candidate.eval()

    def decode_tags(network, items, decode_mode="greedy", transition=None):
        """The tag sequences a model emits for `items`, for behaviour comparison."""
        network.eval()
        tags = []
        with torch.no_grad():
            for offset in range(0, len(items), 64):
                chunk = items[offset:offset + 64]
                batch, _ = collate(chunk, network.encoder.config.pad_token_id)
                logits, _ = network(batch)
                tags += tag_paths(logits.detach().cpu(), [len(item[1]) for item in chunk],
                                  decode_mode, transition)
        return tags

    def checkpoint():
        fresh = Extractor(mode, len(LABELS))
        fresh.load_state_dict(torch.load(out / "best.pt", map_location="cpu",
                                         weights_only=True))
        return fresh.eval()

    shipped = quantized_model()
    reference = checkpoint()
    # A quantized `state_dict` holds packed params (tuples, dtypes), so comparing its
    # internals compares a serialization format. What a reader needs to know is
    # whether the shipped bytes *behave* like a fresh re-quantization of this
    # checkpoint, so the equality claim is made about predictions.
    sample = [reference.encode(row) for row in rows("easy")[:32]]
    shipped_tags = decode_tags(shipped, sample)
    reference_tags = decode_tags(reference, sample)
    requant_tags = decode_tags(torch.quantization.quantize_dynamic(checkpoint(),
                                                                   {nn.Linear},
                                                                   dtype=torch.qint8),
                              sample)
    train_items = [reference.encode(row)
                   for split in ARM_SPLITS[config["args"]["arm"]] for row in rows(split)]
    transition = (legality_matrix() if args.transition == "legal"
                  else transitions(train_items))
    encoded = {split: [reference.encode(row) for row in rows(split)]
               for split in SCORE_SPLITS}
    # Same seed formula and row order as the trainer, so `probe-shuffle` here is the
    # same measurement the fp32 rescores report.
    probe_rng = random.Random(config["args"]["seed"] + 20260915)
    base = rows("probe-base")
    views = [(augment(row, probe_rng) or row) for row in base]
    encoded["probe-shuffle"] = [reference.encode(row) for row in views]
    result = {"checkpoint_sha256": digest(out / "best.pt"),
              "artifact_sha256": digest(shipped_path),
              "artifact_mib": round(shipped_path.stat().st_size / 2 ** 20, 2),
              "shipped_matches_requantization": shipped_tags == requant_tags,
              "shipped_matches_fp32_on_sample": shipped_tags == reference_tags,
              "sample_rows": len(sample), "transition": args.transition, "decodes": {}}
    for decode_mode in ("greedy", "viterbi"):
        result["decodes"][decode_mode] = {
            split: evaluate(shipped, items, decode_mode=decode_mode,
                            transition=transition)[0]
            for split, items in encoded.items()}
    save_json(target, result)
    print(json.dumps({"artifact_sha256": result["artifact_sha256"][:12],
                      "shipped_matches_requantization":
                          result["shipped_matches_requantization"],
                      "shipped_matches_fp32_on_sample":
                          result["shipped_matches_fp32_on_sample"],
                      "greedy": {split: f"{scores['exact']}/{scores['n']}" for split, scores
                                 in result["decodes"]["greedy"].items()},
                      "viterbi": {split: f"{scores['exact']}/{scores['n']}" for split, scores
                                  in result["decodes"]["viterbi"].items()
                                  if split in ("validation", "heldout-layout",
                                               "heldout-lexical", "probe-value")}}, indent=2))


def rescore(args):
    """Re-score an existing checkpoint with a different decoder: no retraining, so
    the comparison isolates decoding from everything else. Also checks that the
    greedy path reproduces the run's recorded numbers before trusting the delta."""
    out = Path(args.rescore)
    config = json.loads((out / "config.json").read_text())
    target = out / f"rescore-{args.transition}.json"
    if target.exists():
        raise SystemExit(f"refusing to overwrite {target}")
    torch.set_num_threads(args.threads)
    device = resolve_device(args.device)
    model = Extractor(config["args"]["mode"], len(LABELS))
    model.load_state_dict(torch.load(out / "best.pt", map_location="cpu", weights_only=True))
    model.to(device)
    train_items = [model.encode(row)
                   for split in ARM_SPLITS[config["args"]["arm"]] for row in rows(split)]
    transition = (legality_matrix() if args.transition == "legal"
                  else transitions(train_items))
    encoded = {split: [model.encode(row) for row in rows(split)] for split in SCORE_SPLITS}
    # The positional probe added after round 1, computed here so a checkpoint
    # trained before it exists can still be compared on it. Same seed formula and
    # same row order as the trainer, so the numbers are the same measurement.
    probe_rng = random.Random(config["args"]["seed"] + 20260915)
    base = rows("probe-base")
    views = [(augment(row, probe_rng) or row) for row in base]
    shuffle_items = [model.encode(row) for row in views]
    result = {"checkpoint_sha256": digest(out / "best.pt"),
              "args": config["args"],
              "transition": args.transition,
              "probe_shuffle_rearranged": sum(1 for row, view in zip(base, views)
                                              if view is not row),
              "decodes": {}}
    for mode in ("greedy", "viterbi"):
        result["decodes"][mode] = {
            split: evaluate(model, encoded[split],
                            decode_mode=mode, transition=transition)[0]
            for split in SCORE_SPLITS}
        result["decodes"][mode]["probe-shuffle"] = evaluate(
            model, shuffle_items, decode_mode=mode, transition=transition)[0]
    recorded = json.loads((out / "results.json").read_text())
    result["greedy_matches_recorded"] = all(
        result["decodes"]["greedy"][split]["exact"] == recorded[split]["exact"]
        for split in SCORE_SPLITS)
    result["delta_viterbi_minus_greedy_exact"] = {
        split: result["decodes"]["viterbi"][split]["exact"]
        - result["decodes"]["greedy"][split]["exact"] for split in SCORE_SPLITS}
    save_json(target, result)
    print(json.dumps({"greedy_matches_recorded": result["greedy_matches_recorded"],
                      "probe_shuffle_rearranged": result["probe_shuffle_rearranged"],
                      "probe_shuffle": {mode: f"{scores['probe-shuffle']['exact']}/"
                                              f"{scores['probe-shuffle']['n']}"
                                        for mode, scores in result["decodes"].items()},
                      "delta": result["delta_viterbi_minus_greedy_exact"]}, indent=2))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--arm", choices=["A-narrow", "B-broad"], default="B-broad")
    parser.add_argument("--mode", choices=["bytes", "characters", "pretrained"], default="characters")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--epochs", type=int, default=6)
    parser.add_argument("--batch", type=int, default=32)
    parser.add_argument("--threads", type=int, default=3)
    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
    parser.add_argument("--decode", choices=["greedy", "viterbi"], default="greedy")
    parser.add_argument("--transition", choices=["legal", "empirical"], default="legal",
                        help="which matrix Viterbi uses: BIO admissibility only "
                             "(default) or empirical tag-transition frequencies")
    parser.add_argument("--augment", choices=["none", "shuffle", "shuffle-values",
                                               "shuffle-paths", "shuffle-rich", "balanced4",
                                               "balanced", "balanced-dose", "shuffle-markers",
                                               "shuffle-boundary"],
                        default="none",
                        help="shuffle atom arrangement per epoch; breaks positional cues. "
                             "`balanced` is the current six-view schedule; `balanced4` is "
                             "round 4's frozen four-view list")
    parser.add_argument("--select", choices=["validation", "multi-view"], default="validation",
                        help="checkpoint objective: seen-layout validation exact-record "
                             "(rounds 1-4) or the mean over fixed derived validation views")
    parser.add_argument("--augment-rate", type=float, default=0.5,
                        help="share of training rows re-arranged each epoch")
    parser.add_argument("--rescore", help="re-score a run directory with both decoders")
    parser.add_argument("--score-artifact",
                        help="score a run's exported model.int8.pt on every split under "
                             "both decoders, and report whether it matches a fresh local "
                             "re-quantization of the checkpoint")
    parser.add_argument("--smoke", action="store_true",
                        help="tiny subsets to check plumbing; never a measurement")
    parser.add_argument("--out")
    args = parser.parse_args()
    if args.rescore:
        return rescore(args)
    if args.score_artifact:
        return score_artifact(args)
    if not args.out:
        parser.error("--out is required unless --rescore or --score-artifact is given")
    train(args)


if __name__ == "__main__":
    main()