Exportar a ONNX

from __future__ import annotations

import contextlib
import logging
import sys
from collections.abc import Iterator
from pathlib import Path
from typing import Any, Literal

import torch
from pydantic import BaseModel, ConfigDict
from torch import Tensor, nn

from rukh import __version__
from rukh.models import MoveDecoder

log = logging.getLogger(__name__)

MODEL_NAME = "model.onnx"
INPUT_NAME = "idx"
OUTPUT_NAME = "logits"
BATCH_AXIS = "batch"
SEQUENCE_AXIS = "sequence"
DEFAULT_OPSET = 18
METADATA_PREFIX = "rukh_"

class ExportResult(BaseModel):
    """What was written and how."""

    model_config = ConfigDict(extra="forbid")

    path: str
    exporter: Literal["dynamo", "legacy"]
    opset: int
    seq_len: int
    block: int
    """The model's context: the exported graph must never be fed more than this many tokens."""
    dynamic_batch: bool
    dynamic_seq: bool
    """What the file really accepts, not what was asked for: verified by running it."""
    dynamic_seq_verified: bool = False
    """Whether ``dynamic_seq`` was checked by running the file (needs ``onnxruntime``)."""
    metadata: dict[str, str] = {}
    """``metadata_props`` written into the file, empty when ``onnx`` is not installed."""
    vocab_size: int
    params: int
    bytes: int
    warning: str | None = None

def target_path(out: Path) -> Path:
    """``out`` as a file: a directory gets ``model.onnx`` inside it."""
    out = Path(out)
    return out if out.suffix == ".onnx" else out / MODEL_NAME

def _dynamic_shapes(dynamic_batch: bool, dynamic_seq: bool) -> dict[str, dict[int, str]] | None:
    axes: dict[int, str] = {}
    if dynamic_batch:
        axes[0] = BATCH_AXIS
    if dynamic_seq:
        axes[1] = SEQUENCE_AXIS
    return {INPUT_NAME: axes} if axes else None


def _dynamic_axes(dynamic_batch: bool, dynamic_seq: bool) -> dict[str, dict[int, str]] | None:
    shapes = _dynamic_shapes(dynamic_batch, dynamic_seq)
    if shapes is None:
        return None
    axes = dict(shapes)
    if dynamic_batch:
        axes[OUTPUT_NAME] = {0: BATCH_AXIS}
    return axes

@contextlib.contextmanager
def _utf8_console() -> Iterator[None]:
    """Let the exporter print its progress ticks on a legacy Windows code page.

    ``torch.onnx.export(dynamo=True)`` writes check marks to stdout. On Windows the console
    is cp1252 by default, so that raises ``UnicodeEncodeError`` inside the exporter and the
    whole export fails for a reason that has nothing to do with the model. Reconfiguring the
    streams to replace unencodable characters keeps the modern exporter usable; without it
    every Windows export silently falls back to the deprecated TorchScript one, which in turn
    produces an fp16 graph that onnxruntime refuses to load.
    """
    streams = [s for s in (sys.stdout, sys.stderr) if hasattr(s, "reconfigure")]
    previous = [(s, s.encoding, s.errors) for s in streams]
    for stream in streams:
        with contextlib.suppress(Exception):
            stream.reconfigure(errors="replace")
    try:
        yield
    finally:
        for stream, encoding, errors in previous:
            with contextlib.suppress(Exception):
                stream.reconfigure(encoding=encoding, errors=errors)

def export_onnx(
    ckpt: Path | MoveDecoder,
    out: Path,
    opset: int = DEFAULT_OPSET,
    dynamic_batch: bool = True,
    seq_len: int = 200,
    dynamic_seq: bool = True,
) -> ExportResult:
    """Export the next-move head of a checkpoint to ONNX and report which exporter ran.

    ``ckpt`` may be a checkpoint path or an already-built model (the tests use the latter).
    ``seq_len`` is the length of the example the exporter traces; with ``dynamic_seq`` the graph
    still accepts any length up to the model's block, which is what the demo feeds it as a game
    grows move by move.
    """
    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 = LastStepLogits(model).eval()
    example = torch.zeros((1, seq_len), dtype=torch.long)
    path = target_path(out)
    path.parent.mkdir(parents=True, exist_ok=True)

    common: dict[str, Any] = {
        "input_names": [INPUT_NAME],
        "output_names": [OUTPUT_NAME],
        "opset_version": opset,
    }
    warning: str | None = None
    exporter: Literal["dynamo", "legacy"] = "dynamo"
    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)
        exporter = "legacy"
        with torch.no_grad():
            torch.onnx.export(
                wrapper,
                (example,),
                str(path),
                dynamo=False,
                dynamic_axes=_dynamic_axes(dynamic_batch, dynamic_seq),
                **common,
            )

    works = verify_dynamic_seq(path, seq_len, model.cfg.block)
    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. Re-export with "
            "dynamic_seq=False and pad the input, or fix the exporter."
        )
    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__,
        },
    )
    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 sequence_lengths(seq_len: int, block: int) -> list[int]:
    """Two lengths to run the exported file at: the traced one and a different, legal one."""
    other = seq_len + 1 if seq_len < block else max(1, seq_len - 1)
    return sorted({seq_len, other})

def verify_dynamic_seq(path: Path, seq_len: int, block: int) -> bool | None:
    """Run the file at two sequence lengths; None when ``onnxruntime`` is not installed."""
    import numpy as np

    try:
        import onnxruntime as ort
    except ImportError:
        log.warning("onnxruntime is not installed: the dynamic sequence axis was not verified")
        return None
    lengths = sequence_lengths(seq_len, block)
    if len(lengths) < 2:  # pragma: no cover - block of 1 is not a usable model
        return None
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    for length in lengths:
        try:
            session.run(None, {INPUT_NAME: np.zeros((1, length), dtype=np.int64)})
        except Exception as exc:  # noqa: BLE001 - any refusal means the axis is not dynamic
            log.info("the exported graph refused a sequence of %d tokens: %s", length, exc)
            return False
    return True

def set_metadata(path: Path, entries: dict[str, str]) -> dict[str, str]:
    """Write ``metadata_props`` verbatim into a file; empty when ``onnx`` is not installed."""
    try:
        import onnx
    except ImportError:
        log.warning("onnx is not installed: the model metadata (block, vocab) was not written")
        return {}
    model = onnx.load(str(path))
    onnx.helper.set_model_props(model, entries)
    onnx.save(model, str(path))
    # The dynamo exporter writes the weights next to the graph as `<name>.onnx.data` and points
    # the initializers at it. `onnx.load` above pulled them into memory and `onnx.save` wrote
    # them back inline, so the sidecar is now dead weight that would ship to the browser (or be
    # published to the Hub) without ever being read. Drop it, but only once the file really is
    # self-contained.
    sidecar = path.with_suffix(path.suffix + ".data")
    if sidecar.is_file() and not any(
        tensor.data_location == onnx.TensorProto.EXTERNAL for tensor in model.graph.initializer
    ):
        sidecar.unlink()
        log.debug("removed the external-data sidecar %s", sidecar.name)
    return entries

def write_metadata(path: Path, props: dict[str, Any]) -> dict[str, str]:
    """Write ``rukh_*`` metadata into the file; empty when ``onnx`` is not installed."""
    return set_metadata(
        path, {f"{METADATA_PREFIX}{key}": str(value) for key, value in props.items()}
    )


def read_metadata(path: Path) -> dict[str, str]:
    """The ``rukh_*`` metadata of an exported file (needs ``onnx``)."""
    import onnx

    model = onnx.load(str(path))
    return {entry.key: entry.value for entry in model.metadata_props}

def _load(ckpt: Path) -> MoveDecoder:
    from rukh.train import load_model

    model, _payload = load_model(ckpt)
    return model

from __future__ import annotations

import logging
from pathlib import Path
from typing import Literal

from pydantic import BaseModel, ConfigDict

log = logging.getLogger(__name__)

FP16_NAME = "model-fp16.onnx"
INT8_NAME = "model-int8.onnx"
QUANTIZED_OPS = ["MatMul", "Gemm"]

class QuantizeResult(BaseModel):
    """One converted file."""

    model_config = ConfigDict(extra="forbid")

    path: str
    kind: Literal["fp16", "int8"]
    method: str
    bytes: int
    source_bytes: int

    @property
    def ratio(self) -> float:
        return self.bytes / self.source_bytes if self.source_bytes else 0.0

def _sibling(path: Path, name: str, out: Path | None) -> Path:
    if out is None:
        return Path(path).with_name(name)
    out = Path(out)
    return out if out.suffix == ".onnx" else out / name

def to_fp16(path: Path, out: Path | None = None) -> QuantizeResult:
    """Convert an fp32 ONNX model to fp16, in place of ``model.onnx``'s sibling by default."""
    import onnx

    source = Path(path)
    target = _sibling(source, FP16_NAME, out)
    target.parent.mkdir(parents=True, exist_ok=True)
    model = onnx.load(str(source))
    try:
        from onnxconverter_common import float16

        converted = float16.convert_float_to_float16(model, keep_io_types=True)
        method = "onnxconverter_common.float16"
    except ImportError:
        converted = _cast_initializers(model)
        method = "onnx.numpy_helper fallback"
        log.warning("onnxconverter_common is not installed; used the initializer-casting fallback")
    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 _cast_initializers(model: object) -> object:
    """Fallback fp16 conversion: cast every float initializer and cast the outputs back.

    Only the weights change type. Every float-typed value in the graph is retyped to fp16 and a
    ``Cast`` node is appended so the model still hands the caller float32 logits, which keeps the
    demo's input and output contract identical to the fp32 file.
    """
    import numpy as np
    import onnx
    from onnx import TensorProto, helper, numpy_helper

    graph = model.graph  # type: ignore[attr-defined]
    for initializer in graph.initializer:
        if initializer.data_type != TensorProto.FLOAT:
            continue
        array = numpy_helper.to_array(initializer).astype(np.float16)
        initializer.CopyFrom(numpy_helper.from_array(array, initializer.name))
    for value in list(graph.value_info) + list(graph.input):
        if value.type.tensor_type.elem_type == TensorProto.FLOAT:
            value.type.tensor_type.elem_type = TensorProto.FLOAT16
    for node in graph.node:
        for attribute in node.attribute:
            if attribute.type == onnx.AttributeProto.TENSOR:
                tensor = attribute.t
                if tensor.data_type == TensorProto.FLOAT:
                    array = numpy_helper.to_array(tensor).astype(np.float16)
                    tensor.CopyFrom(numpy_helper.from_array(array, tensor.name))
    for output in graph.output:
        if output.type.tensor_type.elem_type != TensorProto.FLOAT:
            continue
        inner = f"{output.name}_fp16"
        for node in graph.node:
            node.output[:] = [inner if name == output.name else name for name in node.output]
        graph.node.append(
            helper.make_node("Cast", [inner], [output.name], to=int(TensorProto.FLOAT))
        )
    return model

from __future__ import annotations

import logging
import random
from collections.abc import Sequence
from pathlib import Path

import chess
import numpy as np
import torch
from pydantic import BaseModel, ConfigDict

from rukh.export.onnx import INPUT_NAME, LastStepLogits
from rukh.models import MoveDecoder
from rukh.tokenize.uci_vocab import UciTokenizer

log = logging.getLogger(__name__)

DEFAULT_N = 1_000
MAX_MISMATCHES = 20
DEFAULT_GAMES = "data/uci/year=2025/month=02/games.parquet"
"""The validation month of P1: the parity positions come from here when it exists."""

class ParityResult(BaseModel):
    """How well an exported file reproduces the checkpoint."""

    model_config = ConfigDict(extra="forbid")

    positions: int
    agreement: float
    max_abs_logit_delta: float
    mismatches: list[int]
    """Indices of the first few positions where the chosen token differed."""

def random_prefixes(
    tok: UciTokenizer,
    n: int,
    seed: int = 0,
    min_ply: int = 1,
    max_ply: int = 40,
) -> list[list[int]]:
    """Token-id prefixes of random legal games, for a parity check without any dataset.

    Every prefix is a real sequence of legal moves, so the model is asked the kind of question
    the demo asks it, and the walk is seeded so a parity number can be reproduced exactly.
    """
    rng = random.Random(seed)
    prefixes: list[list[int]] = []
    while len(prefixes) < n:
        board = chess.Board()
        ids = [
            tok.bos_id,
            tok.vocab["<w1800>"],
            tok.vocab["<b1800>"],
        ]
        plies = rng.randint(min_ply, max_ply)
        for _ in range(plies):
            moves = list(board.legal_moves)
            if not moves:
                break
            move = rng.choice(moves)
            ids.append(tok.vocab.get(move.uci(), tok.unk_id))
            board.push(move)
        prefixes.append(ids)
    return prefixes

def validation_prefixes(
    games: Path,
    tok: UciTokenizer,
    n: int = DEFAULT_N,
    block: int = 200,
    seed: int = 0,
) -> list[list[int]]:
    """Token-id prefixes of ``n`` validation positions, cropped exactly as the model sees them."""
    from rukh.eval.legality import sample_positions

    positions = sample_positions(games, n, tok, seed=seed, pool=max(n * 5, 10_000), block=block)
    return [list(position.history) for position in positions]

def parity_positions(
    tok: UciTokenizer,
    n: int = DEFAULT_N,
    block: int = 200,
    seed: int = 0,
    games: Path | None = None,
) -> tuple[list[list[int]], str, str | None]:
    """``(prefixes, source, warning)``: validation positions, or random walks with a warning."""
    from rukh import paths

    parquet = Path(games) if games is not None else paths.resolve(DEFAULT_GAMES)
    if parquet.is_file():
        prefixes = validation_prefixes(parquet, tok, n=n, block=block, seed=seed)
        if prefixes:
            return prefixes, "validation", None
    warning = (
        f"no validation games at {parquet}: parity was measured on random legal walks, which "
        "visit positions no human would reach"
    )
    log.warning("%s", warning)
    return random_prefixes(tok, n, seed=seed), "random-walk", warning

def _session(onnx_path: Path):  # type: ignore[no-untyped-def]
    """An ONNX Runtime session on the CPU provider (tests monkeypatch this)."""
    import onnxruntime as ort

    return ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])

"""Export: ONNX for the browser, fp16 and int8 quantization and PyTorch parity.

``export_all`` is the whole of ``rukh export``: write the fp32 graph, derive the fp16 and int8
files the two browser backends need, and check every file it produced against PyTorch.
"""

from __future__ import annotations

from pathlib import Path

from pydantic import BaseModel, ConfigDict

from rukh.export.onnx import (
    DEFAULT_OPSET,
    INPUT_NAME,
    MODEL_NAME,
    OUTPUT_NAME,
    ExportResult,
    LastStepLogits,
    export_onnx,
    read_metadata,
    sequence_lengths,
    set_metadata,
    target_path,
    verify_dynamic_seq,
    write_metadata,
)
from rukh.export.parity import (
    DEFAULT_GAMES,
    ParityResult,
    parity,
    parity_positions,
    random_prefixes,
    validation_prefixes,
)
from rukh.export.quantize import (
    FP16_NAME,
    INT8_NAME,
    QUANTIZED_OPS,
    QuantizeResult,
    quantize_int8,
    to_fp16,
)

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``."""
    parity_source: str | None = None
    """``validation`` (the positions of ``docs/spec/02`` §6) or ``random-walk``."""
    parity_warning: str | None = None

def export_all(
    ckpt: Path,
    out: Path,
    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,
) -> ExportBundle:
    """Export, quantize and check one checkpoint in a single pass.

    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.
    """
    from rukh.tokenize.uci_vocab import UciTokenizer
    from rukh.train import load_model

    model, _payload = load_model(Path(ckpt))
    bundle = ExportBundle(onnx=export_onnx(model, out, opset=opset, seq_len=seq_len))
    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)
    if check_parity:
        prefixes, source, warning = parity_positions(
            UciTokenizer(), n=positions, block=model.cfg.block, seed=seed, games=games
        )
        bundle.parity_source = source
        bundle.parity_warning = warning
        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
        bundle.parity = {
            kind: parity(model, Path(path), prefixes, n=positions) for kind, path in checks.items()
        }
    return bundle

__all__ = [
    "DEFAULT_GAMES",
    "DEFAULT_OPSET",
    "FP16_NAME",
    "INPUT_NAME",
    "INT8_NAME",
    "MODEL_NAME",
    "OUTPUT_NAME",
    "QUANTIZED_OPS",
    "ExportBundle",
    "ExportResult",
    "LastStepLogits",
    "ParityResult",
    "QuantizeResult",
    "export_all",
    "export_onnx",
    "parity",
    "parity_positions",
    "quantize_int8",
    "random_prefixes",
    "read_metadata",
    "sequence_lengths",
    "set_metadata",
    "target_path",
    "to_fp16",
    "validation_prefixes",
    "verify_dynamic_seq",
    "write_metadata",
]

@app.command("export")
def export_cmd(
    ckpt: Annotated[
        Path,
        typer.Option(
            "--ckpt", exists=True, dir_okay=False, readable=True, help="Checkpoint to export."
        ),
    ],
    out: Annotated[Path, typer.Option("--out", help="Output directory (or .onnx file).")],
    opset: Annotated[int, typer.Option("--opset", help="ONNX opset version.")] = 18,
    seq_len: Annotated[
        int, typer.Option("--seq-len", help="Example length the exporter traces.")
    ] = 200,
    fp16: Annotated[bool, typer.Option("--fp16", help="Also write the fp16 model.")] = False,
    int8: Annotated[bool, typer.Option("--int8", help="Also write the int8 model.")] = False,
    check_parity: Annotated[
        bool, typer.Option("--check-parity", help="Compare every file with PyTorch.")
    ] = False,
    positions: Annotated[
        int, typer.Option("--positions", help="Positions used by the parity check.")
    ] = 1000,
    games: Annotated[
        Path | None,
        typer.Option(
            "--games",
            exists=True,
            dir_okay=False,
            readable=True,
            help="Validation games parquet for the parity positions.",
        ),
    ] = None,
    as_json: Annotated[bool, typer.Option("--json", help="Print the result as JSON only.")] = False,
) -> None:
    """Export the next-move head to ONNX, quantize it and check parity with PyTorch."""
    from rukh.export import export_all

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    try:
        bundle = export_all(
            ckpt,
            out,
            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})")
    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 kind, result in bundle.parity.items():
        typer.echo(
            f"parity {kind}: {result.agreement:.4f} on {result.positions} "
            f"{bundle.parity_source} positions "
            f"(max |delta logits| {result.max_abs_logit_delta:.4g})"
        )
    if bundle.parity_warning:
        typer.echo(f"warning:  {bundle.parity_warning}")