El Hub y las cards: el catálogo, la vuelta y la afirmación pública

MODELS: tuple[Artefact, ...] = (
    Artefact(
        "tiny",
        "rukh-tiny",
        "decoder",
        "checkpoints/tiny/best.pt",
        ("m2", "m6", "m6"),
        "the 5 M decoder M2 iterates on (`configs/train/tiny.yaml`)",
        stage="tiny-greedy",
        measure="decoder",
        lesson="/curso/m2/02-entrenar-y-medir/",
    ),
    Artefact(
        "small",
        "rukh-small",
        "decoder",
        "checkpoints/small-v3/best.pt",
        ("m2", "m6", "m6"),
        "the 39 M course decoder, `small-v3` on the corpus with 24 months of Elite games",
        stage="small-v3-greedy",
        measure="decoder",
        lesson="/curso/m2/02-entrenar-y-medir/",
    ),
    Artefact(
        "medium-v4",
        "rukh-medium",
        "decoder",
        "checkpoints/medium-v4/best.pt",
        ("m4", "m5", "m6"),
        "the 115 M decoder every M4 fine-tune and every M5 alignment run starts from",
        stage="medium-v4-greedy",
        measure="decoder",
        lesson="/curso/m2/04-mas-datos-no-mas-red/",
    ),
    Artefact(
        "encoder-mmm-v4",
        "rukh-encoder-mmm",
        "encoder",
        "checkpoints/encoder-mmm-v4/best.pt",
        ("m3", "m6", "m6"),
        "the masked-move pretraining the published encoder's heads were fine-tuned on",
        lesson="/curso/m3/03-labs-del-encoder/",
    ),
    Artefact(
        "encoder-v4",
        "rukh-encoder",
        "encoder",
        "checkpoints/encoder-heads-v4/step-4000.pt",
        ("m3", "m6", "m6"),
        "the encoder with its three heads at its last step, as published and served by the demo",
        stage="encoder-v4",
        measure="encoder",
        lesson="/curso/m3/03-labs-del-encoder/",
    ),

    Artefact(
        "medium-elo",
        "rukh-medium-elo",
        "decoder",
        "checkpoints/medium-elo/step-3800.pt",
        ("m4", "m6", "m6"),
        "`medium-v4` fine-tuned on the Elo-balanced corpus (the last step, not `best.pt`)",
        stage="medium-elo",
        measure="decoder",
        lesson="/curso/m4/03-labs-de-afinado/",
    ),
    Artefact(
        "medium-masters",
        "rukh-medium-masters",
        "decoder",
        "checkpoints/medium-masters/step-3800.pt",
        ("m4", "m6", "m6"),
        "`medium-v4` fine-tuned on Elite games (the last step, not `best.pt`)",
        stage="medium-masters",
        measure="decoder",
        lesson="/curso/m4/03-labs-de-afinado/",
    ),
    Artefact(
        "lora-e4",
        "rukh-lora-e4",
        "adapter",
        "checkpoints/lora-e4",
        ("m4", "m6", "m6"),
        "the 1.6 MB style adapter that opens 1. e4",
        stage="lora-e4",
        measure="decoder",
        base="checkpoints/medium-v4/best.pt",
        lesson="/curso/m4/03-labs-de-afinado/",
    ),
    Artefact(
        "lora-d4",
        "rukh-lora-d4",
        "adapter",
        "checkpoints/lora-d4",
        ("m4", "m6", "m6"),
        "the 1.6 MB style adapter that opens 1. d4",
        stage="lora-d4",
        measure="decoder",
        base="checkpoints/medium-v4/best.pt",
        lesson="/curso/m4/03-labs-de-afinado/",
    ),
    Artefact(
        "qwen3-pgn-qlora",
        "rukh-qwen3-pgn-qlora",
        "peft",
        "checkpoints/qwen3-pgn-qlora",
        ("m4", "m6", "m6"),
        "the `peft` adapter of Qwen3-0.6B over PGN text, as `trl` wrote it",
        stage="qwen3-pgn-qlora",
        measure="qwen",
        lesson="/curso/m4/03-labs-de-afinado/",
    ),
    Artefact(
        "rm",
        "rukh-rm",
        "reward",
        "checkpoints/rm/reward.pt",
        ("m5", "m6", "m6"),
        "the Bradley-Terry reward model over 69 square tokens",
        lesson="/curso/m5/03-labs-de-alineamiento/",
    ),
    Artefact(
        "medium-v4-dpo-onpolicy",
        "rukh-medium-dpo",
        "decoder",
        "checkpoints/medium-v4-dpo-onpolicy/dpo.pt",
        ("m5", "m6", "m6"),
        "`medium-v4` after DPO on its own pairs, the aligned decoder the demo serves",
        stage="medium-v4-dpo-onpolicy-greedy",
        measure="decoder",
        lesson="/curso/m5/03-labs-de-alineamiento/",
    ),
    Artefact(
        "medium-v4-grpo",
        "rukh-medium-grpo",
        "decoder",
        "checkpoints/medium-v4-grpo/grpo.pt",
        ("m5", "m6", "m6"),
        "`medium-v4` after GRPO against the verifiable reward",
        stage="medium-v4-grpo-greedy",
        measure="decoder",
        lesson="/curso/m5/03-labs-de-alineamiento/",
    ),
)

DATASET_MODULES: dict[str, tuple[str, ...]] = {
    "rukh-games-1800": ("m1", "m2", "m3", "m4", "m5", "m6"),
    "rukh-tokenizer": ("m2", "m3", "m4", "m5", "m6"),
    "rukh-games-elite": ("m2", "m4"),
    "rukh-positions-eval": ("m2", "m3", "m5", "m6"),
    "rukh-puzzles-split": ("m2", "m3", "m4", "m5", "m6"),
    "rukh-elo-bins": (),
    "rukh-pairs-dpo": ("m5",),
    "rukh-pairs-onpolicy": ("m5",),
}
"""Which modules start from each dataset. ``rukh-elo-bins`` is M1's 100-Elo sample of the
1800+ corpus and nothing downstream reads it: M4 builds its own, from the club games it fetches.
"""


DATASET_LESSONS: dict[str, str] = {
    "rukh-games-1800": "/curso/m1/02-labs-del-pipeline/",
    "rukh-tokenizer": "/curso/m1/02-labs-del-pipeline/",
    "rukh-games-elite": "/curso/m2/04-mas-datos-no-mas-red/",
    "rukh-positions-eval": "/curso/m1/02-labs-del-pipeline/",
    "rukh-puzzles-split": "/curso/m1/02-labs-del-pipeline/",
    "rukh-elo-bins": "/curso/m1/02-labs-del-pipeline/",
    "rukh-pairs-dpo": "/curso/m1/02-labs-del-pipeline/",
    "rukh-pairs-onpolicy": "/curso/m5/03-labs-de-alineamiento/",
}
"""The lesson each dataset is built in."""


def _dataset_artefact(spec: DatasetSpec) -> Artefact:
    return Artefact(
        spec.name,
        spec.name,
        "dataset",
        spec.local_dir,
        DATASET_MODULES.get(spec.name, ()),
        spec.description.split(". ")[0].rstrip("."),
        lesson=DATASET_LESSONS.get(spec.name),
    )

def catalogue() -> list[Artefact]:
    """Every artefact ``rukh pull`` knows, models first and then the datasets."""
    return [*MODELS, *(_dataset_artefact(spec) for spec in DATASETS.values())]


def lookup(name: str) -> Artefact:
    """The artefact called ``name``: a run name, a dataset name or a Hub id."""
    for artefact in catalogue():
        if name in (artefact.name, artefact.repo, artefact.repo_id):
            return artefact
    known = ", ".join(a.name for a in catalogue())
    raise KeyError(f"nothing to pull called {name!r}; known: {known}")


def for_module(module: str) -> list[Artefact]:
    """What a reader who skipped everything before ``module`` needs on disk to start it."""
    return [artefact for artefact in catalogue() if module in artefact.modules]

def _download(repo_id: str, filename: str) -> Path:
    import huggingface_hub

    return Path(huggingface_hub.hf_hub_download(repo_id, filename))


def _load_state(repo_id: str) -> dict[str, Any]:
    """The weights of a Hub model, whichever of the two names they were published under."""
    import torch

    for name in ("model.safetensors", "pytorch_model.bin"):
        try:
            weights = _download(repo_id, name)
        except Exception as exc:  # noqa: BLE001 - a missing file only means "try the next one"
            log.debug("%s has no %s (%s)", repo_id, name, exc)
            continue
        if name.endswith(".safetensors"):
            from safetensors.torch import load_file

            return dict(load_file(str(weights)))
        return dict(torch.load(weights, map_location="cpu", weights_only=True))
    raise FileNotFoundError(f"{repo_id} holds neither model.safetensors nor pytorch_model.bin")

def _model_payload(artefact: Artefact, config: dict[str, Any], state: dict[str, Any]) -> dict:
    """Fold a Hub layout back into the checkpoint payload every loader of this project reads.

    ``rukh publish`` writes the model's shape and the run's provenance into ``config.json``;
    this puts them back under the keys ``save_checkpoint`` uses. The heads of an encoder need
    the pooling they were trained with, which the card's config carries and ``build_heads``
    reads from the run config.
    """
    from rukh.models import DecoderConfig, EncoderConfig

    if artefact.kind == "reward":
        return {
            "model": state,
            "encoder_config": config["encoder"],
            "pooling": config.get("pooling"),
            "config": {"pulled_from": artefact.repo_id},
        }
    shape = EncoderConfig if artefact.kind == "encoder" else DecoderConfig
    model_cfg = {key: value for key, value in config.items() if key in shape.model_fields}
    run_cfg: dict[str, Any] = {"pulled_from": artefact.repo_id}
    if artefact.kind == "encoder" and config.get("pooling"):
        run_cfg["pooling"] = config["pooling"]
    return {
        "step": int(config.get("step") or 0),
        "model_state": state,
        "opt_state": None,
        "cfg": run_cfg,
        "model_cfg": model_cfg,
        "vocab_hash": config.get("vocab_hash"),
        "data_manifest_sha": config.get("data_manifest_sha"),
        "git_sha": config.get("git_sha"),
        "best_val": None,
    }

def pull(artefact: Artefact, force: bool = False) -> Path:
    """Bring ``artefact`` to its local path and return that path.

    An artefact already on disk is left alone unless ``force``: a reader's own training run,
    sitting at the same stable path, must never be replaced silently by the published one.
    """
    import torch

    target = paths.resolve(artefact.target)
    if target.exists() and not force:
        log.info("%s is already there; --force replaces it", target)
        return target
    if artefact.kind == "dataset":
        import huggingface_hub

        spec = DATASETS[artefact.name]
        target.mkdir(parents=True, exist_ok=True)
        huggingface_hub.snapshot_download(
            artefact.repo_id,
            repo_type=spec.repo_type,
            local_dir=str(target),
            allow_patterns=[*spec.patterns, "manifest.json"],
        )
        return target
    if artefact.kind == "peft":
        import huggingface_hub

        target.mkdir(parents=True, exist_ok=True)
        huggingface_hub.snapshot_download(artefact.repo_id, local_dir=str(target))
        return target
    if artefact.kind == "adapter":
        target.mkdir(parents=True, exist_ok=True)
        for name in ("adapter.safetensors", "adapter_config.json"):
            source = _download(artefact.repo_id, name)
            (target / name).write_bytes(source.read_bytes())
        return target
    config = json.loads(_download(artefact.repo_id, "config.json").read_text(encoding="utf-8"))
    state = _load_state(artefact.repo_id)
    target.parent.mkdir(parents=True, exist_ok=True)
    tmp = target.with_suffix(target.suffix + ".tmp")
    torch.save(_model_payload(artefact, config, state), tmp)
    tmp.replace(target)
    return target

            patterns=["pairs.parquet", "dpo-prompts.parquet", "manifest.json"],

@app.command("pull")
def pull_cmd(
    names: Annotated[
        list[str] | None,
        typer.Argument(help="Run names, dataset names or Hub ids (see --list)."),
    ] = None,
    module: Annotated[
        str | None,
        typer.Option("--module", help="Pull everything a module starts from, e.g. m4."),
    ] = None,
    list_all: Annotated[
        bool, typer.Option("--list", help="Print the catalogue, by module, and exit.")
    ] = False,
    force: Annotated[
        bool, typer.Option("--force", help="Replace what is already on disk.")
    ] = False,
) -> None:
    """Bring published models and datasets to the paths the configs and the lessons read.

    A pulled model lands under its run name (`checkpoints/medium-v4/best.pt`), which is the
    spelling every config uses and the one a run you trained yourself also answers to.
    """
    from rukh.hub import catalogue, for_module, lookup, pull

    if list_all:
        for artefact in catalogue():
            modules = ",".join(artefact.modules) or "-"
            typer.echo(f"{artefact.name:<24} {modules:<15} {artefact.target}")
            typer.echo(f"{'':<24} {artefact.repo_id:<15} {artefact.note}")
        return
    wanted = []
    if module is not None:
        wanted.extend(for_module(module))
        if not wanted:
            typer.echo(f"error: nothing is catalogued for module {module!r}", err=True)
            raise typer.Exit(code=2)
    try:
        wanted.extend(lookup(name) for name in names or [])
    except KeyError as exc:
        typer.echo(f"error: {exc.args[0]}", err=True)
        raise typer.Exit(code=2) from exc
    if not wanted:
        typer.echo("error: give a name, --module or --list (see rukh pull --help)", err=True)
        raise typer.Exit(code=2)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    for chatty in ("httpx", "huggingface_hub", "urllib3"):
        logging.getLogger(chatty).setLevel(logging.WARNING)  # one line per artefact, not per GET
    for artefact in wanted:
        try:
            path = pull(artefact, force=force)
        except (FileNotFoundError, OSError) as exc:
            typer.echo(f"error: {artefact.name}: {exc}", err=True)
            raise typer.Exit(code=1) from exc
        typer.echo(f"{artefact.name:<24} {artefact.repo_id:<28} -> {path}")

"""The links every card carries back to the course: its lesson, its row in the table, the demo.

A model card that does not say where its number came from is a number without a method. The
single results table (`/proyecto/` on the course site) is where every stage is measured the same
way, and the lesson is where it was built; both are looked up in the catalogue of ``rukh.hub`` so
a repo id is enough and nothing is typed twice.
"""

from __future__ import annotations

from typing import Any

from rukh.hub import catalogue

TABLE_PATH = "/proyecto/"

LESSON_TITLES: dict[str, str] = {
    "/curso/m1/02-labs-del-pipeline/": "M1 · Del PGN al tensor: los labs",
    "/curso/m2/02-entrenar-y-medir/": "M2 · El decoder: entrenarlo y medirlo",
    "/curso/m2/04-mas-datos-no-mas-red/": "M2 · Más datos, no más red: de small a medium-v4",
    "/curso/m3/03-labs-del-encoder/": "M3 · El encoder: los labs",
    "/curso/m4/03-labs-de-afinado/": "M4 · Fine-tuning: los labs",
    "/curso/m5/03-labs-de-alineamiento/": "M5 · Alineamiento: los labs",
}


def course_links(
    repo_id: str,
    course_url: str = "https://lab.rukh.borjaglez.com",
    demo_url: str = "https://rukh.borjaglez.com",
) -> dict[str, Any]:
    """What the card's "In the course" section needs for ``repo_id``.

    ``lesson_url`` and ``stage`` are ``None`` for a repository the catalogue does not know, and
    the section then names only the table: a card never claims a lesson it was not built in.
    """
    found = next((a for a in catalogue() if a.repo_id == repo_id), None)
    lesson = found.lesson if found else None
    from rukh.publish.model import DEMO_STAGES

    name = repo_id.split("/")[-1].removeprefix("rukh-")
    demo_stage = DEMO_STAGES.get(name)
    return {
        "table_url": f"{course_url}{TABLE_PATH}",
        "lesson_url": f"{course_url}{lesson}" if lesson else None,
        "lesson_title": LESSON_TITLES.get(lesson or "", lesson),
        "stage": found.stage if found else None,
        "demo_url": f"{demo_url}/?stage={demo_stage}" if demo_stage else None,
        "reproduce": f"uv run rukh pull {found.name}" if found else None,
    }

    # Lazy: `rukh.publish.links` reads the catalogue of `rukh.hub`, which imports this module.
    from rukh.publish.links import course_links

## In the course

{% if course.lesson_url -%}
- Built in [{{ course.lesson_title }}]({{ course.lesson_url }}), the lesson that runs every command behind this repository.
{% endif -%}
- Measured in the [single results table]({{ course.table_url }}){% if course.stage %} as the stage `{{ course.stage }}`{% endif %}: every stage of the course, the same suite, the same day.
{% if course.demo_url -%}
- Play it in the browser: [{{ course.demo_url }}]({{ course.demo_url }}).
{% endif -%}
{% if course.reproduce -%}
- Bring it to the paths the configs read: `{{ course.reproduce }}`.
{% endif %}

{% include '_course.md.jinja' %}

---
license: {{ license }}
library_name: rukh
pipeline_tag: feature-extraction
language:
  - en
datasets:
{%- for dataset in datasets %}
  - {{ dataset }}
{%- endfor %}
tags:
  - chess
  - rukh
  - encoder
  - masked-move-modeling
---

# {{ repo_id }}

A bidirectional transformer written from scratch and **pretrained only**: it has no value,
blunder or result head yet. It read games written in UCI with 15 % of the moves hidden and
learned to put them back by looking at both sides of the hole, which is masked language
modeling with moves for words. This is the `{{ stage }}` stage of [Rukh]({{ repository_url }}),
a course that builds a chess language model end to end: {{ "{:,}".format(params) }} parameters
over {{ config.get("n_layer", 8) }} layers of width {{ config.get("d_model", 384) }}, on the
`{{ config.get("input", "moves") }}` input scheme.

It exists on the Hub for one reason: it is where the encoder with heads
([`{{ heads_repo }}`](https://huggingface.co/{{ heads_repo }})) starts from, and a reader of the
course who skips the pretraining lab can put it where the fine-tuning configs look for it:

```sh
uv run rukh pull {{ stage }}      # -> checkpoints/{{ stage }}/best.pt
uv run rukh train heads --config configs/train/encoder-heads-v4.yaml --mode last-n
```

Read how it was built: [{{ course_url }}]({{ course_url }})

## Results

The only thing a pretraining run measures is how well it fills the holes it was given.

| Metric | Value |
|---|---|
| Masked-move loss on held-out games | {{ best_val }} |
| Masked-move top-1 on held-out games | {{ val_top1 }} |
| Training step of this checkpoint | {{ step }} |

These are not the numbers of the model that goes into the demo: the heads are fine-tuned on
Stockfish labels afterwards, and their F1, correlation and result accuracy live on the card of
[`{{ heads_repo }}`](https://huggingface.co/{{ heads_repo }}).

## Training recipe

{% if recipe %}
| Parameter | Value |
|---|---|
{%- for name, value in recipe %}
| `{{ name }}` | `{{ value }}` |
{%- endfor %}
{% else %}
The MLflow run for this checkpoint was not available when the card was generated; the shape of
the model is in `config.json`.
{% endif %}
{% if run_id %}MLflow run: `{{ run_id }}`.{% endif %}

```json
{{ config_json }}
```

## Data

Trained on {% for dataset in datasets %}[`{{ dataset }}`](https://huggingface.co/datasets/{{ dataset }}){% if not loop.last %}, {% endif %}{% endfor %},
derived from the [Lichess open database](https://database.lichess.org) (CC0), on the same
packed token stream the project's decoder trains on, so the two models saw the same games.
Validation uses a held-out slice of a month the training split does not read.

## Limitations

- It predicts hidden moves; it does not evaluate positions. Ask it how good a position is and
  the answer is whatever an untrained linear head returns.
- Its masked-move accuracy is not a measure of chess strength: a move is easy to fill in when
  the moves around it leave one legal option.
- Fine-tuning the whole network on the labels afterwards made it *worse* than fine-tuning the
  last two blocks; the course records that measurement, and the heads configs default to it.

{% include '_course.md.jinja' %}

## License

Weights: Apache-2.0. Data: CC0 (Lichess). Code: MIT, at [{{ repository_url }}]({{ repository_url }}).

ONNX_DIRS: dict[str, str] = {
    "tiny": "artifacts/onnx/tiny",
    "small": "artifacts/onnx/small-v3",
    "medium-v4": "artifacts/onnx/medium-v4",
    "encoder-v4": "artifacts/onnx/encoder-v4",
    "medium-elo": "artifacts/onnx/medium-elo",
    "medium-masters": "artifacts/onnx/medium-masters",
    "medium-v4-dpo-onpolicy": "artifacts/onnx/medium-dpo",
    "medium-v4-grpo": "artifacts/onnx/medium-grpo",
}
"""The ONNX export each model's card reports parity from, when it is on disk."""

EFFECTS_DIR = "artifacts/publish/effects"


class CardRefresh(BaseModel):
    """One repository's card: where it was staged and whether it went up."""

    model_config = ConfigDict(extra="forbid")

    name: str
    repo_id: str
    repo_type: str
    card: str | None = None
    status: str
    """``staged`` (dry run), ``uploaded`` or ``failed``."""
    detail: str | None = None

def stage_card(artefact: Artefact, cfg: ModelPublishConfig) -> Path:
    """Render ``artefact``'s card through its own publisher, in dry-run mode; returns the card."""
    target = resolve_run(paths.resolve(artefact.target))
    if artefact.kind == "dataset":
        from rukh.data.publish import PublishConfig, publish

        result = publish(artefact.name, PublishConfig(owner=cfg.owner), dry_run=True)
        return Path(result.card_path)
    if artefact.kind == "adapter":
        from rukh.publish.adapter import AdapterEffect, publish_adapter

        effect_path = paths.resolve(EFFECTS_DIR) / f"{artefact.name}.json"
        effect = (
            AdapterEffect.model_validate_json(effect_path.read_text(encoding="utf-8"))
            if effect_path.is_file()
            else None
        )
        result = publish_adapter(
            target, artefact.repo_id, f"{cfg.owner}/rukh-medium", cfg, effect=effect, dry_run=True
        )
        return Path(result.card_path)
    if artefact.kind == "peft":
        from rukh.publish.adapter import publish_qwen_adapter

        result = publish_qwen_adapter(target, artefact.repo_id, cfg, dry_run=True)
        return Path(result.card_path)
    if artefact.kind == "reward":
        from rukh.publish.reward import publish_reward

        result = publish_reward(target, artefact.repo_id, cfg, dry_run=True)
        return Path(result.folder) / README_NAME
    from rukh.publish.model import publish_model

    onnx = ONNX_DIRS.get(artefact.name)
    onnx_dir = paths.resolve(onnx) if onnx and paths.resolve(onnx).is_dir() else None
    result = publish_model(
        target, artefact.repo_id, cfg, onnx_dir=onnx_dir, stage=artefact.stage, dry_run=True
    )
    return Path(result.card_path)

def refresh_cards(
    only: set[str] | None = None,
    dry_run: bool = False,
    cfg: ModelPublishConfig | None = None,
    api: Any | None = None,
) -> list[CardRefresh]:
    """Regenerate the card of every repository (or of ``only``); upload it unless ``dry_run``."""
    cfg = cfg or ModelPublishConfig()
    wanted = [a for a in catalogue() if not only or a.name in only or a.repo_id in only]
    client = None
    outcomes: list[CardRefresh] = []
    for artefact in wanted:
        repo_type = (
            "dataset"
            if artefact.kind == "dataset" and artefact.name != "rukh-tokenizer"
            else "model"
        )
        outcome = CardRefresh(
            name=artefact.name, repo_id=artefact.repo_id, repo_type=repo_type, status="staged"
        )
        outcomes.append(outcome)
        try:
            card = stage_card(artefact, cfg)
            outcome.card = card.as_posix()
            if not dry_run:
                client = client or api or _api()
                client.upload_file(
                    path_or_fileobj=str(card),
                    path_in_repo=README_NAME,
                    repo_id=artefact.repo_id,
                    repo_type=repo_type,
                    commit_message="docs(card): regenerated from the current table and course",
                )
                outcome.status = "uploaded"
        except Exception as exc:  # noqa: BLE001 - one card failing must not stop the rest
            log.exception("card of %s failed", artefact.name)
            outcome.status = "failed"
            outcome.detail = f"{type(exc).__name__}: {exc}"
    return outcomes

"""``rukh publish collection``: the ``Rukh`` collection on the Hub, kept equal to the catalogue.

A collection is the one page that lists everything the course trained, in the order the course
built it. It is derived from ``rukh.hub`` so it cannot drift from what ``rukh pull`` can bring
back: adding a model to the catalogue and running this is what "publishing the collection" is.
Idempotent on purpose -- the Hub refuses duplicates when asked (``exists_ok``), and a second run
changes nothing.
"""

from __future__ import annotations

import logging
from typing import Any

from pydantic import BaseModel, ConfigDict

from rukh.hub import catalogue
from rukh.publish.model import ModelPublishConfig, _api

log = logging.getLogger(__name__)

TITLE = "Rukh"
DESCRIPTION = (
    "A chess language model built from scratch as a course on generative and agentic AI: every "
    "model, adapter and dataset, in course order."
)
"""The Hub caps a collection's description at 150 characters."""

def collection_items() -> list[tuple[str, str]]:
    """``(repo id, item type)`` for every catalogued artefact, models first, in course order."""
    items: list[tuple[str, str]] = []
    for artefact in catalogue():
        item_type = (
            "dataset"
            if artefact.kind == "dataset" and artefact.name != "rukh-tokenizer"
            else "model"
        )
        if (artefact.repo_id, item_type) not in items:
            items.append((artefact.repo_id, item_type))
    return items


def sync_collection(
    cfg: ModelPublishConfig | None = None, dry_run: bool = False, api: Any | None = None
) -> CollectionSync:
    """Create the collection if it does not exist and add every catalogued repository to it."""
    cfg = cfg or ModelPublishConfig()
    items = collection_items()
    if dry_run:
        return CollectionSync(slug=None, created=False, items=[r for r, _ in items], dry_run=True)
    client = api or _api()
    existing = next((c for c in client.list_collections(owner=cfg.owner) if c.title == TITLE), None)
    created = existing is None
    collection = existing or client.create_collection(
        TITLE, namespace=cfg.owner, description=DESCRIPTION, private=False, exists_ok=True
    )
    for repo_id, item_type in items:
        client.add_collection_item(collection.slug, repo_id, item_type, exists_ok=True)
        log.info("collection %s: %s (%s)", collection.slug, repo_id, item_type)
    return CollectionSync(
        slug=collection.slug, created=created, items=[r for r, _ in items], dry_run=False
    )

@publish_app.command("cards")
def publish_cards_cmd(
    only: Annotated[
        str | None, typer.Option("--only", help="Comma-separated names or repo ids.")
    ] = None,
    dry_run: Annotated[
        bool, typer.Option("--dry-run", help="Stage every card locally; upload nothing.")
    ] = False,
) -> None:
    """Regenerate every card from the current table and course, and upload only the card.

    Each repository is staged through its own publisher (nothing leaves the machine), then
    `README.md` alone is uploaded: changing a paragraph must not republish 460 MB of weights.
    Needs the checkpoints on disk (`rukh pull`) and the evaluations (`rukh eval nightly`).
    """
    from rukh.publish.cards import refresh_cards

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    wanted = {part.strip() for part in (only or "").split(",") if part.strip()} or None
    outcomes = refresh_cards(only=wanted, dry_run=dry_run)
    for outcome in outcomes:
        typer.echo(f"{outcome.repo_id:<34} {outcome.status:<9} {outcome.card or outcome.detail}")
    if any(outcome.status == "failed" for outcome in outcomes):
        raise typer.Exit(code=1)


@publish_app.command("collection")
def publish_collection_cmd(
    dry_run: Annotated[
        bool, typer.Option("--dry-run", help="List the items; touch no network.")
    ] = False,
) -> None:
    """Create or update the `Rukh` collection on the Hub with every catalogued repository."""
    from rukh.publish.collection import sync_collection

    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    result = sync_collection(dry_run=dry_run)
    typer.echo(f"collection: {result.slug or '(dry run)'}{' · created' if result.created else ''}")
    for item in result.items:
        typer.echo(f"  {item}")

    The evaluation stores ``sha256`` of the checkpoint file and, since P6, of its tensors alone
    (``weights_sha``). A different file with the same weights -- the copy ``rukh pull`` brings
    back, which ``rukh publish`` wrote without the optimizer -- passes on the second; a
    ``step-*.pt`` whose weights differ from ``best.pt`` is still refused (D-135).
    """
    if not evaluation:
        return
    measured = str(evaluation.get("model_sha") or "")
    if not measured:
        return
    from rukh.eval.cache import file_sha, weights_sha

    actual = file_sha(Path(ckpt))
    if measured == actual:
        return
    measured_weights = str(evaluation.get("weights_sha") or "")
    if measured_weights and measured_weights == weights_sha(Path(ckpt)):
        return
    raise ValueError(
        f"stage {stage!r} was measured on a different checkpoint: its results.json records "
        f"model_sha {measured[:12]} and {Path(ckpt).as_posix()} hashes to {actual[:12]}"
        + (f" (weights {measured_weights[:12]} vs the file's own)" if measured_weights else "")
        + ". Evaluate this checkpoint under this stage, or publish the one that was evaluated."
    )

def pretrained_encoder_card_context(
    repo_id: str,
    stage: str,
    cfg: ModelPublishConfig,
    config: dict[str, Any],
    payload: dict[str, Any],
    run: RunSummary | None,
    files: list[str],
) -> dict[str, Any]:
    """The card of a masked-move pretraining run: what it measured, and what starts from it.

    A pretraining checkpoint has no blunder F1 and no value correlation to report; rendering it
    through the heads' card would say "trained from scratch on Stockfish labels" about a model
    that never saw a label. What it has is the held-out masked-move loss the run selected
    ``best.pt`` on, which the checkpoint carries, and the top-1 of the same evaluation when the
    MLflow run is reachable.
    """
    metrics = run.metrics if run else {}
    top1 = next((metrics[key] for key in ("val_top1", "val/top1") if key in metrics), None)
    best_val = payload.get("best_val")

    from rukh.models import PositionEncoder

    if isinstance(model, PositionEncoder):
        # Pretraining only: no heads on it, so none of the heads' numbers apply.
        card = render_card(
            pretrained_encoder_card_context(repo_id, name, cfg, config, payload, run, files),
            PRETRAINED_ENCODER_CARD_TEMPLATE,
        )
    elif kind == "encoder":
        card = render_card(
            encoder_card_context(
                repo_id, name, cfg, config, evaluation, run, files, measured_parity
            ),
            ENCODER_CARD_TEMPLATE,
        )

def test_every_card_template_carries_the_course_section() -> None:
    templates = [p for p in CARDS_DIR.glob("*.md.jinja") if not p.name.startswith("_")]
    assert len(templates) >= 7
    for template in templates:
        text = template.read_text(encoding="utf-8")
        assert "{% include '_course.md.jinja' %}" in text, template.name
        # Before the licence: the links are part of the card, not an afterthought under it.
        assert text.index("_course.md.jinja") < text.index("## Licen"), template.name

def test_the_links_name_the_lesson_the_table_and_the_demo_when_they_apply() -> None:
    medium = course_links("chorcat/rukh-medium")
    assert (
        medium["lesson_url"] == "https://lab.rukh.borjaglez.com/curso/m2/04-mas-datos-no-mas-red/"
    )
    assert medium["lesson_title"].startswith("M2")
    assert medium["table_url"] == "https://lab.rukh.borjaglez.com/proyecto/"
    assert medium["stage"] == "medium-v4-greedy"
    assert medium["reproduce"] == "uv run rukh pull medium-v4"
    pairs = course_links("chorcat/rukh-pairs-dpo")
    assert pairs["lesson_url"].endswith("/curso/m1/02-labs-del-pipeline/")
    assert pairs["stage"] is None and pairs["demo_url"] is None
    # A repository the catalogue does not know gets the table and nothing it cannot back.
    unknown = course_links("chorcat/rukh-nope", course_url="https://x", demo_url="https://d")
    assert unknown == {
        "table_url": "https://x/proyecto/",
        "lesson_url": None,
        "lesson_title": None,
        "stage": None,
        "demo_url": None,
        "reproduce": None,
    }

def test_model_targets_are_run_names_the_configs_use() -> None:
    """A pulled model must sit exactly where the shipped configs look for it."""
    from rukh.paths import package_root

    configs = (path for path in (package_root() / "configs").rglob("*.yaml") if path.name[0] != ".")
    referenced = set()
    for config in configs:
        for line in config.read_text(encoding="utf-8").splitlines():
            if "checkpoints/" in line and not line.strip().startswith("#"):
                referenced.add(line.split("checkpoints/", 1)[1].split()[0].strip("'\""))
    targets = {a.target.removeprefix("checkpoints/") for a in MODELS}
    # Runs the reader always trains themselves, because they take minutes and the lesson is in
    # running them: the 15 M masked-move encoder of M3 and the off-policy DPO arm of M5.
    local_only = {"encoder-mmm/best.pt", "medium-v4-dpo-offpolicy/dpo.pt"}
    # Every other checkpoint a shipped config starts from is something `rukh pull` can provide.
    missing = referenced - targets - local_only
    assert not missing, missing