File size: 29,069 Bytes

c83c617

"""V4 inference bundle, FP32/INT8 Core ML packages and gates.

Stateless graph: int32 [1,T], 1<=T<=128 -> raw float32 [1,T,8192] logits.
KV graph: one token at an absolute position with explicit FP32 caches
(7,1,6,128,64) from `vimeml.deployment.v4.KVStepGraphV4`. The caller owns
snapshots and rollback. V1/V2 formats and entry points remain unchanged.
"""
import copy
import json
import shutil
import time
from pathlib import Path

import numpy as np
import sentencepiece as spm
import torch

from vimeml.benchmarks.comparison import paired
from vimeml.deployment import v2
from vimeml.deployment.bundle import environment, fresh_directory, read_json, write_json
from vimeml.deployment.coreml import CoreMLForward, coremltools, inspect_spec
from vimeml.deployment.graph import trace_graph
from vimeml.deployment.v4 import CACHE_SHAPE, trace_kv_graph
from vimeml.deployment.validation import compare_rows, numeric_error
from vimeml.training.data import file_sha
from vimeml.training.evaluate_ime import common_prefix_length
from vimeml.training.infer import ROOT
from vimeml.training.infer_v4 import JapaneseLMV4
from vimeml.training.model_v4 import GPTV4Config, TinyGPTV4

BUNDLE_FORMAT = "vimeml_inference_bundle_v4"
PACKAGE_FORMAT = "vimeml_coreml_v4"
KV_FORMAT = "vimeml_coreml_v4_kv_v1"
ARCHITECTURE = "tiny_gpt_v4"
REFERENCE_FORMAT = "vimeml_v4_mac_preparation_reference_v1"
VOCAB = 8192
TOLERANCE = 3e-4
QUANTILES = (0.5, 0.9, 0.99, 0.999)


def export_bundle(checkpoint, tokenizer_dir, output):
    checkpoint, tokenizer_dir, output = map(Path, (checkpoint, tokenizer_dir, output))
    if output.exists():
        raise ValueError("Bundle output already exists.")
    saved = torch.load(checkpoint, map_location="cpu", weights_only=True)
    if (saved.get("format") != "vimeml_tiny_gpt_v4" or saved.get("architecture") != ARCHITECTURE
            or not saved.get("deployment_only") or not saved.get("source_checkpoint_sha256")):
        raise ValueError("Expected a V4 deployment checkpoint with source provenance.")
    # JapaneseLMV4 binds token data, tokenizer and checkpoint tokenizer identity.
    lm = JapaneseLMV4(checkpoint, tokenizer_dir, "cpu")
    state = saved["model"]
    if not torch.equal(state["lm_head.weight"], state["token_embedding.weight"]):
        raise ValueError("Expected identical tied embedding/head tensors.")
    if any(not torch.isfinite(p).all() for p in lm.model.parameters()):
        raise ValueError("Nonfinite learned weight.")
    token_manifest = ROOT / saved["config"]["token_dir"] / "manifest.json"
    weights = {name: value.detach().cpu().float().contiguous() for name, value in state.items()
               if name != "lm_head.weight"}
    fresh_directory(output)
    torch.save(weights, output / "weights.pt")
    write_json(output / "config.json", lm.model.configuration())
    shutil.copyfile(tokenizer_dir / "tokenizer.model", output / "tokenizer.model")
    shutil.copyfile(token_manifest, output / "token-manifest.json")
    result = {"format": BUNDLE_FORMAT, "status": "complete", "architecture": ARCHITECTURE,
              "precision": "fp32", "parameter_count": lm.model.parameter_count(), "special_ids": lm.special,
              "forbidden_generation_ids": list(lm.forbidden),
              "checkpoint_sha256": saved["source_checkpoint_sha256"],
              "deployment_checkpoint_sha256": file_sha(checkpoint),
              "checkpoint_step": saved["step"], "checkpoint_epoch_index": saved["epoch"],
              "training_signatures": saved["signatures"], "tokenizer_identity": saved["tokenizer_identity"],
              "tied_lm_head": True, "files": v2.inventory(output, hashes=True), "environment": environment(),
              "policy": "Unique FP32 inference weights; no optimizer/RNG. Joint suffix logP sum; no EOS/truncation."}
    write_json(output / "manifest.json", result)
    return result


def verify_bundle(path, hashes=True):
    path = Path(path)
    manifest = read_json(path / "manifest.json")
    if (manifest.get("format") != BUNDLE_FORMAT or manifest.get("status") != "complete"
            or manifest.get("architecture") != ARCHITECTURE or not manifest.get("tied_lm_head")):
        raise ValueError("Unsupported/incomplete V4 bundle.")
    if set(manifest["files"]) != {"weights.pt", "config.json", "tokenizer.model", "token-manifest.json"}:
        raise ValueError("Unexpected inference file list.")
    for name, entry in manifest["files"].items():
        if (path / name).stat().st_size != entry["bytes"] or (hashes and file_sha(path / name) != entry["sha256"]):
            raise ValueError(f"Bundle checksum mismatch: {name}")
    return manifest


class BundleLMV4(JapaneseLMV4):
    """PyTorch CPU FP32 model from a V4 bundle; V4 prefix and forbidden-ID rules."""
    def __init__(self, path):
        path = Path(path)
        manifest = verify_bundle(path, hashes=False)
        self.device = torch.device("cpu")
        self.processor = spm.SentencePieceProcessor(model_file=str(path / "tokenizer.model"))
        self.special = {name: getattr(self.processor, f"{name}_id")() for name in ("pad", "unk", "bos", "eos")}
        self.model = TinyGPTV4(GPTV4Config(**read_json(path / "config.json"))).eval()
        state = torch.load(path / "weights.pt", map_location="cpu", weights_only=True)
        if "lm_head.weight" in state:
            raise ValueError("Shared embedding must be stored only once.")
        loaded = self.model.load_state_dict(state, strict=False)
        if loaded.missing_keys != ["lm_head.weight"] or loaded.unexpected_keys:
            raise ValueError(f"Unexpected inference weights: {loaded}")
        if (self.special != manifest["special_ids"] or self.processor.vocab_size() != VOCAB
                or self.model.parameter_count() != manifest["parameter_count"]):
            raise ValueError("Bundle model/tokenizer contract mismatch.")
        self.forbidden = tuple(manifest["forbidden_generation_ids"])
        self.metadata = {"architecture": ARCHITECTURE, "checkpoint_sha256": manifest["checkpoint_sha256"],
                         "checkpoint_step": manifest["checkpoint_step"],
                         "deployment_checkpoint_sha256": manifest["deployment_checkpoint_sha256"],
                         "tokenizer_sha256": manifest["files"]["tokenizer.model"]["sha256"],
                         "bundle_manifest_sha256": file_sha(path / "manifest.json"),
                         "model": self.model.configuration(), "precision": "fp32", "device": "cpu",
                         "environment": environment()}


def finish(output, model, fmt, metadata):
    output = Path(output)
    fresh_directory(output)
    model.save(str(output / "model.mlpackage"))
    files = v2.inventory(output / "model.mlpackage")
    write_json(output / "manifest.json", {"format": fmt, "status": "complete", "architecture": ARCHITECTURE,
        **metadata, "environment": environment(), "spec": inspect_spec(model), "package_files": files,
        "package_bytes": sum(entry["bytes"] for entry in files.values()),
        "note": "Logical package bytes; compiled size, resident memory and device footprint are separate."})


def verify_package(path, fmt=PACKAGE_FORMAT):
    path = Path(path)
    manifest = read_json(path / "manifest.json")
    if (manifest.get("format") != fmt or manifest.get("status") != "complete"
            or manifest.get("architecture") != ARCHITECTURE):
        raise ValueError("Incomplete/unsupported V4 Core ML package.")
    if v2.inventory(path / "model.mlpackage") != manifest["package_files"]:
        raise ValueError("Package file inventory changed.")
    return manifest


def convert(bundle, output):
    ct = coremltools()
    if Path(output).exists():
        raise ValueError("Conversion output exists.")
    lm = BundleLMV4(bundle)
    model = ct.convert(trace_graph(lm.model), source="pytorch", convert_to="mlprogram",
        minimum_deployment_target=ct.target.iOS18, compute_precision=ct.precision.FLOAT32,
        inputs=[ct.TensorType(name="input_ids", shape=(1, ct.RangeDim(lower_bound=1, upper_bound=128, default=16)),
                              dtype=np.int32)],
        outputs=[ct.TensorType(name="logits", dtype=np.float32)], skip_model_load=True)
    model.short_description = "VimeML V4.2 RMSNorm/SwiGLU 8K; FP32 computation; no KV cache."
    model.user_defined_metadata["bundle_manifest_sha256"] = lm.metadata["bundle_manifest_sha256"]
    finish(output, model, PACKAGE_FORMAT, {"kind": "fp32", "minimum_ios": 18, "bundle": lm.metadata,
        "interface": {"input_ids": "int32 [1,T], 1<=T<=128; right PAD",
                      "logits": "float32 [1,T,8192]; full vocabulary; no softmax"},
        "precision_policy": "All learned parameters and computation FP32; explicit causal attention.",
        "conversion_source_sha256": file_sha(Path(__file__))})


def convert_kv(bundle, output):
    ct = coremltools()
    if Path(output).exists():
        raise ValueError("Conversion output exists.")
    lm = BundleLMV4(bundle)
    model = ct.convert(trace_kv_graph(lm.model), source="pytorch", convert_to="mlprogram",
        minimum_deployment_target=ct.target.iOS18, compute_precision=ct.precision.FLOAT32,
        inputs=[ct.TensorType(name="input_ids", shape=(1, 1), dtype=np.int32),
                ct.TensorType(name="position", shape=(1,), dtype=np.int32),
                ct.TensorType(name="key_cache", shape=CACHE_SHAPE, dtype=np.float32),
                ct.TensorType(name="value_cache", shape=CACHE_SHAPE, dtype=np.float32)],
        outputs=[ct.TensorType(name="logits", dtype=np.float32),
                 ct.TensorType(name="new_key_cache", dtype=np.float32),
                 ct.TensorType(name="new_value_cache", dtype=np.float32)], skip_model_load=True)
    model.short_description = "VimeML V4.2 single-token KV step; FP32 computation and caches."
    model.user_defined_metadata["bundle_manifest_sha256"] = lm.metadata["bundle_manifest_sha256"]
    finish(output, model, KV_FORMAT, {"kind": "fp32", "minimum_ios": 18, "bundle": lm.metadata,
        "cache_shape": list(CACHE_SHAPE), "cache_dtype": "float32", "compute_units": "CPU_ONLY",
        "cache_bytes_per_snapshot": 2 * int(np.prod(CACHE_SHAPE)) * 4,
        "interface": {"input_ids": "int32 [1,1]", "position": "int32 [1], 0<=position<=127; absolute slot",
                      "key_cache/value_cache": "float32 (7,1,6,128,64); slots >= position are ignored",
                      "logits": "float32 [1,1,8192] for the new token; no softmax",
                      "new_key_cache/new_value_cache": "inputs with slot `position` replaced"},
        "precision_policy": "All learned parameters, computation and caches FP32.",
        "policy": "Batch1; caller owns immutable snapshots and rollback; no sliding or cross-request retention.",
        "conversion_source_sha256": file_sha(Path(__file__))})


def compress(source, output, alignment, kv=False):
    ct = coremltools()
    source, output, alignment = map(Path, (source, output, alignment))
    if output.exists():
        raise ValueError("Compression output exists.")
    fmt = KV_FORMAT if kv else PACKAGE_FORMAT
    original = verify_package(source, fmt)
    gate = read_json(alignment)
    expected = "vimeml_kv_alignment_v4" if kv else "vimeml_alignment_v4"
    if (original["kind"] != "fp32" or gate.get("format") != expected or not gate.get("passed")
            or gate.get("coreml_manifest_sha256") != file_sha(source / "manifest.json")):
        raise ValueError("First pass alignment for this exact uncompressed V4 FP32 package.")
    from coremltools.optimize import coreml as opt
    model = ct.models.MLModel(str(source / "model.mlpackage"), skip_model_load=True)
    # Identical to the deployed V2.1 recipe.
    config = opt.OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8",
                                        granularity="per_block", block_size=32, weight_threshold=2048)
    optimization, selection = v2.parameter_selection(opt, model, config)
    print(f"Selected {len(selection['selected_constants'])} parameter matrices; "
          f"excluded {len(selection['excluded_constants'])} structural constants.", flush=True)
    result = opt.linear_quantize_weights(model, config=optimization)
    if not any(name.startswith("constexpr_") for name in inspect_spec(result)["operations"]):
        raise ValueError("No compressed weights found.")
    carried = {key: original[key] for key in original if key in (
        "minimum_ios", "bundle", "interface", "cache_shape", "cache_dtype", "compute_units",
        "cache_bytes_per_snapshot", "policy")}
    finish(output, result, fmt, {**carried, "kind": "linear8_fp32_compute",
        "source_manifest_sha256": file_sha(source / "manifest.json"),
        "fp32_alignment_sha256": file_sha(alignment), "conversion_source_sha256": file_sha(Path(__file__)),
        "compression": {"method": "linear_symmetric", "bits": 8, "block_size": 32, "granularity": "per_block",
                        "weight_threshold": 2048, "selection": selection},
        "precision_policy": "INT8 parameter matrices, FP32 compute; norms/masks unchanged; no activation quantization."})


class CoreMLLMV4(BundleLMV4):
    def __init__(self, bundle, directory):
        super().__init__(bundle)
        manifest = verify_package(directory)
        if manifest["bundle"]["bundle_manifest_sha256"] != self.metadata["bundle_manifest_sha256"]:
            raise ValueError("Core ML package belongs to another bundle.")
        self.model = CoreMLForward(directory, self.model.config, "CPU_ONLY")
        self.metadata.update(precision=manifest["kind"], compute_units="CPU_ONLY",
                             coreml_manifest_sha256=file_sha(Path(directory) / "manifest.json"))


def empty_cache():
    return tuple(np.zeros(CACHE_SHAPE, dtype=np.float32) for _ in range(2))


class KVRuntimeV4:
    """Single-token Core ML KV steps. Inputs are never modified; outputs are owned copies."""
    def __init__(self, bundle, directory):
        lm = BundleLMV4(bundle)
        self.config, self.metadata = lm.model.config, dict(lm.metadata)
        self.manifest = verify_package(directory, KV_FORMAT)
        if self.manifest["bundle"]["bundle_manifest_sha256"] != self.metadata["bundle_manifest_sha256"]:
            raise ValueError("KV model belongs to another bundle.")
        if (self.manifest["cache_shape"] != list(CACHE_SHAPE) or self.manifest["cache_dtype"] != "float32"
                or self.manifest["compute_units"] != "CPU_ONLY"):
            raise ValueError("KV package cache/compute contract differs.")
        ct = coremltools()
        self.model = ct.models.MLModel(str(Path(directory) / "model.mlpackage"),
                                       compute_units=ct.ComputeUnit.CPU_ONLY)
        self.metadata.update(precision=self.manifest["kind"], compute_units="CPU_ONLY",
                             coreml_manifest_sha256=file_sha(Path(directory) / "manifest.json"))
        self.calls = 0

    def step(self, token, position, state):
        if (type(position) is not int or not 0 <= position < 128
                or not 0 <= int(token) < VOCAB or len(state) != 2
                or any(not isinstance(a, np.ndarray) or a.shape != CACHE_SHAPE or a.dtype != np.float32
                       for a in state)):
            raise ValueError("Invalid KV step input.")
        result = self.model.predict({"input_ids": np.array([[token]], np.int32),
                                     "position": np.array([position], np.int32),
                                     "key_cache": state[0], "value_cache": state[1]})
        self.calls += 1
        logits = np.array(result["logits"], dtype=np.float32, copy=True)
        caches = tuple(np.array(result[name], dtype=np.float32, copy=True)
                       for name in ("new_key_cache", "new_value_cache"))
        if (logits.shape != (1, 1, VOCAB) or not np.isfinite(logits).all()
                or any(a.shape != CACHE_SHAPE for a in caches)):
            raise ValueError("Invalid Core ML cache/logits output.")
        return logits, caches

    def run(self, ids, state=None, start=0):
        """Process tokens one at a time from `start`; returns [1,Q,V] logits and the final snapshot."""
        ids = [int(i) for i in np.asarray(ids).reshape(-1)]
        if not ids or start < 0 or start + len(ids) > 128 or (state is None and start != 0):
            raise ValueError("Invalid KV position range.")
        state = empty_cache() if state is None else state
        parts = []
        for offset, token in enumerate(ids):
            logits, state = self.step(token, start + offset, state)
            parts.append(logits)
        return np.concatenate(parts, axis=1), state


class CachedForwardV4:
    """Evaluation adapter: share the exact common token prefix within one batch.

    Rows arrive right-padded with PAD (id 0, never produced by encode). Trailing
    PAD positions are not computed; their logits are zero and never scored.
    """
    def __init__(self, runtime, pad):
        self.runtime, self.config, self.pad = runtime, runtime.config, pad

    def __call__(self, inputs):
        rows = [list(row) for row in inputs.numpy().astype(np.int64)]
        width = len(rows[0])
        for row in rows:
            while len(row) > 1 and row[-1] == self.pad:
                row.pop()
        common = common_prefix_length(rows)
        prefix = state = None
        if common:
            prefix, state = self.runtime.run(rows[0][:common])
        output = np.zeros((len(rows), width, VOCAB), dtype=np.float32)
        for index, row in enumerate(rows):
            if common == len(row):
                logits = prefix
            elif common:
                suffix, _ = self.runtime.run(row[common:], state, common)
                logits = np.concatenate((prefix, suffix), axis=1)
            else:
                logits, _ = self.runtime.run(row)
            output[index, :len(row)] = logits[0]
        return torch.from_numpy(output)


class CachedLMV4(BundleLMV4):
    def __init__(self, bundle, directory):
        super().__init__(bundle)
        self.runtime = KVRuntimeV4(bundle, directory)
        self.model = CachedForwardV4(self.runtime, self.special["pad"])
        self.metadata = self.runtime.metadata


def error_distribution(expected, actual, atol=TOLERANCE, rtol=TOLERANCE):
    """Strict error plus quantiles and per-position argmax agreement."""
    result = numeric_error(expected, actual, atol, rtol)
    delta = np.abs(expected.astype(np.float64) - actual.astype(np.float64))
    result["quantiles_abs"] = {str(q): float(np.quantile(delta, q)) for q in QUANTILES}
    result["outside_fraction"] = result["outside_tolerance"] / result["values"]
    top_e, top_a = expected.argmax(-1), actual.argmax(-1)
    result["positions"] = int(top_e.size)
    result["argmax_equal_positions"] = int((top_e == top_a).sum())
    return result


def read_reference(reference, lm):
    reference = Path(reference)
    info = read_json(reference / "reference.json")
    if (info.get("format") != REFERENCE_FORMAT or info["checkpoint_step"] != lm.metadata["checkpoint_step"]
            or info["checkpoint_sha256"] != lm.metadata["checkpoint_sha256"]
            or info["tokenizer_sha256"] != lm.metadata["tokenizer_sha256"]):
        raise ValueError("Reference belongs to another checkpoint/tokenizer.")
    with np.load(reference / "logits.npz", allow_pickle=False) as arrays:
        return info, {name: arrays[name] for name in arrays.files}


def summarize_errors(items):
    deltas = [item["max_abs"] for item in items]
    return {"checks": len(items), "max_abs": max(deltas), "failed_checks": sum(
        item["outside_tolerance"] > 0 for item in items)}


@torch.inference_mode()
def align(lm, reference, output):
    """Stateless fixture gate: valid prefixes within 3e-4 and causality invariants."""
    output = Path(output)
    if output.exists():
        raise ValueError("Alignment output exists.")
    info, arrays = read_reference(reference, lm)
    errors, actuals = [], {}
    for name, case in info["cases"].items():
        ids = arrays[name + "_input_ids"]
        actual = lm.model(torch.from_numpy(ids.copy())).numpy()
        actuals[name] = actual
        valid = case["valid_prefix"]
        expected = arrays[name + "_logits"]
        error = error_distribution(expected[:, :valid], actual[:, :valid])
        error.update(id=name, length=case["length"], valid_prefix=valid,
                     last_top1_equal=int(actual[0, valid - 1].argmax()) == int(expected[0, valid - 1].argmax()))
        errors.append(error)
    invariants = {name: numeric_error(actuals["short"], actuals[name][:, :16], TOLERANCE, TOLERANCE)
                  for name in ("right_pad", "changed_future")}
    report = {"format": "vimeml_alignment_v4",
        "passed": all(item["outside_tolerance"] == 0 for item in [*errors, *invariants.values()]),
        "model": lm.metadata, "coreml_manifest_sha256": lm.metadata.get("coreml_manifest_sha256"),
        "reference": info, "tolerances": {"atol": TOLERANCE, "rtol": TOLERANCE},
        "summary": summarize_errors(errors), "logits": errors, "invariants": invariants,
        "note": "Strict PyTorch FP32 fixture comparison. Quantized ranking quality is evaluated separately; failures are preserved."}
    fresh_directory(output)
    write_json(output / "alignment.json", report)
    return report


@torch.inference_mode()
def align_kv(runtime, bundle, reference, output, stateless=None):
    """KV gate against frozen PyTorch fixtures, or against a same-weight stateless package.

    one_shot: uninterrupted single-token prefill vs full-sequence logits.
    one_token: the same steps vs PyTorch `forward_cached` single-token logits.
    split_chunk: prefill resumed from copied intermediate snapshots.
    Branches: shared immutable snapshot, A/B/A rollback, dirty unused tail.
    """
    output = Path(output)
    if output.exists():
        raise ValueError("Alignment output exists.")
    lm = BundleLMV4(bundle)
    info, arrays = read_reference(reference, lm)
    old = None
    if stateless:
        old = CoreMLLMV4(bundle, stateless)
        if old.metadata["precision"] != runtime.metadata["precision"]:
            raise ValueError("Compare KV with the stateless package of the same precision.")
    full = (lambda ids: old.model(torch.tensor([list(ids)])).numpy()) if old else (
        lambda ids: lm.model(torch.tensor([list(ids)])).numpy())
    errors, invariants, saved = [], [], {}
    for name, case in info["cases"].items():
        ids = [int(i) for i in arrays[name + "_input_ids"][0]]
        valid = case["valid_prefix"]
        expected = full(ids) if old else arrays[name + "_logits"]
        actual, _ = runtime.run(ids)
        saved[name] = actual
        errors.append({"case": name, "check": "one_shot", **error_distribution(expected[:, :valid], actual[:, :valid])})
        if name == "long" and not old:
            errors.append({"case": name, "check": "one_token",
                           **error_distribution(arrays["long_kv_step_logits"], actual)})
        if len(ids) > 3:
            for split in (3, len(ids) // 2):
                first, state = runtime.run(ids[:split])
                snapshot = tuple(a.copy() for a in state)
                second, _ = runtime.run(ids[split:], snapshot, split)
                chunked = np.concatenate((first, second), axis=1)
                errors.append({"case": name, "check": "split_chunk", "chunks": [split, len(ids) - split],
                               **error_distribution(expected[:, :valid], chunked[:, :valid])})
                invariants.append({"case": name, "check": "split_equals_one_shot", "chunks": [split, len(ids) - split],
                                   **numeric_error(actual, chunked, 0, 0)})
    for name in ("right_pad", "changed_future"):
        invariants.append({"case": name, "check": "valid_prefix_vs_short",
                           **numeric_error(saved["short"], saved[name][:, :16], TOLERANCE, TOLERANCE)})
    prefix = [int(i) for i in arrays["short_input_ids"][0][:7]]
    _, state = runtime.run(prefix)
    original = tuple(a.copy() for a in state)
    branches = []
    for suffix in ([9], [47, 81, 102], [9]):
        actual, _ = runtime.run(suffix, state, len(prefix))
        expected = full([*prefix, *suffix])[:, len(prefix):]
        errors.append({"check": "branch", "suffix": suffix, **error_distribution(expected, actual)})
        branches.append(actual)
    immutable = all(np.array_equal(a, b) for a, b in zip(state, original, strict=True))
    invariants.append({"check": "rollback_A_B_A", **numeric_error(branches[0], branches[2], 0, 0)})
    dirty = tuple(a.copy() for a in state)
    for a in dirty:
        a[:, :, :, len(prefix):, :] = 999
    clean_logits, _ = runtime.run([9], state, len(prefix))
    dirty_logits, _ = runtime.run([9], dirty, len(prefix))
    invariants.append({"check": "unused_tail", **numeric_error(clean_logits, dirty_logits, 0, 0)})
    report = {"format": "vimeml_kv_alignment_v4",
        "passed": immutable and all(e["outside_tolerance"] == 0 for e in errors + invariants),
        "coreml_manifest_sha256": runtime.metadata["coreml_manifest_sha256"], "model": runtime.metadata,
        "comparison": "same-weight stateless Core ML" if old else "frozen PyTorch FP32 fixtures",
        "baseline_manifest_sha256": old.metadata["coreml_manifest_sha256"] if old else None,
        "tolerances": {"atol": TOLERANCE, "rtol": TOLERANCE}, "immutable_branch_inputs": immutable,
        "summary": summarize_errors(errors), "errors": errors, "invariants": invariants,
        "kv_calls": runtime.calls}
    fresh_directory(output)
    write_json(output / "alignment.json", report)
    return report


def top1(rows):
    return sum(bool(r["orders"]["lm_context_sum"]) and r["orders"]["lm_context_sum"][0] in r["answers"] for r in rows)


def evaluate_tracks(lm, output, reference=None, reference_label=None):
    """Score all four frozen pools; optionally compare per case with a reference run."""
    from vimeml.benchmarks.evaluate_ajimee import rerank, summarize
    from vimeml.posttraining.infer import TRACKS, read_snapshot
    output = Path(output)
    if output.exists():
        raise ValueError("Evaluation output exists.")
    frozen = {track: read_snapshot(track) for track in TRACKS}
    fresh_directory(output)
    summary = {"format": "vimeml_v4_track_evaluation_v1", "model": lm.metadata,
               "reference": str(reference) if reference else None, "reference_label": reference_label,
               "tracks": {}, "blind_read": False, "test_read": False,
               "semantics": "full-vocabulary suffix logP sum; joint context+candidate tokenization; no final EOS, truncation, or tuning; stable ties"}
    for track, original in frozen.items():
        rows = copy.deepcopy(original)
        started = time.perf_counter()
        rerank(lm, rows)
        elapsed = time.perf_counter() - started
        destination = output / track
        destination.mkdir()
        (destination / "scores.jsonl").write_text("".join(
            json.dumps(r, ensure_ascii=False, allow_nan=False) + "\n" for r in rows), encoding="utf-8")
        result = {"cases": len(rows), "top1": top1(rows), "elapsed_seconds": elapsed,
                  "metrics": summarize(rows)}
        if reference:
            before = [json.loads(line) for line in (Path(reference) / track / "scores.jsonl")
                      .read_text(encoding="utf-8").splitlines() if line.strip()]
            orders = compare_rows(before, rows)
            pairs = paired(before, rows)
            write_json(destination / "comparison.json", {"reference_label": reference_label,
                                                        "orders": orders, "paired": pairs})
            result["comparison"] = {"reference_top1": top1(before),
                **{k: v for k, v in orders.items() if k != "changes"},
                "order_changed_ids": [c["id"] for c in orders["changes"]],
                "top1_changed_ids": [c["id"] for c in orders["changes"] if c["top1_changed"]],
                "paired": pairs["paired"],
                "improved_ids": [c["id"] for c in pairs["changed_cases"] if c["change"] == "improved"],
                "regressed_ids": [c["id"] for c in pairs["changed_cases"] if c["change"] == "regressed"]}
        write_json(destination / "metrics.json", result)
        summary["tracks"][track] = {k: v for k, v in result.items() if k != "metrics"}
        print(json.dumps({"track": track, "top1": result["top1"], "cases": len(rows),
                          "comparison": {k: result["comparison"][k] for k in (
                              "reference_top1", "order_changed_cases", "top1_changed_cases", "max_score_sum_abs")}
                          if reference else None}, ensure_ascii=False), flush=True)
    for track in TRACKS:
        read_snapshot(track)
    write_json(output / "summary.json", summary)
    return summary