Exportar el encoder: dos salidas, tres precisiones y una model card

The encoder travels the same road with a different wrapper: ``EncoderHeads`` returns the two
outputs the demo needs (``value`` and ``blunder``) and nothing else, the ``squares`` scheme has
no dynamic sequence axis at all (it is always 69 tokens) and the metadata carries ``rukh_kind``
and ``rukh_heads`` so a file can say what it is. Everything else — the exporter, the fallback,
the verification, the metadata, the sidecar — is the same code: the decoder's path is untouched.

MODEL_NAME = "model.onnx"
INPUT_NAME = "idx"
OUTPUT_NAME = "logits"
VALUE_OUTPUT = "value"
BLUNDER_OUTPUT = "blunder"
ENCODER_OUTPUTS = (VALUE_OUTPUT, BLUNDER_OUTPUT)
"""What the encoder graph returns: the evaluation bar and the blunder alert of the demo."""
BATCH_AXIS = "batch"
SEQUENCE_AXIS = "sequence"
DEFAULT_OPSET = 18
METADATA_PREFIX = "rukh_"

class EncoderHeads(nn.Module):
    """Wraps a fine-tuned encoder so that ``forward(idx)`` returns ``(value, blunder)``.

    The third head (``result``) is left out on purpose: the demo draws an evaluation bar and a
    blunder alert, and every tensor the graph returns is one the browser has to read back over
    the worker boundary. ``blunder`` comes out as a probability rather than a logit, so the page
    compares it against 0.5 instead of carrying a sigmoid of its own.
    """

    def __init__(self, model: MultiHead) -> None:
        super().__init__()
        self.model = model

    def forward(self, idx: Tensor) -> tuple[Tensor, Tensor]:
        pooled = self.model.pooled(idx)
        return self.model.value(pooled), torch.sigmoid(self.model.blunder(pooled))

def _write_graph(
    wrapper: nn.Module,
    example: Tensor,
    path: Path,
    opset: int,
    outputs: Sequence[str],
    dynamic_batch: bool,
    dynamic_seq: bool,
) -> tuple[Literal["dynamo", "legacy"], str | None]:
    """Write one ONNX file with the modern exporter, falling back to the legacy one.

    Both graphs of the project (the decoder's next move and the encoder's two heads) are written
    here so that the fallback, the console workaround and the axis declarations exist once.
    """
    common: dict[str, Any] = {
        "input_names": [INPUT_NAME],
        "output_names": list(outputs),
        "opset_version": opset,
    }
    try:
        with torch.no_grad(), _utf8_console():
            torch.onnx.export(
                wrapper,
                (example,),
                str(path),
                dynamo=True,
                dynamic_shapes=_dynamic_shapes(dynamic_batch, dynamic_seq),
                **common,
            )
    except Exception as exc:  # noqa: BLE001 - any exporter failure must fall back, not stop
        warning = f"the dynamo exporter failed ({type(exc).__name__}: {exc}); used the legacy one"
        log.warning("%s", warning)
        with torch.no_grad():
            torch.onnx.export(
                wrapper,
                (example,),
                str(path),
                dynamo=False,
                dynamic_axes=_dynamic_axes(dynamic_batch, dynamic_seq, outputs),
                **common,
            )
        return "legacy", warning
    return "dynamo", None

def export_encoder_onnx(
    ckpt: Path | MultiHead,
    out: Path,
    opset: int = DEFAULT_OPSET,
    dynamic_batch: bool = True,
) -> ExportResult:
    """Export the ``value`` and ``blunder`` heads of a fine-tuned encoder to ONNX.

    The length of the input is the scheme's, not the caller's: ``squares`` is always 69 tokens
    and there is no sequence axis to make dynamic, while ``moves`` reads a growing game exactly
    as the decoder does. The example the exporter traces is a sequence of real tokens rather
    than zeros, because ``<pad>`` everywhere is a position the encoder legitimately refuses.
    """
    model = ckpt if isinstance(ckpt, MultiHead) else _load_heads(Path(ckpt))
    model = model.eval()
    cfg = model.encoder.cfg
    seq_len = cfg.seq
    dynamic_seq = cfg.input == "moves"
    wrapper = EncoderHeads(model).eval()
    # Two rows, not one, and no zeros. `torch.export` specialises a dimension whose example is
    # 1 (the 0/1 specialisation), so a batch of one is baked into the graph and the demo's
    # second position fails inside a reshape; tracing at two keeps the axis dynamic and the file
    # still accepts a batch of one. `<pad>` everywhere would be an entirely masked row, which
    # the encoder refuses on purpose, so the example is made of real token ids.
    example = torch.ones((2, seq_len), dtype=torch.long)
    path = target_path(out)
    path.parent.mkdir(parents=True, exist_ok=True)

    exporter, warning = _write_graph(
        wrapper, example, path, opset, list(ENCODER_OUTPUTS), dynamic_batch, dynamic_seq
    )
    works = verify_dynamic_seq(path, seq_len, cfg.seq) if dynamic_seq else None
    if dynamic_seq and works is False:
        raise ValueError(
            f"{path} was exported with a dynamic sequence axis but only runs at length "
            f"{seq_len}: the {exporter} exporter baked the length in."
        )
    really_dynamic = dynamic_seq if works is None else works
    clear_value_info(path)
    metadata = write_metadata(
        path,
        {
            "kind": "encoder",
            "heads": ",".join(ENCODER_OUTPUTS),
            "input": cfg.input,
            "pooling": model.pooling,
            "blunder": "probability",
            "block": seq_len,
            "vocab_size": cfg.tokens,
            "seq_len": seq_len,
            "dynamic_batch": dynamic_batch,
            "dynamic_seq": really_dynamic,
            "exporter": exporter,
            "version": __version__,
        },
    )

    return ExportResult(
        path=path.as_posix(),
        kind="encoder",
        outputs=list(ENCODER_OUTPUTS),
        exporter=exporter,
        opset=opset,
        seq_len=seq_len,
        block=seq_len,
        dynamic_batch=dynamic_batch,
        dynamic_seq=really_dynamic,
        dynamic_seq_verified=works is not None,
        metadata=metadata,
        vocab_size=cfg.tokens,
        params=sum(parameter.numel() for parameter in model.parameters()),
        bytes=path.stat().st_size,
        warning=warning,
    )

def clear_value_info(path: Path) -> int:
    """Drop the graph's annotations for intermediate values; returns how many were removed.

    The modern exporter records a shape for every intermediate tensor, and some of those are the
    shape of the traced example rather than of the dynamic graph. Nothing runs them:
    onnxruntime executes the file regardless. But ``quantize_dynamic`` re-runs shape inference
    in strict mode first and refuses a file whose annotations disagree with what it infers, so
    the int8 build of the encoder dies on an annotation instead of on a weight. They are
    optional by the specification, so the encoder's graph goes out without them.
    """
    try:
        import onnx
    except ImportError:  # pragma: no cover - onnx is a hard dependency of the exporter
        return 0
    model = onnx.load(str(path))
    removed = len(model.graph.value_info)
    del model.graph.value_info[:]
    onnx.save(model, str(path))
    return removed

def _load_heads(ckpt: Path) -> MultiHead:
    from rukh.train import load_heads

    model, _payload = load_heads(ckpt)
    return model

def _align_cast_outputs(model: object) -> int:
    """Make every ``Cast`` say the type its own output is declared to be; count the repairs.

    ``convert_float_to_float16`` retypes the values of the graph but does not touch the ``to``
    attribute of a ``Cast`` node, so a cast that used to produce float32 keeps saying so while
    its declared output is now float16. onnxruntime refuses to load that file — the symptom is a
    ``Type Error: Type (tensor(float16)) of output arg ... does not match expected type
    (tensor(float))`` — and the encoder hits it because pooling casts its padding mask to the
    hidden dtype. Only intermediate values are considered: the graph's own inputs and outputs
    keep the types ``keep_io_types`` promised the caller.
    """
    from onnx import TensorProto

    graph = model.graph  # type: ignore[attr-defined]
    declared = {value.name: value.type.tensor_type.elem_type for value in graph.value_info}
    repaired = 0
    for node in graph.node:
        if node.op_type != "Cast" or not node.output:
            continue
        wanted = declared.get(node.output[0])
        if wanted not in (TensorProto.FLOAT, TensorProto.FLOAT16):
            continue
        for attribute in node.attribute:
            if attribute.name == "to" and attribute.i != wanted:
                attribute.i = int(wanted)
                repaired += 1
    return repaired

    repaired = _align_cast_outputs(converted)
    if repaired:
        log.info(
            "retyped %d Cast node(s) the fp16 conversion left disagreeing with the graph", repaired
        )
    onnx.save(converted, str(target))
    return QuantizeResult(
        path=target.as_posix(),
        kind="fp16",
        method=method,
        bytes=target.stat().st_size,
        source_bytes=source.stat().st_size,
    )

def quantize_int8(path: Path, out: Path | None = None) -> QuantizeResult:
    """Dynamically quantize the MatMul and Gemm weights of an ONNX model to int8."""
    from onnxruntime.quantization import QuantType, quantize_dynamic

    source = Path(path)
    target = _sibling(source, INT8_NAME, out)
    target.parent.mkdir(parents=True, exist_ok=True)
    quantize_dynamic(
        model_input=str(source),
        model_output=str(target),
        weight_type=QuantType.QInt8,
        op_types_to_quantize=QUANTIZED_OPS,
    )
    return QuantizeResult(
        path=target.as_posix(),
        kind="int8",
        method="onnxruntime.quantization.quantize_dynamic",
        bytes=target.stat().st_size,
        source_bytes=source.stat().st_size,
    )

BLUNDER_THRESHOLD = 0.5
"""The exported graph returns a probability; this is where the demo's alert switches on."""

class EncoderParityResult(BaseModel):
    """How well an exported encoder reproduces the checkpoint's two heads."""

    model_config = ConfigDict(extra="forbid")

    positions: int
    agreement: float
    """Share of positions where the ``blunder`` decision (at 0.5) is the same."""
    max_abs_value_delta: float
    max_abs_blunder_delta: float
    mismatches: list[int]
    """Indices of the first few positions where the decision differed."""

def label_positions(
    model: MultiHead,
    cfg: LabelsConfig | None = None,
    n: int = DEFAULT_N,
    split: str = "val",
) -> list[list[int]]:
    """Token sequences of ``n`` held-out labelled positions, in the encoder's own scheme.

    Empty when the labelled table has not been built: an encoder parity check has no honest
    fallback (a random board is not a position the demo will ever evaluate), so the caller
    reports that it could not be measured instead of measuring something else.
    """
    import polars as pl

    from rukh.data.labels import LabelsConfig as Labels
    from rukh.data.labels import build_labels, game_moves
    from rukh.train.heads import LabelledPositions

    labels = cfg if cfg is not None else Labels()
    try:
        frame = build_labels(labels)
    except FileNotFoundError as exc:
        log.warning("no labelled positions for the parity check: %s", exc)
        return []
    rows = frame.filter(pl.col("split") == split).sort(["order", "game_id", "ply"]).head(n)
    if not rows.height:
        return []
    # The scheme is the model's, not a constant: a ``moves`` encoder tokenized with
    # ``fen_to_tokens`` would be compared against its export on ids from another vocabulary, and
    # the parity number would be meaningless rather than wrong-looking.
    scheme = str(model.encoder.cfg.input)
    moves = game_moves(rows, labels.games_dir) if scheme == "moves" else None
    dataset = LabelledPositions(rows, scheme, moves, model.encoder.cfg.block)
    return [dataset.tokens(index) for index in range(len(dataset))]

def encoder_parity_positions(
    model: MultiHead, cfg: LabelsConfig | None = None, n: int = DEFAULT_N, split: str = "val"
) -> tuple[list[list[int]], str, str | None]:
    """``(positions, source, warning)`` for the encoder: validation labels, or nothing.

    ``model`` is here to say which scheme the positions have to be tokenized in; it is never
    run."""
    positions = label_positions(model, cfg, n=n, split=split)
    if positions:
        return positions, "validation-labels", None
    warning = (
        "no labelled validation positions (build them with `rukh data evals`): the encoder "
        "parity check was skipped rather than measured on positions nobody will evaluate"
    )
    log.warning("%s", warning)
    return [], "none", warning

def encoder_parity(
    model: MultiHead,
    onnx_path: Path,
    positions: Sequence[Sequence[int]],
    n: int = DEFAULT_N,
    threshold: float = BLUNDER_THRESHOLD,
) -> EncoderParityResult:
    """Compare the blunder decision and the value of the checkpoint and the exported file."""
    from rukh.export.onnx import BLUNDER_OUTPUT, VALUE_OUTPUT, EncoderHeads

    used = list(positions)[:n]
    if not used:
        raise ValueError("parity needs at least one position")
    wrapper = EncoderHeads(model.eval()).eval()
    session = _session(Path(onnx_path))

    agreed = 0
    worst_value = worst_blunder = 0.0
    mismatches: list[int] = []
    for index, tokens in enumerate(used):
        idx = np.asarray([list(tokens)], dtype=np.int64)
        with torch.no_grad():
            value, blunder = wrapper(torch.from_numpy(idx))
        reference = {
            VALUE_OUTPUT: float(value[0]),
            BLUNDER_OUTPUT: float(blunder[0]),
        }
        outputs = session.run(None, {INPUT_NAME: idx})
        # Named by the graph itself, not by position: zipping against ``ENCODER_OUTPUTS`` would
        # keep comparing happily if the exporter ever swapped the two heads, and the parity
        # check would then be measuring `value` against `blunder` and calling it agreement.
        exported = {
            output.name: float(np.asarray(array).reshape(-1)[0])
            for output, array in zip(session.get_outputs(), outputs, strict=True)
        }
        missing = [name for name in (VALUE_OUTPUT, BLUNDER_OUTPUT) if name not in exported]
        if missing:
            raise ValueError(f"the exported graph has no {missing} output: {sorted(exported)}")
        worst_value = max(worst_value, abs(reference[VALUE_OUTPUT] - exported[VALUE_OUTPUT]))
        worst_blunder = max(
            worst_blunder, abs(reference[BLUNDER_OUTPUT] - exported[BLUNDER_OUTPUT])
        )
        same = (reference[BLUNDER_OUTPUT] >= threshold) == (exported[BLUNDER_OUTPUT] >= threshold)
        if same:
            agreed += 1
        elif len(mismatches) < MAX_MISMATCHES:
            mismatches.append(index)
    return EncoderParityResult(
        positions=len(used),
        agreement=agreed / len(used),
        max_abs_value_delta=worst_value,
        max_abs_blunder_delta=worst_blunder,
        mismatches=mismatches,
    )

KINDS = ("decoder", "encoder")
"""What ``rukh export --kind`` accepts."""


class ExportBundle(BaseModel):
    """Everything one ``rukh export`` produced."""

    model_config = ConfigDict(extra="forbid")

    onnx: ExportResult
    fp16: QuantizeResult | None = None
    int8: QuantizeResult | None = None
    parity: dict[str, ParityResult] = {}
    """Parity per file kind: ``fp32``, ``fp16``, ``int8``."""
    heads_parity: dict[str, EncoderParityResult] = {}
    """The encoder's parity per file kind: the blunder decision and the drift of ``value``."""
    parity_source: str | None = None
    """``validation`` (the positions of ``docs/spec/02`` §6), ``random-walk``, or, for the
    encoder, ``validation-labels``."""
    parity_warning: str | None = None

def _quantize(bundle: ExportBundle, fp16: bool, int8: bool) -> dict[str, str]:
    """Derive the browser's files from the fp32 graph and return every file to check."""
    if fp16:
        bundle.fp16 = to_fp16(Path(bundle.onnx.path))
    if int8:
        bundle.int8 = quantize_int8(Path(bundle.onnx.path))
    # The demo loads the fp16 or the int8 file, not the fp32 one, so the context length has to
    # travel with them too; neither converter promises to keep the metadata of its input.
    for quantized in (bundle.fp16, bundle.int8):
        if quantized is not None and bundle.onnx.metadata:
            set_metadata(Path(quantized.path), bundle.onnx.metadata)
    checks = {"fp32": bundle.onnx.path}
    if bundle.fp16 is not None:
        checks["fp16"] = bundle.fp16.path
    if bundle.int8 is not None:
        checks["int8"] = bundle.int8.path
    return checks

def export_all(
    ckpt: Path,
    out: Path,
    kind: str = "decoder",
    opset: int = DEFAULT_OPSET,
    seq_len: int = 200,
    fp16: bool = False,
    int8: bool = False,
    check_parity: bool = False,
    positions: int = 1_000,
    seed: int = 0,
    games: Path | None = None,
    labels: LabelsConfig | None = None,
    split: str = "val",
) -> ExportBundle:
    """Export, quantize and check one checkpoint in a single pass.

    For the decoder, parity is measured on validation positions (``games``, or the validation
    month of P1 when it is on disk) and falls back to random legal walks with a warning when
    there are none. For the encoder it is measured on the held-out labelled positions of
    ``rukh.data.labels`` and is skipped, with a warning, when those have not been built: a
    random board is not a position the demo will ever be asked to evaluate.
    """
    if kind not in KINDS:
        raise ValueError(f"unknown kind {kind!r}; expected one of {', '.join(KINDS)}")
    if kind == "encoder":
        return _export_encoder(ckpt, out, opset, fp16, int8, check_parity, positions, labels, split)

def _export_encoder(
    ckpt: Path,
    out: Path,
    opset: int,
    fp16: bool,
    int8: bool,
    check_parity: bool,
    positions: int,
    labels: LabelsConfig | None,
    split: str,
) -> ExportBundle:
    """``export_all(kind="encoder")``: the two heads, the same files, its own parity."""
    from rukh.train import load_heads

    model, _payload = load_heads(Path(ckpt))
    bundle = ExportBundle(onnx=export_encoder_onnx(model, out, opset=opset))
    checks = _quantize(bundle, fp16, int8)
    if check_parity:
        items, source, warning = encoder_parity_positions(model, labels, n=positions, split=split)
        bundle.parity_source = source
        bundle.parity_warning = warning
        if items:
            bundle.heads_parity = {
                name: encoder_parity(model, Path(path), items, n=positions)
                for name, path in checks.items()
            }
    return bundle

__all__ = [
    "BLUNDER_OUTPUT",
    "DEFAULT_GAMES",
    "DEFAULT_OPSET",
    "ENCODER_OUTPUTS",
    "FP16_NAME",
    "INPUT_NAME",
    "INT8_NAME",
    "KINDS",
    "MODEL_NAME",
    "OUTPUT_NAME",
    "QUANTIZED_OPS",
    "VALUE_OUTPUT",
    "EmbedResult",
    "EncoderHeads",
    "EncoderParityResult",
    "ExportBundle",
    "ExportResult",
    "LastStepLogits",
    "ParityResult",
    "QuantizeResult",
    "clear_value_info",
    "embed_positions",
    "encoder_parity",
    "encoder_parity_positions",
    "export_all",
    "export_encoder_onnx",
    "export_onnx",
    "fen4",
    "label_positions",
    "parity",
    "parity_positions",
    "quantize_int8",
    "random_prefixes",
    "read_metadata",
    "sequence_lengths",
    "set_metadata",
    "target_path",
    "to_fp16",
    "validation_prefixes",
    "verify_dynamic_seq",
    "write_metadata",
]

"""Mean-pooled position embeddings: the encoder's secondary output, for phase 2.

``docs/spec/02`` §"Componente 2" asks for one vector per position so that similar positions can
be retrieved later (the A1 agent of phase 2). The vector is ``PositionEncoder.pool(idx, "mean")``
over the ``squares`` tokens of a FEN: the average of the real tokens, which is why padding never
enters it.

A matrix of floats is useless without knowing which row is which position, and a ``.npy`` file
has no room to say so. Every run therefore writes two files: the array, and a parquet sidecar
with ``row`` and ``fen4`` in the array's own order. The four-field FEN is the key P1 uses for
positions and evaluations, so the sidecar joins straight back to them.
"""

from __future__ import annotations

import logging
from pathlib import Path

import polars as pl
from pydantic import BaseModel, ConfigDict

log = logging.getLogger(__name__)

FEN_COLUMN = "fen"
SIDECAR_SUFFIX = ".parquet"
POOLING = "mean"


class EmbedResult(BaseModel):
    """What one embedding run wrote."""

    model_config = ConfigDict(extra="forbid")

    positions: int
    dim: int
    array: str
    sidecar: str
    pooling: str = POOLING
    input: str = "squares"
    checkpoint: str

def fen4(fen: str) -> str:
    """The four-field FEN: placement, side to move, castling and en passant."""
    return " ".join(str(fen).split(" ")[:4])


def _encoder(ckpt: Path, device: str | None = None):  # type: ignore[no-untyped-def]
    """The ``PositionEncoder`` inside a checkpoint, whether or not it carries the heads."""
    from rukh.models.heads import MultiHead
    from rukh.train import load_any

    model, _payload, kind = load_any(Path(ckpt), map_location=device or "cpu")
    if kind != "encoder":
        raise ValueError(f"{ckpt} is a {kind} checkpoint; embeddings need an encoder")
    return model.encoder if isinstance(model, MultiHead) else model

def embed_positions(
    ckpt: Path,
    positions: Path,
    out: Path,
    batch_size: int = 256,
    device: str | None = None,
    column: str = FEN_COLUMN,
) -> EmbedResult:
    """Embed every position of a parquet and write the array plus its ``fen4`` sidecar."""
    import numpy as np
    import torch

    from rukh.models.squares import fen_to_tokens
    from rukh.train import pick_device

    source = Path(positions)
    frame = pl.read_parquet(source)
    if column not in frame.columns:
        raise ValueError(f"{source} has no '{column}' column (found {', '.join(frame.columns)})")
    fens = [str(value) for value in frame[column].to_list()]
    if not fens:
        raise ValueError(f"{source} holds no positions")

    where = torch.device(device or pick_device())
    encoder = _encoder(Path(ckpt), str(where)).to(where).eval()
    if encoder.cfg.input != "squares":
        raise ValueError(
            f"the checkpoint reads the {encoder.cfg.input!r} scheme; embeddings of a FEN table "
            "need the 'squares' scheme, which is the one the labelled positions carry"
        )

    vectors: list[np.ndarray] = []
    with torch.no_grad():
        for start in range(0, len(fens), max(1, batch_size)):
            chunk = fens[start : start + max(1, batch_size)]
            idx = torch.tensor([fen_to_tokens(fen) for fen in chunk], dtype=torch.long)
            pooled = encoder.pool(idx.to(where), POOLING)
            vectors.append(pooled.detach().float().cpu().numpy())
    array = np.concatenate(vectors) if vectors else np.zeros((0, encoder.cfg.d_model), "float32")

    target = Path(out)
    if target.suffix != ".npy":
        target = target / "embeddings.npy"
    target.parent.mkdir(parents=True, exist_ok=True)
    np.save(target, array.astype("float32"))
    sidecar = target.with_suffix(SIDECAR_SUFFIX)
    pl.DataFrame(
        {
            "row": pl.Series("row", range(len(fens)), dtype=pl.UInt32),
            "fen4": [fen4(fen) for fen in fens],
        }
    ).write_parquet(sidecar)
    log.info("wrote %d embeddings of %d dimensions to %s", array.shape[0], array.shape[1], target)
    return EmbedResult(
        positions=int(array.shape[0]),
        dim=int(array.shape[1]),
        array=target.as_posix(),
        sidecar=sidecar.as_posix(),
        input=encoder.cfg.input,
        checkpoint=Path(ckpt).as_posix(),
    )

    kind: Annotated[
        str, typer.Option("--kind", help="What to export: decoder or encoder.")
    ] = "decoder",

    """Export a model to ONNX, quantize it and check parity with PyTorch.

    The decoder exports its next-move head; ``--kind encoder`` exports the two heads the demo
    reads from a position, ``value`` and ``blunder``.
    """
    from rukh.export import KINDS, export_all

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    if kind not in KINDS:
        typer.echo(f"error: --kind must be one of {', '.join(KINDS)}", err=True)
        raise typer.Exit(code=2)

    try:
        bundle = export_all(
            ckpt,
            out,
            kind=kind,
            opset=opset,
            seq_len=seq_len,
            fp16=fp16,
            int8=int8,
            check_parity=check_parity,
            positions=positions,
            games=games,
        )
    except (ImportError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    if as_json:
        typer.echo(bundle.model_dump_json(indent=2))
        return

    typer.echo(f"exporter: {bundle.onnx.exporter} (opset {bundle.onnx.opset})")
    typer.echo(f"kind:     {bundle.onnx.kind} -> {', '.join(bundle.onnx.outputs)}")
    if bundle.onnx.warning:
        typer.echo(f"warning:  {bundle.onnx.warning}")
    checked = "verified" if bundle.onnx.dynamic_seq_verified else "not verified"
    typer.echo(
        f"shapes:   batch dynamic={bundle.onnx.dynamic_batch}, "
        f"sequence dynamic={bundle.onnx.dynamic_seq} ({checked}), block={bundle.onnx.block}"
    )
    typer.echo(f"fp32:     {bundle.onnx.path} ({bundle.onnx.bytes} bytes)")
    for quantized in (bundle.fp16, bundle.int8):
        if quantized is not None:
            typer.echo(
                f"{quantized.kind}:     {quantized.path} ({quantized.bytes} bytes, "
                f"{quantized.ratio:.2f} of fp32, {quantized.method})"
            )
    for name, result in bundle.parity.items():
        typer.echo(
            f"parity {name}: {result.agreement:.4f} on {result.positions} "
            f"{bundle.parity_source} positions "
            f"(max |delta logits| {result.max_abs_logit_delta:.4g})"
        )
    for name, heads in bundle.heads_parity.items():
        typer.echo(
            f"parity {name}: {heads.agreement:.4f} of the blunder decisions on "
            f"{heads.positions} {bundle.parity_source} positions "
            f"(max |delta value| {heads.max_abs_value_delta:.4g})"
        )
    if bundle.parity_warning:
        typer.echo(f"warning:  {bundle.parity_warning}")

encoder_app = typer.Typer(help="Encoder utilities: position embeddings.", no_args_is_help=True)
app.add_typer(encoder_app, name="encoder")

@encoder_app.command("embed")
def encoder_embed_cmd(
    positions: Annotated[
        Path,
        typer.Option(
            "--positions",
            exists=True,
            dir_okay=False,
            readable=True,
            help="Parquet with a `fen` column (the P1 positions or evaluations table).",
        ),
    ],
    out: Annotated[Path, typer.Option("--out", help="Where to write the .npy array.")],
    ckpt: Annotated[
        Path,
        typer.Option(
            "--ckpt", exists=True, dir_okay=False, readable=True, help="Encoder checkpoint."
        ),
    ],
    batch_size: Annotated[
        int, typer.Option("--batch-size", help="Positions per forward pass.")
    ] = 256,
    device: Annotated[
        str | None, typer.Option("--device", help="Where to run: cuda, cpu... ")
    ] = None,
    as_json: Annotated[bool, typer.Option("--json", help="Print the result as JSON only.")] = False,
) -> None:

    """Write mean-pooled position embeddings plus the `fen4` sidecar that names their rows."""
    from rukh.export import embed_positions

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    try:
        result = embed_positions(ckpt, positions, out, batch_size=batch_size, device=device)
    except (FileNotFoundError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    if as_json:
        typer.echo(result.model_dump_json(indent=2))
        return
    typer.echo(f"positions: {result.positions}")
    typer.echo(f"dim:       {result.dim} ({result.pooling} pooling, {result.input} scheme)")
    typer.echo(f"array:     {result.array}")
    typer.echo(f"sidecar:   {result.sidecar} (row, fen4)")

Two kinds of model go out through this one path. A decoder is published with its UCI
vocabulary and the metrics of ``rukh eval``; an encoder (the ``MultiHead`` of M3: the position
encoder plus the value, blunder and result heads) is published with the ``squares`` vocabulary
it reads and the metrics of ``rukh eval encoder``, including the F1 of the material baseline it
has to beat. ``checkpoint_kind`` decides which one a checkpoint is, from the weights themselves,
and the two differ only in the card, the config and the vocabulary file.

CARD_TEMPLATE = "model.md.jinja"
ENCODER_CARD_TEMPLATE = "encoder.md.jinja"
CONFIG_NAME = "config.json"
README_NAME = "README.md"
SAFETENSORS_NAME = "model.safetensors"
TORCH_NAME = "pytorch_model.bin"
VOCAB_PATH = "tokenizer/vocab.json"
ONNX_DIR = "onnx"
ONNX_FILES = ("model-fp16.onnx", "model-int8.onnx", "model.onnx")
REPO_TYPE = "model"
ARCHITECTURES = {"decoder": "MoveDecoder", "encoder": "PositionEncoder"}
MODEL_TYPES = {"decoder": "rukh-move-decoder", "encoder": "rukh-position-encoder"}
ENCODER_HEADS = ("value", "blunder", "result")

def tied_names(model: Any) -> list[str]:
    """The tied heads of a model: the names whose storage is the embedding's, under another name.

    Asked of the tensors rather than of the config, because the encoder's tie lives one level
    down (``encoder.mlm_head.weight``) and only exists for the ``moves`` scheme.
    """
    state = model.state_dict()
    return [
        name
        for name, source in TIED_SOURCES.items()
        if name in state and source in state and state[name].data_ptr() == state[source].data_ptr()
    ]


def publish_state(model: Any) -> dict[str, Any]:
    """The state dict as it is published: on the CPU, contiguous and with no aliased tensor.

    With tied embeddings ``lm_head.weight`` *is* ``tokens.weight``; it is dropped here so the
    file holds every tensor exactly once. Loading re-ties it (see the module docstring).
    """
    dropped = set(tied_names(model))
    return {
        key: value.detach().cpu().contiguous()
        for key, value in model.state_dict().items()
        if key not in dropped
    }

def write_config(
    payload: dict[str, Any],
    stage: str,
    params: int,
    folder: Path,
    kind: str = "decoder",
) -> dict[str, Any]:
    """Write ``config.json``: the model's shape plus the provenance of the checkpoint."""
    model_cfg = dict(payload.get("model_cfg") or {})
    scheme = str(model_cfg.get("input") or "moves")
    encoder = kind == "encoder"
    config = {
        "architectures": [ARCHITECTURES.get(kind, ARCHITECTURES["decoder"])],
        "model_type": MODEL_TYPES.get(kind, MODEL_TYPES["decoder"]),
        "library_name": "rukh",
        "rukh_version": __version__,
        "stage": stage,
        "step": payload.get("step"),
        "params": params,
        "tokenizer": scheme if encoder else "uci",
        "vocab_hash": payload.get("vocab_hash"),
        "data_manifest_sha": payload.get("data_manifest_sha"),
        "git_sha": payload.get("git_sha"),
        **(
            {
                "heads": list(ENCODER_HEADS),
                "pooling": _pooling(payload),
                "pretrained_from": _pretrained_from(payload),
            }
            if encoder
            else {}
        ),
        **model_cfg,
    }

def _pooling(payload: dict[str, Any]) -> str:
    """How the published encoder pools its tokens, as its fine-tuning run recorded it."""
    return str((payload.get("cfg") or {}).get("pooling") or "mean")


def _pretrained_from(payload: dict[str, Any]) -> str | None:
    """The masked-move checkpoint the heads were fine-tuned from, or ``None`` for from scratch.

    ``encoder_ckpt`` is ``None`` in a perfectly legitimate run — training the encoder from
    scratch is the baseline the pretraining has to beat — so the card cannot claim masked move
    modeling unconditionally; this is the fact it gates that claim on.
    """
    found = (payload.get("cfg") or {}).get("encoder_ckpt")
    return str(found) if found else None

def write_vocab(folder: Path, kind: str, scheme: str) -> str:
    """Write the vocabulary the model actually reads, and return its path in the repository.

    A decoder (and an encoder on the ``moves`` scheme) reads the P1 UCI enumeration; an encoder
    on the ``squares`` scheme reads the 47 fixed square tokens, which are just as much part of
    the released model: without them the 69 integers of an input mean nothing.
    """
    if kind == "encoder" and scheme == "squares":
        from rukh.models.squares import SQUARE_TOKENS, SQUARE_VOCAB, vocab_hash

        path = folder / VOCAB_PATH
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = {
            "version": 1,
            "scheme": "squares",
            "size": len(SQUARE_VOCAB),
            "sequence": SQUARE_TOKENS,
            "vocab_hash": vocab_hash(),
            "tokens": list(SQUARE_VOCAB),
        }
        path.write_text(
            json.dumps(payload, indent=0, ensure_ascii=True) + "\n", encoding="utf-8", newline="\n"
        )
        return VOCAB_PATH
    UciTokenizer().export(folder / VOCAB_PATH)
    return VOCAB_PATH

def encoder_card_context(
    repo_id: str,
    stage: str,
    cfg: ModelPublishConfig,
    config: dict[str, Any],
    evaluation: dict[str, Any] | None,
    run: RunSummary | None,
    files: list[str],
) -> dict[str, Any]:
    """Everything the encoder card needs: its metrics, the baseline's, and the input scheme."""
    from rukh.eval.encoder import base_rate_caveat, sentence

    measured = evaluation or {}
    encoder = measured.get("encoder_blunder") or {}
    fixed = measured.get("encoder_blunder_fixed") or {}
    ranked = measured.get("encoder_blunder_ranking") or {}
    baseline = measured.get("heuristic_blunder") or {}
    value = measured.get("encoder_value") or {}
    margin = measured.get("f1_margin")
    tuned = measured.get("threshold_tuned")
    configured = measured.get("threshold_fixed")
    metrics = [
        (
            "Blunder F1" + ("" if tuned is None else f" (tuned, `p >= {float(tuned):.4g}`)"),
            _percent(encoder.get("f1")),
        ),
        ("Blunder precision", _percent(encoder.get("precision"))),
        ("Blunder recall", _percent(encoder.get("recall"))),
        (
            "Blunder F1"
            + ("" if configured is None else f" (fixed, `p >= {float(configured):.4g}`)"),
            _percent(fixed.get("f1")),
        ),
        ("Blunder ROC AUC", _ratio(ranked.get("roc_auc"))),
        ("Blunder average precision", _ratio(ranked.get("average_precision"))),
        ("Blunder base rate", _percent(measured.get("blunder_base_rate"))),
        ("Blunder F1, material baseline", _percent(baseline.get("f1"))),
        (
            "Margin over the baseline",
            "n/a" if margin is None else f"{float(margin):+.1f} F1 points",
        ),
        ("Value vs Stockfish cp, Pearson", _ratio(value.get("pearson"))),
        ("Value vs Stockfish cp, Spearman", _ratio(value.get("spearman"))),
        ("Result accuracy", _percent(measured.get("result_accuracy"))),
    ]

    curve = [
        (f"{float(point.get('fraction', 0)) * 100:.0f} %", point.get("train_labels"))
        for point in measured.get("label_curve", [])
        if isinstance(point, dict)
    ]

    return {
        "repo_id": repo_id,
        "stage": stage,
        "license": cfg.license,
        "datasets": cfg.datasets,
        "demo_url": f"{cfg.demo_url}/?stage={stage}",
        "course_url": cfg.course_url,
        "repository_url": cfg.repository_url,
        "params": config.get("params", 0),
        "config": config,
        "config_json": json.dumps(config, indent=2, ensure_ascii=False),
        "scheme": config.get("tokenizer", "squares"),
        "pooling": config.get("pooling", "mean"),
        "pretrained_from": config.get("pretrained_from"),
        "heads": list(config.get("heads", ENCODER_HEADS)),
        "metrics": metrics,
        "curve": curve,
        "positions": measured.get("items"),
        "blunder_items": measured.get("blunder_items"),
        "blunder_caveat": sentence(base_rate_caveat(measured.get("blunder_base_rate"))),
        "tune_items": measured.get("tune_items"),
        "score_items": measured.get("score_items"),
        "tune_games": measured.get("tune_games"),
        "score_games": measured.get("score_games"),
        "threshold_tuned": tuned,
        "evaluated_on": measured.get("date"),
        "notes": [str(note) for note in measured.get("notes", []) if str(note).strip()],
        "run_id": run.run_id if run else None,
        "recipe": sorted((run.params if run else {}).items()),
        "files": files,
        "has_onnx": any(name.startswith(f"{ONNX_DIR}/") for name in files),
        "rukh_version": __version__,
    }

        "blunder_caveat": sentence(base_rate_caveat(measured.get("blunder_base_rate"))),
        "tune_items": measured.get("tune_items"),
        "score_items": measured.get("score_items"),
        "tune_games": measured.get("tune_games"),
        "score_games": measured.get("score_games"),
        "threshold_tuned": tuned,
        "evaluated_on": measured.get("date"),
        "notes": [str(note) for note in measured.get("notes", []) if str(note).strip()],

def render_card(context: dict[str, Any], template: str = CARD_TEMPLATE) -> str:
    """Render the English model card."""
    env = Environment(
        loader=FileSystemLoader(str(CARDS_DIR)),
        undefined=StrictUndefined,
        autoescape=False,
        keep_trailing_newline=True,
    )
    return env.get_template(template).render(**context)

def publish_model(
    ckpt: Path,
    repo: str,
    cfg: ModelPublishConfig | None = None,
    onnx_dir: Path | None = None,
    stage: str | None = None,
    run_id: str | None = None,
    dry_run: bool = False,
) -> ModelPublishResult:
    """Stage (and unless ``dry_run``, upload) one trained model as a Hub model repository."""
    from rukh.train import load_any

    cfg = cfg or ModelPublishConfig()
    ckpt = Path(ckpt)
    repo_id = repo if "/" in repo else f"{cfg.owner}/{repo}"
    name = stage or repo_id.split("/")[-1].removeprefix("rukh-")

    model, payload, kind = load_any(ckpt)
    state = publish_state(model)
    params = (
        model.num_params(non_embedding=False)
        if hasattr(model, "num_params")
        else sum(parameter.numel() for parameter in model.parameters())
    )

    folder = resolve(cfg.publish_dir) / repo_id
    folder.mkdir(parents=True, exist_ok=True)
    weights_name, weights_format = write_weights(state, folder)
    config = write_config(payload, name, params, folder, kind)
    write_vocab(folder, kind, str(config.get("input") or "moves"))
    files = [weights_name, CONFIG_NAME, VOCAB_PATH, *copy_onnx(onnx_dir, folder)]

    run = read_run(run_id, run_name=ckpt.parent.name)
    evaluation = read_eval(name, cfg)
    card = (
        render_card(
            encoder_card_context(repo_id, name, cfg, config, evaluation, run, files),
            ENCODER_CARD_TEMPLATE,
        )
        if kind == "encoder"
        else render_card(card_context(repo_id, name, cfg, config, evaluation, run, files))
    )

    if not dry_run:
        api = _api()
        api.create_repo(repo_id, repo_type=REPO_TYPE, exist_ok=True)
        api.upload_folder(
            repo_id=repo_id,
            folder_path=str(folder),
            repo_type=REPO_TYPE,
            commit_message=f"Publish {name}",
        )
    return ModelPublishResult(
        repo_id=repo_id,
        kind=kind,
        stage=name,
        dry_run=dry_run,
        folder=folder.as_posix(),
        card_path=card_path.as_posix(),
        files=[*files, README_NAME],
        weights_format=weights_format,
        run_id=run.run_id if run else None,
        params=params,
        tied_embeddings=bool(tied_names(model)),
    )