Adaptadores y publicación: 1,6 MB que llegan al navegador

"""Exporting the decoder with its LoRA factors as **inputs of the graph**.

Merging an adapter and exporting the result is the obvious way to serve a style in the browser,
and it throws away the only thing LoRA bought: one style is 221 MB of ONNX for ``medium``, two
styles are 442 MB, and the 1.6 MB file stops meaning anything. So the graph written here keeps
the correction outside the weights. ``A`` and ``B`` arrive as two tensors on every call, stacked
over the layers (``rukh.models.lora.AdapterLayout``), and the browser changes style by uploading
1.6 MB instead of downloading a second model.

Three properties make the file honest, and all three are checked by the tests:

* **Zeros are the base model.** A graph that takes the factors as inputs and is fed zeros
  computes exactly what the plain export computes, to the last bit of float32. So the adaptable
  file *replaces* the ordinary one instead of being a second copy of it, and a demo with no
  style selected is not running a different model.
* **The factors are really inputs.** Feeding a different adapter changes the logits without
  touching the file. That is the whole claim, and it is one assertion away from being tested.
* **Parity with PyTorch.** The same adapter through ``apply_lora`` and through this graph has to
  give the same next move; ``rukh.export.parity`` measures it over real positions.

What it does not do is take an adapter of any shape. The factors travel as one tensor each, so
every block must adapt the same matrix with the same slice widths -- true of query, key and
value on the fused ``qkv``, false as soon as the MLP joins in, whose ``fc`` is four times wider.
That case still has ``merge_lora`` and an ordinary export; it just cannot be swapped at runtime.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING, Any

import torch
from torch import Tensor, nn
from torch.nn import functional as F

from rukh import __version__
from rukh.export.onnx import (
    DEFAULT_OPSET,
    INPUT_NAME,
    OUTPUT_NAME,
    ExportResult,
    _load,
    _write_graph,
    target_path,
    verify_dynamic_seq,
    write_metadata,
)
from rukh.models import MoveDecoder
from rukh.models.lora import (
    ADAPTER_CONFIG,
    AdapterLayout,
    LoraConfig,
    LoRALinear,
    adapter_layout,
    apply_lora,
    lora_modules,
    stack_adapter,
)

if TYPE_CHECKING:
    from collections.abc import Sequence

__all__ = [
    "ADAPTER_INPUTS",
    "A_INPUT",
    "B_INPUT",
    "WEB_ADAPTER_FILE",
    "WEB_ADAPTER_META",
    "AdaptableDecoder",
    "adaptable_inputs",
    "browser_lora",
    "build_adaptable",
    "export_adaptable_onnx",
    "write_web_adapter",
    "write_web_adapter_from_file",
    "zero_inputs",
]

A_INPUT = "lora_a"
B_INPUT = "lora_b"
ADAPTER_INPUTS = (A_INPUT, B_INPUT)

WEB_ADAPTER_FILE = "web/adapter.bin"
"""Where a published adapter keeps the copy the browser downloads: raw little-endian float32,
``A`` first and then ``B``, in the order the layers run."""

WEB_ADAPTER_META = "web/adapter.json"
"""Its shapes and provenance, because a flat buffer of floats says nothing about itself. The
sidecar is always the buffer's own name with a ``.json`` suffix, so a file written under another
name keeps its metadata next to it and not next to somebody else's."""


def browser_lora(cfg: LoraConfig) -> LoraConfig:
    """The same adapter shape with ``alpha = r``, which is the scaling the browser graph uses.

    The correction is ``(alpha / r) B A``, and that factor has to live in exactly one place. It
    could be the graph -- and then every adapter fed to it would have to have been trained with
    the same ``alpha`` and ``r``, with nothing in a flat buffer of floats to check it against --
    or it could be the file. It is the file: ``write_web_adapter`` multiplies ``B`` by the
    adapter's own scale before writing it, and the graph the browser runs multiplies by one. An
    adapter of any rank and any alpha is then correct in that graph without anyone matching two
    numbers by hand.
    """
    return cfg.model_copy(update={"alpha": cfg.r})

def zero_inputs(layout: AdapterLayout) -> dict[str, Any]:
    """The feed that makes the graph the base model: an adapter of all zeros."""
    import numpy as np

    return {
        A_INPUT: np.zeros(layout.shape_a, dtype=np.float32),
        B_INPUT: np.zeros(layout.shape_b, dtype=np.float32),
    }


def adaptable_inputs(model: nn.Module) -> dict[str, Any]:
    """The feed that applies the adapter currently loaded into ``model``, as numpy arrays."""
    a, b = stack_adapter(model)
    return {A_INPUT: a.numpy(), B_INPUT: b.numpy()}

def export_adaptable_onnx(
    ckpt: Path | MoveDecoder,
    out: Path,
    lora: LoraConfig | None = None,
    opset: int = DEFAULT_OPSET,
    dynamic_batch: bool = True,
    seq_len: int = 200,
    dynamic_seq: bool = True,
) -> ExportResult:
    """Export the next-move head with the LoRA factors as two extra inputs of the graph.

    The result is a drop-in replacement for ``export_onnx``'s file for any caller willing to
    feed the two tensors: with zeros it is the same model, bit for bit. ``lora`` describes the
    shape of the adapters the file will accept and defaults to the project's own recipe
    (``r = 8`` on the query and value ranges of the fused ``qkv``).
    """
    model = ckpt if isinstance(ckpt, MoveDecoder) else _load(Path(ckpt))
    model = model.eval()
    if seq_len > model.cfg.block:
        raise ValueError(f"seq_len={seq_len} is longer than the model's block {model.cfg.block}")
    wrapper, layout = build_adaptable(model, lora or LoraConfig())
    example = (
        torch.zeros((1, seq_len), dtype=torch.long),
        torch.zeros(layout.shape_a),
        torch.zeros(layout.shape_b),
    )
    path = target_path(out)
    path.parent.mkdir(parents=True, exist_ok=True)

    inputs = (INPUT_NAME, *ADAPTER_INPUTS)
    exporter, warning = _write_graph(
        wrapper, example, path, opset, [OUTPUT_NAME], dynamic_batch, dynamic_seq, inputs
    )
    works = verify_dynamic_seq(path, seq_len, model.cfg.block, zero_inputs(layout))
    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
    metadata = write_metadata(
        path,
        {
            "block": model.cfg.block,
            "vocab_size": model.cfg.vocab_size,
            "seq_len": seq_len,
            "dynamic_batch": dynamic_batch,
            "dynamic_seq": really_dynamic,
            "exporter": exporter,
            "version": __version__,
            "adapter_inputs": ",".join(ADAPTER_INPUTS),
            "adapter_shape_a": ",".join(str(n) for n in layout.shape_a),
            "adapter_shape_b": ",".join(str(n) for n in layout.shape_b),
            "adapter_module": layout.module,
            "adapter_r": layout.r,
            "adapter_scale": layout.scale,
        },
    )
    return ExportResult(
        path=path.as_posix(),
        exporter=exporter,
        opset=opset,
        seq_len=seq_len,
        block=model.cfg.block,
        dynamic_batch=dynamic_batch,
        dynamic_seq=really_dynamic,
        dynamic_seq_verified=works is not None,
        metadata=metadata,
        vocab_size=model.cfg.vocab_size,
        params=model.num_params(non_embedding=False),
        bytes=path.stat().st_size,
        warning=warning,
    )

def _write_web(
    a: Tensor,
    b: Tensor,
    scale: float,
    module: str,
    r: int,
    out: Path | str,
    base_repo: str | None,
) -> Path:
    """The serialisation itself, shared by the two sources an adapter can come from."""
    target = Path(out)
    target.parent.mkdir(parents=True, exist_ok=True)
    with target.open("wb") as handle:
        handle.write(a.cpu().float().numpy().astype("<f4").tobytes())
        handle.write((b.cpu().float() * scale).numpy().astype("<f4").tobytes())
    meta = target.with_suffix(".json")
    meta.write_text(
        json.dumps(
            {
                "format": "rukh-lora-web-1",
                "dtype": "float32",
                "order": [A_INPUT, B_INPUT],
                "shape_a": list(a.shape),
                "shape_b": list(b.shape),
                "module": module,
                "r": r,
                "scale_applied": scale,
                "base_model": base_repo,
                "bytes": target.stat().st_size,
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
        newline="\n",
    )
    return target


def write_web_adapter(model: nn.Module, out: Path | str, base_repo: str | None = None) -> Path:
    """Write the browser's copy of an adapter: one flat float32 buffer and its shapes.

    ``model`` is a decoder with the adapter already applied and loaded (``load_adapter``). What
    comes out is deliberately dull -- no safetensors header to parse in the worker, no shape
    inside the bytes -- because the page fetches it before it can do anything else and every
    dependency it needs is one more thing that can fail between the Hub and the board.

    ``B`` is written **scaled** (see ``browser_lora``), so the graph multiplies by one.
    """
    a, b = stack_adapter(model)
    layout = adapter_layout(model)
    return _write_web(a, b, layout.scale, layout.module, layout.r, out, base_repo)

def write_web_adapter_from_file(
    adapter: Path | str, out: Path | str, base_repo: str | None = None
) -> Path:
    """The same file, read straight from ``adapter.safetensors`` without a base model.

    Publishing an adapter should not need 440 MB of weights on disk to write 1.6 MB of floats,
    and the saved file already has everything: the keys carry the block number and the slice
    index, and ``adapter_config.json`` next to it carries the rank and the alpha. The order the
    tensors are stacked in is the order the blocks run, which is what the graph indexes by.
    """
    from safetensors.torch import load_file

    source = Path(adapter)
    cfg = LoraConfig.model_validate_json(
        (source.with_name(ADAPTER_CONFIG)).read_text(encoding="utf-8")
    )
    tensors = load_file(source.as_posix())
    a_keys = sorted(
        (key for key in tensors if ".lora_A." in key), key=lambda key: _adapter_sort_key(key)
    )
    if not a_keys:
        raise ValueError(f"{source} has no `lora_A` tensors")
    modules = sorted({key.split(".lora_A.")[0] for key in a_keys}, key=_adapter_sort_key)
    suffixes = {name.split(".", 2)[-1] for name in modules}
    if len(suffixes) != 1:
        raise ValueError(f"adapters on more than one matrix per block ({sorted(suffixes)})")
    per_layer = len(a_keys) // len(modules)
    a = torch.stack(
        [
            torch.stack([tensors[f"{name}.lora_A.{index}"] for index in range(per_layer)])
            for name in modules
        ]
    )
    b = torch.stack(
        [
            torch.stack([tensors[f"{name}.lora_B.{index}"] for index in range(per_layer)])
            for name in modules
        ]
    )
    return _write_web(a, b, cfg.scale, suffixes.pop(), cfg.r, out, base_repo)


def _adapter_sort_key(key: str) -> tuple[int, int]:
    """``blocks.10.attn.qkv.lora_A.1`` -> ``(10, 1)``: numeric, so 10 does not sort before 2."""
    parts = key.split(".")
    block = next((int(part) for part in parts if part.isdigit()), 0)
    index = int(parts[-1]) if parts[-1].isdigit() and ".lora_" in key else 0
    return block, index

def restore_lora_modules(model: MoveDecoder, cfg: LoraConfig) -> int:
    """Put ordinary ``LoRALinear`` wrappers back where ``build_adaptable`` left holders.

    Only the tests need this -- exporting is the end of a model's life in a process -- but they
    need it enough to keep it here rather than to reach into the internals from outside.
    """
    restored = 0
    for name, module in list(model.named_modules()):
        if not isinstance(module, SliceLoraFromInputs):
            continue
        wrapper = LoRALinear(module.base, cfg.r, cfg.alpha, cfg.dropout, module.slices)
        parent, _, leaf = name.rpartition(".")
        setattr(_module(model, parent), leaf, wrapper)
        restored += 1
    return restored

    """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.

    ``adapter_inputs`` writes the graph that takes its LoRA factors as two extra inputs
    (``rukh.export.adapter``) instead of a graph with fixed weights. Its parity is then read
    twice: once with an adapter of zeros, which has to reproduce the checkpoint exactly, and, if
    ``adapter`` names a trained one, once with that adapter fed to both sides.
    """
    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)
    from rukh.tokenize.uci_vocab import UciTokenizer
    from rukh.train import load_model

    model, _payload = load_model(Path(ckpt))
    block = model.cfg.block
    if adapter_inputs:
        # `export_adaptable_onnx` swaps the adapted matrices for holders, so the model it is
        # given stops being a plain decoder; everything downstream reloads from the checkpoint.
        lora = lora or LoraConfig()
        bundle = ExportBundle(
            onnx=export_adaptable_onnx(model, out, lora, opset=opset, seq_len=seq_len),
            adapter_inputs=True,
        )
        base_feed = zero_inputs(_layout_of(Path(ckpt), lora))
    else:
        bundle = ExportBundle(onnx=export_onnx(model, out, opset=opset, seq_len=seq_len))
        base_feed = None
    checks = _quantize(bundle, fp16, int8)
    if check_parity:
        prefixes, source, warning = parity_positions(
            UciTokenizer(), n=positions, block=block, seed=seed, games=games
        )
        bundle.parity_source = source
        bundle.parity_warning = warning
        # The reference for the zero-adapter parity is the plain checkpoint: with no correction
        # the adaptable file has to be the model it was exported from, and nothing else.
        plain = load_model(Path(ckpt))[0] if adapter_inputs else model
        bundle.parity = {
            name: parity(plain, Path(path), prefixes, n=positions, extra=base_feed)
            for name, path in checks.items()
        }
        if adapter_inputs and adapter is not None:
            bundle.adapter_name = Path(adapter).name
            bundle.adapter_parity = _adapter_parity(ckpt, adapter, checks, prefixes, positions)
        bundle.parity_path = write_parity(bundle, checks, out)
    return bundle

def adapter_file(adapter: Path | str) -> Path:
    """``adapter.safetensors`` inside a training run, or the file itself when given directly.

    ``rukh export --adapter`` takes the *folder* a run wrote, because that is what a reader has in
    front of them, while ``load_adapter`` takes the weights file and finds the config next to it.
    Handing it the folder makes it look for ``adapter_config.json`` one level too high and fail on
    a path nobody wrote, so the translation happens once and here.
    """
    path = Path(adapter)
    return path / ADAPTER_FILE if path.is_dir() else path


def _layout_of(ckpt: Path, lora: LoraConfig) -> AdapterLayout:
    """The shape of the adapters a checkpoint's adaptable graph accepts."""
    from rukh.train import load_model

    model, _payload = load_model(Path(ckpt))
    apply_lora(model, lora)
    return adapter_layout(model)


def _adapter_parity(
    ckpt: Path,
    adapter: Path,
    checks: dict[str, str],
    prefixes: list[list[int]],
    positions: int,
) -> dict[str, ParityResult]:
    """Parity again, with a real adapter on both sides: the claim the swapping rests on."""
    from rukh.models.lora import load_adapter
    from rukh.train import load_model

    adapted, _payload = load_model(Path(ckpt))
    load_adapter(adapted, adapter_file(adapter))
    feed = adaptable_inputs(adapted)
    return {
        name: parity(adapted, Path(path), prefixes, n=positions, extra=feed)
        for name, path in checks.items()
    }

"""The adaptable graph: zeros are the base model, and the factors really are inputs."""

from __future__ import annotations

import pathlib

import numpy as np
import pytest
import torch

from rukh.export import export_onnx
from rukh.export.adapter import (
    A_INPUT,
    B_INPUT,
    WEB_ADAPTER_META,
    adaptable_inputs,
    browser_lora,
    build_adaptable,
    export_adaptable_onnx,
    restore_lora_modules,
    write_web_adapter,
    zero_inputs,
)
from rukh.export.onnx import INPUT_NAME, read_metadata
from rukh.models import DecoderConfig, MoveDecoder
from rukh.models.lora import LoraConfig, adapter_layout, apply_lora, lora_modules

ort = pytest.importorskip("onnxruntime")

TOY = DecoderConfig(vocab_size=64, n_layer=2, n_head=2, d_model=32, block=16, dropout=0.0)
LORA = LoraConfig(r=4, alpha=8, targets=("q", "v"))
SEQ = 8  # the toy block is 16; the exporter refuses an example longer than the model


def _toy(seed: int = 0) -> MoveDecoder:
    torch.manual_seed(seed)
    return MoveDecoder(TOY).eval()


def _trained_adapter(model: MoveDecoder) -> None:
    """Give every ``B`` something other than zero, which is what training would do."""
    torch.manual_seed(7)
    with torch.no_grad():
        for _, adapter in lora_modules(model):
            for b in adapter.b:
                b.normal_(std=0.05)


def _session(path) -> ort.InferenceSession:
    return ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])


def test_the_graph_declares_the_two_factors_as_inputs(tmp_path):
    result = export_adaptable_onnx(_toy(), tmp_path / "model.onnx", LORA, seq_len=SEQ)
    session = _session(result.path)
    assert [i.name for i in session.get_inputs()] == [INPUT_NAME, A_INPUT, B_INPUT]
    metadata = read_metadata(result.path)
    assert metadata["rukh_adapter_inputs"] == f"{A_INPUT},{B_INPUT}"
    assert metadata["rukh_adapter_shape_a"] == "2,2,4,32"
    assert metadata["rukh_adapter_shape_b"] == "2,2,32,4"
    assert metadata["rukh_adapter_module"] == "attn.qkv"

def test_zeros_are_the_base_model(tmp_path):
    """The claim that lets the adaptable file replace the ordinary one instead of joining it."""
    plain = export_onnx(_toy(), tmp_path / "plain" / "model.onnx", seq_len=SEQ)
    adaptable = export_adaptable_onnx(_toy(), tmp_path / "lora" / "model.onnx", LORA, seq_len=SEQ)
    idx = torch.randint(1, 64, (1, 9)).numpy()
    layout = adapter_layout(_with_lora(_toy()))

    base = _session(plain.path).run(None, {INPUT_NAME: idx})[0]
    adapted = _session(adaptable.path).run(None, {INPUT_NAME: idx, **zero_inputs(layout)})[0]
    assert np.abs(base - adapted).max() < 1e-5


def test_the_adapter_changes_the_logits_without_touching_the_file(tmp_path):
    model = _with_lora(_toy())
    _trained_adapter(model)
    feed = adaptable_inputs(model)
    result = export_adaptable_onnx(_toy(), tmp_path / "model.onnx", LORA, seq_len=SEQ)
    layout = adapter_layout(_with_lora(_toy()))
    idx = torch.randint(1, 64, (1, 9)).numpy()

    session = _session(result.path)
    without = session.run(None, {INPUT_NAME: idx, **zero_inputs(layout)})[0]
    with_adapter = session.run(None, {INPUT_NAME: idx, **feed})[0]
    assert np.abs(with_adapter - without).max() > 1e-3


def test_the_graph_agrees_with_pytorch_on_the_same_adapter(tmp_path):
    """Parity: the file and ``apply_lora`` have to compute the same next move."""
    model = _with_lora(_toy())
    _trained_adapter(model)
    idx = torch.randint(1, 64, (1, 9))
    with torch.no_grad():
        expected = model(idx)[0][:, -1, :].numpy()

    result = export_adaptable_onnx(_toy(), tmp_path / "model.onnx", LORA, seq_len=SEQ)
    feed = adaptable_inputs(model)
    got = _session(result.path).run(None, {INPUT_NAME: idx.numpy(), **feed})[0]
    assert np.abs(got - expected).max() < 1e-4


def test_two_adapters_give_two_different_answers_from_one_file(tmp_path):
    """Swapping is the whole feature: one download, as many styles as there are files."""
    result = export_adaptable_onnx(_toy(), tmp_path / "model.onnx", LORA, seq_len=SEQ)
    session = _session(result.path)
    idx = torch.randint(1, 64, (1, 9)).numpy()
    answers = []
    for seed in (1, 2):
        model = _with_lora(_toy())
        torch.manual_seed(seed)
        with torch.no_grad():
            for _, adapter in lora_modules(model):
                for b in adapter.b:
                    b.normal_(std=0.1)
        answers.append(session.run(None, {INPUT_NAME: idx, **adaptable_inputs(model)})[0])
    assert np.abs(answers[0] - answers[1]).max() > 1e-3

def test_the_sequence_axis_survives_the_extra_inputs(tmp_path):
    result = export_adaptable_onnx(_toy(), tmp_path / "model.onnx", LORA, seq_len=8)
    assert result.dynamic_seq is True
    assert result.dynamic_seq_verified is True


def test_an_adapter_that_cannot_be_stacked_says_so():
    """The MLP's ``fc`` is four times as wide, so the factors stop being one tensor."""
    model = _toy()
    apply_lora(model, LoraConfig(r=4, alpha=8, targets=("q", "mlp_in")))
    with pytest.raises(ValueError, match="more than one matrix per block"):
        adapter_layout(model)


def test_the_holders_can_be_turned_back_into_ordinary_adapters():
    model = _toy()
    build_adaptable(model, LORA)
    assert not list(lora_modules(model))
    assert restore_lora_modules(model, LORA) == TOY.n_layer
    assert len(list(lora_modules(model))) == TOY.n_layer


def _with_lora(model: MoveDecoder) -> MoveDecoder:
    apply_lora(model, LORA)
    return model

def test_the_browser_file_feeds_the_unit_scale_graph_and_matches_pytorch(tmp_path):
    """The whole browser path end to end: write the flat file, read it back, run the graph.

    The scale lives in the file and not in the graph (``browser_lora``), so this is also the
    check that the convention holds: a ``B`` written scaled, fed to a graph that multiplies by
    one, has to reproduce the model that applies ``alpha / r`` itself.
    """
    import json

    model = _with_lora(_toy())
    _trained_adapter(model)
    idx = torch.randint(1, 64, (1, 9))
    with torch.no_grad():
        expected = model(idx)[0][:, -1, :].numpy()

    binary = write_web_adapter(model, tmp_path / "web" / "adapter.bin", base_repo="chorcat/toy")
    meta = json.loads((tmp_path / "web" / pathlib.Path(WEB_ADAPTER_META).name).read_text("utf-8"))
    assert meta["shape_a"] == [2, 2, 4, 32]
    assert meta["scale_applied"] == LORA.scale

    flat = np.frombuffer(binary.read_bytes(), dtype="<f4")
    size_a = int(np.prod(meta["shape_a"]))
    feed = {
        A_INPUT: flat[:size_a].reshape(meta["shape_a"]).copy(),
        B_INPUT: flat[size_a:].reshape(meta["shape_b"]).copy(),
    }

    result = export_adaptable_onnx(_toy(), tmp_path / "model.onnx", browser_lora(LORA), seq_len=SEQ)
    got = _session(result.path).run(None, {INPUT_NAME: idx.numpy(), **feed})[0]
    assert np.abs(got - expected).max() < 1e-4


def test_the_web_file_is_the_same_from_the_model_and_from_the_saved_adapter(tmp_path):
    """Publishing must not need 440 MB of base weights on disk to write 1.6 MB of floats."""
    from rukh.export.adapter import write_web_adapter_from_file
    from rukh.models.lora import save_adapter

    model = _with_lora(_toy())
    _trained_adapter(model)
    save_adapter(model, tmp_path / "adapter.safetensors", LORA)

    from_model = write_web_adapter(model, tmp_path / "a" / "adapter.bin", "chorcat/toy")
    from_file = write_web_adapter_from_file(
        tmp_path / "adapter.safetensors", tmp_path / "b" / "adapter.bin", "chorcat/toy"
    )
    assert from_model.read_bytes() == from_file.read_bytes()

def test_the_quantized_files_still_take_the_adapter(tmp_path):
    """int8 and fp16 are what the browser actually loads, so the swap has to survive both.

    Dynamic quantization rewrites the `MatMul`s whose other side is an initialiser, and the
    adapter's two are not: both of their inputs come from outside the graph. Nothing guarantees
    the converters leave them alone, so it is checked rather than assumed.
    """
    from rukh.export.quantize import quantize_int8, to_fp16

    model = _with_lora(_toy())
    _trained_adapter(model)
    feed = adaptable_inputs(model)
    result = export_adaptable_onnx(_toy(), tmp_path / "model.onnx", LORA, seq_len=SEQ)
    layout = adapter_layout(_with_lora(_toy()))
    idx = torch.randint(1, 64, (1, 9)).numpy()

    for quantized in (to_fp16(pathlib.Path(result.path)), quantize_int8(pathlib.Path(result.path))):
        session = _session(quantized.path)
        assert [i.name for i in session.get_inputs()] == [INPUT_NAME, A_INPUT, B_INPUT]
        without = session.run(None, {INPUT_NAME: idx, **zero_inputs(layout)})[0]
        with_adapter = session.run(None, {INPUT_NAME: idx, **feed})[0]
        assert np.abs(with_adapter - without).max() > 1e-3, quantized.kind


def test_the_adapter_parity_takes_the_folder_a_run_wrote(tmp_path):
    """`--adapter` names a training run's folder; `load_adapter` wants the weights file.

    The translation used to be missing, so the export looked for `adapter_config.json` one level
    above the run and died on a path nobody had written -- after fifteen minutes of exporting.
    """
    from rukh.export import adapter_file
    from rukh.models.lora import ADAPTER_FILE, save_adapter

    model = _with_lora(_toy())
    run = tmp_path / "lora-e4-20260920-160133"
    save_adapter(model, run / ADAPTER_FILE, LORA)

    assert adapter_file(run) == run / ADAPTER_FILE
    assert adapter_file(run / ADAPTER_FILE) == run / ADAPTER_FILE
    # And the file it points at is the one `load_adapter` can actually read.
    plain = _toy()
    from rukh.models.lora import load_adapter

    assert load_adapter(plain, adapter_file(run)).r == LORA.r

def _align_constant_outputs(model: object) -> int:
    """Retype every constant whose value still says float32 while its output says float16.

    The same fault as ``_align_cast_outputs`` in a different node. ``convert_float_to_float16``
    retypes the graph's values but leaves the ``value`` attribute of ``Constant`` and
    ``ConstantOfShape`` alone, so a node that produces a block of zeros keeps producing float32
    ones into a graph that now expects float16, and onnxruntime refuses the file with the same
    ``Type Error`` message.

    The decoder that takes its LoRA factors as inputs is what found this: the correction is
    assembled with ``cat`` over the projection's full width, and the ranges nobody adapts are a
    ``ConstantOfShape`` of zeros. Without this the fp16 build of every swappable model would fail
    to load, and the fp32 one -- which nothing serves -- would be fine.

    A ``ConstantOfShape`` with **no** ``value`` attribute is the same fault written as an absence:
    the specification says its default is a float32 zero, so the node the exporter wrote without
    one quietly stops matching a graph that is now float16. That case gets the attribute it was
    relying on the default for.
    """
    import numpy as np
    from onnx import TensorProto, helper, numpy_helper

    graph = model.graph  # type: ignore[attr-defined]
    declared = {value.name: value.type.tensor_type.elem_type for value in graph.value_info}
    kinds = {TensorProto.FLOAT: np.float32, TensorProto.FLOAT16: np.float16}
    repaired = 0
    for node in graph.node:
        if node.op_type not in {"Constant", "ConstantOfShape"} or not node.output:
            continue
        wanted = declared.get(node.output[0])
        if wanted not in kinds:
            continue
        value = next((entry for entry in node.attribute if entry.name == "value"), None)
        if value is None:
            if node.op_type != "ConstantOfShape" or wanted == TensorProto.FLOAT:
                continue  # `Constant` always has one; the default is already float32
            fill = numpy_helper.from_array(np.zeros(1, dtype=kinds[wanted]), "value")
            node.attribute.append(helper.make_attribute("value", fill))
            repaired += 1
            continue
        if value.t.data_type not in kinds or value.t.data_type == wanted:
            continue
        array = numpy_helper.to_array(value.t).astype(kinds[wanted])
        value.t.CopyFrom(numpy_helper.from_array(array, value.t.name))
        repaired += 1
    return repaired

def parity(
    ckpt: Path | MoveDecoder,
    onnx_path: Path,
    positions: Sequence[Sequence[int]],
    n: int = DEFAULT_N,
    extra: dict[str, object] | None = None,
) -> ParityResult:
    """Compare the argmax token and the logits of the checkpoint and the exported file.

    ``extra`` is the rest of the feed, for a graph that takes more than the token ids: the
    adaptable decoder wants its two LoRA tensors, and the model on the PyTorch side is expected
    to carry the same adapter already (``load_adapter``) so that both sides answer the same
    question. Zeros on one side and a trained adapter on the other would measure the adapter,
    not the export.
    """

"""Publishing a LoRA adapter, which is not a model and must not pretend to be one.

A model repository in this project carries weights, three ONNX exports, a vocabulary and a card
whose numbers come from ``rukh eval``. An adapter has none of that. It is two matrices per adapted
projection -- 1.6 MB against the 440 MB of the weights it mounts on -- and on its own it does
nothing at all: without the exact base checkpoint it is a pile of numbers with no model to correct.

So the card leads with the base model and the command that loads it, and the metrics it publishes
are the ones that say what the adapter *changed*: the share of first moves it shifted, which is the
whole visible point of a style adapter, and the Elo of base-plus-adapter against the base alone,
which is what that style cost in strength. Publishing an adapter with the base model's Elo and no
mention of the adapter's own effect would be publishing somebody else's number.

Two copies of the same numbers go out. ``adapter.safetensors`` is for PyTorch and is what
``load_adapter`` reads. ``web/adapter.bin`` is the same factors as one flat little-endian float32
buffer, with the ``alpha / r`` already folded in, for the demo: the browser feeds it straight into
the adaptable ONNX graph (``rukh.export.adapter``) and changes style without downloading another
221 MB of model. The two are written from the same tensors by the same function, so they cannot
drift apart.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict

from rukh.config import BaseConfig
from rukh.export.adapter import WEB_ADAPTER_FILE, WEB_ADAPTER_META, write_web_adapter_from_file
from rukh.models.lora import ADAPTER_CONFIG, ADAPTER_FILE, LoraConfig
from rukh.paths import resolve
from rukh.publish.model import README_NAME, ModelPublishConfig, _api, read_eval, render_card

log = logging.getLogger(__name__)

__all__ = ["AdapterCost", "AdapterPublishResult", "publish_adapter", "publish_qwen_adapter"]

REPO_TYPE = "model"
ADAPTER_CARD_TEMPLATE = "adapter.md.jinja"

class AdapterCost(BaseConfig):
    """One metric read on the base model and on the adapted one: same suite, same seed.

    The card's most important table, and the one an adapter published on its own would not have.
    A style that changes the opening is easy; a style that changes the opening and nothing else
    is the claim worth making, and it is only a claim if the "nothing else" is measured.
    """

    metric: str
    base: float
    adapted: float
    unit: Literal["percent", "elo", "bits"] = "percent"

    def render(self) -> tuple[str, str, str]:
        """``(metric, base, adapted)`` as the card prints them."""
        if self.unit == "percent":
            return self.metric, f"{self.base * 100:.1f} %", f"{self.adapted * 100:.1f} %"
        if self.unit == "elo":
            return self.metric, f"{self.base:.0f}", f"{self.adapted:.0f}"
        return self.metric, f"{self.base:.4f} bits", f"{self.adapted:.4f} bits"


class AdapterEffect(BaseConfig):
    """What the adapter measurably did, as opposed to what it was meant to do."""

    first_move: str | None = None
    """UCI of the opening move the style is about, e.g. ``e2e4``."""
    share_before: float | None = None
    share_after: float | None = None
    """Probability the model gives ``first_move`` from the opening position, base and adapted.

    Read straight off the softmax over the twenty legal first moves, with no sampling, no
    temperature and no seed -- so two readings agree to the last decimal and the number says
    something about the weights rather than about a draw. Counting self-play openings instead
    would answer the same question with sampling noise on top, which for a card is strictly
    worse: a published number that moves between runs invites the reader to average it."""
    cost: list[AdapterCost] = []
    """Every metric measured on both sides, in the order the card prints them."""
    entropy_before: float | None = None
    entropy_after: float | None = None
    """First-move entropy in bits, the same measurement in one number instead of one move."""
    train_games: int | None = None
    """Games in the style slice the adapter was trained on."""
    predicate: str | None = None
    """How that slice was selected, verbatim from the style manifest."""


class AdapterPublishResult(BaseModel):
    """What was staged and, unless this was a dry run, uploaded."""

    model_config = ConfigDict(extra="forbid")

    repo_id: str
    base_repo: str
    stage: str
    dry_run: bool
    folder: str
    card_path: str
    files: list[str]
    params: int
    bytes: int


def adapter_params(lora: LoraConfig, n_layer: int, d_model: int) -> int:
    """``2 * r * d_model`` per adapted matrix per layer: the number the card leads with."""
    return n_layer * len(lora.targets) * 2 * lora.r * d_model

def publish_adapter(
    run_dir: Path | str,
    repo: str,
    base_repo: str,
    cfg: ModelPublishConfig | None = None,
    stage: str | None = None,
    effect: AdapterEffect | None = None,
    base_config: dict[str, Any] | None = None,
    dry_run: bool = False,
) -> AdapterPublishResult:
    """Stage (and unless ``dry_run``, upload) one LoRA adapter as its own small repository."""
    cfg = cfg or ModelPublishConfig()
    source = Path(run_dir)
    weights = source / ADAPTER_FILE
    if not weights.is_file():
        raise FileNotFoundError(f"no adapter at {weights}")
    lora = LoraConfig.model_validate_json((source / ADAPTER_CONFIG).read_text(encoding="utf-8"))

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

    folder = resolve(cfg.publish_dir) / repo_id
    folder.mkdir(parents=True, exist_ok=True)
    (folder / ADAPTER_FILE).write_bytes(weights.read_bytes())
    web = write_web_adapter_from_file(weights, folder / WEB_ADAPTER_FILE, base_id)
    (folder / ADAPTER_CONFIG).write_text(
        json.dumps(
            {"lora": lora.model_dump(mode="json"), "base_model": base_id, "format": "rukh-lora-1"},
            indent=2,
        )
        + "\n",
        encoding="utf-8",
        newline="\n",
    )

    config = base_config or {}
    params = adapter_params(lora, int(config.get("n_layer", 16)), int(config.get("d_model", 768)))
    card = render_card(
        {
            "repo_id": repo_id,
            "base_repo": base_id,
            "stage": name,
            "license": cfg.license,
            "lora": lora.model_dump(mode="json"),
            "targets": ", ".join(lora.targets),
            "params": params,
            "bytes": weights.stat().st_size,
            "effect": (effect or AdapterEffect()).model_dump(mode="json"),
            "cost": [item.render() for item in (effect or AdapterEffect()).cost],
            "web_file": WEB_ADAPTER_FILE,
            "web_bytes": web.stat().st_size,
            "demo_url": cfg.demo_url,
            "course_url": cfg.course_url,
            "repository_url": cfg.repository_url,
            "datasets": cfg.datasets,
        },
        ADAPTER_CARD_TEMPLATE,
    )
    card_path = folder / README_NAME
    card_path.write_text(card, encoding="utf-8", newline="\n")

    files = [ADAPTER_FILE, ADAPTER_CONFIG, README_NAME, WEB_ADAPTER_FILE, WEB_ADAPTER_META]
    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 AdapterPublishResult(
        repo_id=repo_id,
        base_repo=base_id,
        stage=name,
        dry_run=dry_run,
        folder=folder.as_posix(),
        card_path=card_path.as_posix(),
        files=files,
        params=params,
        bytes=weights.stat().st_size,
    )

QWEN_CARD_TEMPLATE = "qwen-adapter.md.jinja"


def _percent(value: float | None) -> str:
    return "n/a" if value is None else f"{value * 100:.1f} %"


def qwen_card_context(
    repo_id: str,
    run_dir: Path,
    cfg: ModelPublishConfig,
    evaluation: dict[str, Any] | None,
) -> dict[str, Any]:
    """Everything the Qwen adapter's card says, read off the run and the evaluation.

    Nothing here is written by hand. ``run.json`` is what ``rukh train qwen`` recorded about
    itself -- including whether four bits were actually used -- and the evaluation is the report
    of ``rukh eval qwen``, which ran the project's own harness. A card that quoted the intended
    recipe instead of the executed one would be the wrong kind of document.
    """
    run = json.loads((run_dir / "run.json").read_text(encoding="utf-8"))
    peft_config = json.loads((run_dir / ADAPTER_CONFIG).read_text(encoding="utf-8"))
    results: list[tuple[str, str]] = []
    failures: list[tuple[str, str]] = []
    if evaluation:
        written = evaluation.get("written") or {}
        asked = int(written.get("asked") or 0)

        def share(count: int) -> str:
            return f"{100 * count / asked:.2f} %" if asked else "n/a"

        elo = evaluation.get("elo") or {}
        results = [
            ("Legal moves written, no mask", share(int(written.get("legal") or 0))),
            ("Top-1 next move", _percent(evaluation.get("top1"))),
            ("Puzzles solved", _percent((evaluation.get("puzzles") or {}).get("rate"))),
            ("Estimated Elo", _elo_text(elo)),
        ]
        failures = [
            ("legal", share(int(written.get("legal") or 0))),
            ("illegal here", share(int(written.get("illegal") or 0))),
            ("not a move at all", share(int(written.get("unparseable") or 0))),
            ("ambiguous SAN", share(int(written.get("ambiguous") or 0))),
            ("nothing written", share(int(written.get("empty") or 0))),
        ]
    return {
        "repo_id": repo_id,
        "base_model": str(peft_config.get("base_model_name_or_path", run.get("model"))),
        "license": cfg.license,
        "lora": peft_config,
        "targets": ", ".join(sorted(peft_config.get("target_modules") or [])),
        "trainable_params": int(run.get("trainable_params") or 0),
        "total_params": int(run.get("total_params") or 1),
        "train_samples": int(run.get("train_samples") or 0),
        "steps": int(run.get("steps") or 0),
        "four_bit": bool(run.get("four_bit")),
        "weights_memory_mb": run.get("weights_memory_mb"),
        "peak_memory_mb": run.get("peak_memory_mb"),
        "results": results,
        "failures": failures,
        "course_url": cfg.course_url,
        "repository_url": cfg.repository_url,
    }

def _elo_text(elo: dict[str, Any]) -> str:
    if not elo:
        return "n/a"
    if elo.get("separated"):
        if elo.get("elo_lower") is not None:
            return f"> {elo['elo_lower']:.0f} (one-sided 95 % bound; every game won)"
        if elo.get("elo_upper") is not None:
            return f"< {elo['elo_upper']:.0f} (one-sided 95 % bound; every game lost)"
    low, high = elo.get("ci_low"), elo.get("ci_high")
    if low is None or high is None:
        return f"{elo.get('elo', 0):.0f} (no interval)"
    return f"{elo['elo']:.0f} (95 % CI {low:.0f}-{high:.0f})"


def publish_qwen_adapter(
    run_dir: Path | str,
    repo: str,
    cfg: ModelPublishConfig | None = None,
    stage: str = "qwen3-pgn-qlora",
    dry_run: bool = False,
) -> AdapterPublishResult:
    """Upload a ``peft`` adapter folder as it stands, plus a card built from the run's own record.

    Unlike the project's own adapters this one is **not** re-staged: ``peft`` already wrote a
    valid ``adapter_config.json`` and ``adapter_model.safetensors``, and the Hub knows that
    format. Rewriting them into our format would make the file unusable with the two lines of
    ``peft`` that every reader would actually type.
    """
    cfg = cfg or ModelPublishConfig()
    source = Path(run_dir)
    weights = source / "adapter_model.safetensors"
    if not weights.is_file():
        raise FileNotFoundError(f"no peft adapter at {weights}")
    repo_id = repo if "/" in repo else f"{cfg.owner}/{repo}"

    evaluation = read_eval(stage, cfg)
    context = qwen_card_context(repo_id, source, cfg, evaluation)
    (source / README_NAME).write_text(
        render_card(context, QWEN_CARD_TEMPLATE), encoding="utf-8", newline="\n"
    )
    files = sorted(path.name for path in source.iterdir() if path.is_file())

    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(source),
            repo_type=REPO_TYPE,
            commit_message=f"Publish {stage}",
        )
    return AdapterPublishResult(
        repo_id=repo_id,
        base_repo=str(context["base_model"]),
        stage=stage,
        dry_run=dry_run,
        folder=source.as_posix(),
        card_path=(source / README_NAME).as_posix(),
        files=files,
        params=int(context["trainable_params"]),
        bytes=weights.stat().st_size,
    )