El baseline público: el nanoGPT de Karvonen en nuestra escalera

@dataclass(frozen=True)
class GPTConfig:
    """nanoGPT's ``model_args`` as the checkpoint carries them; ``dropout`` is unused here."""

    n_layer: int = 8
    n_head: int = 8
    n_embd: int = 512
    block_size: int = 1023
    bias: bool = False
    vocab_size: int = 32
    dropout: float = 0.0

class _Attention(nn.Module):
    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        self.n_head, self.n_embd = cfg.n_head, cfg.n_embd
        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias)
        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias)

    def forward(self, x: Tensor) -> Tensor:
        batch, length, width = x.shape
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        shape = (batch, length, self.n_head, width // self.n_head)
        q, k, v = (t.view(shape).transpose(1, 2) for t in (q, k, v))
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.c_proj(y.transpose(1, 2).contiguous().view(batch, length, width))


class _MLP(nn.Module):
    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        self.c_fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=cfg.bias)
        self.c_proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=cfg.bias)

    def forward(self, x: Tensor) -> Tensor:
        return self.c_proj(F.gelu(self.c_fc(x)))


class _Block(nn.Module):
    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        self.ln_1 = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.attn = _Attention(cfg)
        self.ln_2 = nn.LayerNorm(cfg.n_embd, bias=cfg.bias)
        self.mlp = _MLP(cfg)

    def forward(self, x: Tensor) -> Tensor:
        x = x + self.attn(self.ln_1(x))
        return x + self.mlp(self.ln_2(x))

class GPT(nn.Module):
    """Karpathy's nanoGPT, exactly as far as the state dict needs: same names, same shapes."""

    def __init__(self, cfg: GPTConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.transformer = nn.ModuleDict(
            {
                "wte": nn.Embedding(cfg.vocab_size, cfg.n_embd),
                "wpe": nn.Embedding(cfg.block_size, cfg.n_embd),
                "h": nn.ModuleList([_Block(cfg) for _ in range(cfg.n_layer)]),
                "ln_f": nn.LayerNorm(cfg.n_embd, bias=cfg.bias),
            }
        )
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight  # tied, as nanoGPT ties them

    def forward(self, idx: Tensor) -> Tensor:
        """Logits over the vocabulary at every position of ``idx`` (``batch, length``)."""
        length = idx.shape[1]
        if length > self.cfg.block_size:
            raise ValueError(f"sequence of {length} exceeds block size {self.cfg.block_size}")
        pos = torch.arange(length, device=idx.device)
        x = self.transformer.wte(idx) + self.transformer.wpe(pos)
        for block in self.transformer.h:
            x = block(x)
        return self.lm_head(self.transformer.ln_f(x))

def strip_compiled(state: dict[str, Any]) -> dict[str, Any]:
    """The state dict without the ``_orig_mod.`` that ``torch.compile`` leaves on every key."""
    return {key.removeprefix(COMPILED_PREFIX): value for key, value in state.items()}


def load_karvonen(checkpoint: Path | str, device: str | None = None) -> GPT:
    """The model of a nanoGPT checkpoint (``model_args`` plus ``model``), in eval mode."""
    payload = torch.load(Path(checkpoint), map_location="cpu", weights_only=True)
    args = {k: v for k, v in payload["model_args"].items() if k in GPTConfig.__dataclass_fields__}
    model = GPT(GPTConfig(**args))
    model.load_state_dict(strip_compiled(payload["model"]))
    return model.to(device or "cpu").eval()

def ensure_karvonen(download: bool = True, local_dir: Path | str | None = None) -> Path:
    """The checkpoint on disk (downloading it and ``meta.pkl`` when allowed), and its path.

    The Hub repository is somebody else's, so this is not a ``rukh pull`` artefact -- nothing here
    is published by the course, nothing can be. ``huggingface_hub`` caches the 100 MB file once.
    """
    directory = Path(local_dir) if local_dir is not None else paths.resolve(LOCAL_DIR)
    checkpoint = directory / CHECKPOINT_FILE
    meta = directory / "meta.pkl"
    if not checkpoint.is_file():
        if not download:
            raise FileNotFoundError(f"{checkpoint} is not on disk (rukh eval karvonen pulls it)")
        import huggingface_hub

        log.info("downloading %s/%s", HF_REPO, CHECKPOINT_FILE)
        directory.mkdir(parents=True, exist_ok=True)
        huggingface_hub.hf_hub_download(HF_REPO, CHECKPOINT_FILE, local_dir=str(directory))
    if not meta.is_file() and download:
        import urllib.request

        log.info("downloading %s", META_URL)
        with urllib.request.urlopen(META_URL, timeout=60) as response:  # noqa: S310 - https, pinned
            meta.write_bytes(response.read())
    if meta.is_file():
        read_meta(meta)
    return checkpoint

class KarvonenPlayer:
    """A ``rukh.infer.Player`` that writes SAN one character at a time from the transcript.

    Same accounting as ``QwenPlayer`` (``QwenStats``, ``parse_san_answer``, first legal move as the
    fallback and never a second sample), so its failure breakdown is comparable with that row.
    """

    def __init__(
        self,
        model: GPT,
        stats: QwenStats | None = None,
        temperature: float = 0.0,
        top_k: int | None = None,
        seed: int | None = None,
        max_new_chars: int = MAX_NEW_CHARS,
        max_plies: int = DECODER_PLIES,
    ) -> None:
        self.model = model
        self.stats = stats if stats is not None else QwenStats()
        self.temperature, self.top_k, self.max_new_chars = temperature, top_k, max_new_chars
        self.max_plies = max_plies
        self.device = next(model.parameters()).device
        self.generator: torch.Generator | None = None
        if seed is not None:
            self.generator = torch.Generator(device="cpu").manual_seed(seed)

    @property
    def limit(self) -> int:
        """Plies this player can carry: the decoder's number, so every row plays the same game."""
        return self.max_plies

    def start(self, white_elo: int, black_elo: int) -> None:
        """Nothing to set: this model has no rating header, which is part of what it is."""

    def observe(self, move: chess.Move) -> None:
        """Nothing to record: the board carries the history and the transcript is built from it."""

class KarvonenResult(BaseModel):
    """One evaluation of the public baseline on the decoder's own suite."""

    model_config = ConfigDict(extra="forbid")

    stage: str
    source: str
    """Where the weights come from: the Hub repository and file of their author."""
    checkpoint: str
    model_sha: str
    params: int
    date: str
    device: str
    positions: int
    written: QwenStats
    """How its answers came out, by failure mode, over the validation positions only."""
    top1: float | None = None
    puzzles: PuzzleResult | None = None
    puzzles_blind: int = 0
    """Puzzle prompts that carried no game prefix and were answered from the FEN position alone,
    which this model cannot read. A rate over many of these is a floor, not a measurement."""
    elo: EloResult | None = None
    notes: list[str] = []
    config: dict[str, Any] = {}

def evaluate_karvonen(
    cfg: EvalConfig,
    stage: str = DEFAULT_STAGE,
    use_cache: bool = True,
    device: str | None = None,
    checkpoint: Path | str | None = None,
) -> KarvonenResult:
    """Run legality, next-move accuracy, puzzles and Elo on the published nanoGPT."""
    from rukh.engine import EngineNotFound
    from rukh.eval.cache import EvalCache, config_sha, file_sha
    from rukh.train import pick_device

    ckpt = Path(checkpoint) if checkpoint is not None else ensure_karvonen()
    where = device or pick_device()
    model = load_karvonen(ckpt, device=where)
    sha = file_sha(ckpt)
    tok = UciTokenizer()
    notes: list[str] = [
        "this model has no rating header: header_elo does not condition this row",
    ]
    positions: list[Position] = _positions(cfg, tok, notes)
    player = KarvonenPlayer(model, temperature=cfg.temperature, top_k=cfg.top_k, seed=cfg.seed)

    log.info("legality and accuracy over %s positions", min(len(positions), cfg.accuracy_positions))
    scored = positions[: cfg.accuracy_positions]
    hits = 0
    for position in scored:
        move, _ = player.propose(_board_of(position))
        if move is not None and position.target is not None and move.uci() == position.target:
            hits += 1
    written = player.stats.model_copy(deep=True)

    cache_path = paths.resolve(cfg.cache_db) if use_cache else None
    games_cache = EvalCache(
        cache_path, sha, enabled=use_cache, config_sha=config_sha(cfg.cache_fields("games"))
    )
    puzzles_cache = EvalCache(
        cache_path, sha, enabled=use_cache, config_sha=config_sha(cfg.cache_fields("puzzles"))
    )
    try:
        puzzle_path = paths.resolve(cfg.puzzles)
        if puzzle_path.is_file():
            items = load_puzzles(
                puzzle_path, cfg.puzzles_per_band, seed=cfg.seed, split=cfg.puzzle_split
            )
            log.info("puzzles: %s", len(items))
            source = karvonen_source(player, tok)
            result.puzzles = run_puzzles(
                source,
                tok,
                items,
                cache=puzzles_cache,
                header_elo=cfg.header_elo,
                force_header=cfg.force_header,
            )
            result.puzzles_blind = source.blind
            if source.blind:
                notes.append(
                    f"{source.blind} puzzle prompts had no game prefix and were answered from "
                    "the position alone, which this model cannot read: the rate is a floor"
                )
        else:
            notes.append(f"puzzles not found at {puzzle_path}: puzzle suite skipped")

        if cfg.elo_games:
            log.info("elo: %s games per rung", cfg.elo_games)
            try:
                records = play_rungs(
                    None,
                    tok,
                    cfg.elo_rungs,
                    cfg.elo_games,
                    cfg.sampling(),
                    move_time=cfg.elo_move_time,
                    max_plies=cfg.elo_max_plies,
                    cache=games_cache,
                    header_elo=cfg.header_elo,
                    player=player,
                    nodes=cfg.elo_nodes,
                )
            except EngineNotFound as exc:
                notes.append(f"Elo skipped: {exc}")
            else:
                result.elo = estimate(records, samples=cfg.bootstrap, seed=cfg.seed)
    finally:
        games_cache.close()
        puzzles_cache.close()

                max_plies=cfg.elo_max_plies,
                cache=cache,
                header_elo=cfg.header_elo,
                nodes=cfg.elo_nodes,
                player=player,
            )

    date: str
    run_id: str | None = None
    baseline: bool = False
    """A model somebody else trained, measured here as a control on the harness itself. Off for
    every model of the course, and for every row written before the flag existed."""

def row_of(result: KarvonenResult) -> WebRow:
    """The row this model contributes to the single results table, flagged as a baseline."""
    elo = result.elo
    interval = (
        [elo.ci_low, elo.ci_high]
        if elo and not elo.separated and elo.ci_low is not None and elo.ci_high is not None
        else None
    )
    return WebRow(
        stage=result.stage,
        params=result.params,
        legality=result.written.legal_rate or None,
        top1=result.top1,
        puzzles=result.puzzles.by_band() if result.puzzles else {},
        elo=elo.elo if elo else None,
        elo_ci=interval,
        elo_lower=elo.elo_lower if elo else None,
        elo_upper=elo.elo_upper if elo else None,
        elo_separated=bool(elo.separated) if elo else False,
        date=result.date,
        baseline=True,
    )

    if result.puzzles is not None:
        lines += [
            "## Puzzles by band",
            "",
            f"Prompts answered without a game prefix: {result.puzzles_blind}.",
            "",
            "| Band | Attempted | Solved | Rate |",
            "|---|---:|---:|---:|",
        ]
        lines += [
            f"| {band.band} | {band.attempted} | {band.solved} | {band.rate * 100:.1f} % |"
            for band in result.puzzles.bands
        ]
        lines.append("")
    if result.elo is not None:
        lines += ["## Games against Stockfish", "", "| Rung | Games | Score |", "|---|---:|---:|"]
        lines += [f"| {rung.name} | {rung.games} | {rung.score:.3f} |" for rung in result.elo.rungs]
        lines.append("")
    if result.notes:
        lines += ["## Notes", ""] + [f"- {note}" for note in result.notes] + [""]
    return "\n".join(lines)

def run_karvonen_suite(
    cfg: EvalConfig,
    stage: str = DEFAULT_STAGE,
    use_cache: bool = True,
    device: str | None = None,
    checkpoint: Path | str | None = None,
) -> tuple[KarvonenResult, ReportPaths]:
    """Evaluate and report, so the row lands in the same table as every other stage."""
    result = evaluate_karvonen(
        cfg, stage=stage, use_cache=use_cache, device=device, checkpoint=checkpoint
    )
    report = write_report_files(
        result.stage,
        render_markdown(result),
        json.loads(result.model_dump_json()),
        paths.resolve(cfg.out_dir),
        paths.resolve(cfg.web_results) if cfg.web_results else None,
        row_of(result),
    )
    return result, report

@eval_app.command("karvonen")
def eval_karvonen_cmd(
    suite: Annotated[str, typer.Option("--suite", help="Suite name: full or quick.")] = "full",
    config: Annotated[
        Path | None,
        typer.Option(
            "--config", exists=True, dir_okay=False, readable=True, help="Suite YAML override."
        ),
    ] = None,
    stage: Annotated[str, typer.Option("--stage", help="Row name in the results table.")] = (
        "karvonen-8l"
    ),
    checkpoint: Annotated[
        Path | None,
        typer.Option(
            "--checkpoint",
            exists=True,
            dir_okay=False,
            help="A nanoGPT checkpoint to measure instead of the published 8-layer one.",
        ),
    ] = None,
    no_cache: Annotated[bool, typer.Option("--no-cache", help="Recompute everything.")] = False,
    device: Annotated[str | None, typer.Option("--device", help="Where to run.")] = None,
) -> None:
    """Put Adam Karvonen's public chess nanoGPT through the decoder's own suite.

    The one row of the table nobody here trained: same positions, same puzzles, same Stockfish
    ladder as every other stage, so a reader can check the harness against a model with
    published numbers of its own. The 100 MB checkpoint is downloaded from his Hub repository
    on first use; nothing of it is published back.
    """
    from rukh.eval import load_suite
    from rukh.eval.karvonen_suite import run_karvonen_suite
    from rukh.eval.report import elo_line
    from rukh.eval.suite import SUITES

    if suite not in SUITES:
        typer.echo(f"error: --suite must be one of {', '.join(SUITES)}", err=True)
        raise typer.Exit(code=2)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    cfg = load_suite(suite, config)
    try:
        result, report = run_karvonen_suite(
            cfg, stage=stage, use_cache=not no_cache, device=device, checkpoint=checkpoint
        )
    except (FileNotFoundError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    stats = result.written
    typer.echo(f"stage:    {result.stage} ({result.params:,} parameters, baseline)")
    typer.echo(f"source:   {result.source}")
    typer.echo(f"legal:    {stats.legal_rate:.4f} of {stats.asked} answers")
    typer.echo(
        f"failures: {stats.illegal} illegal, {stats.unparseable} not a move, "
        f"{stats.ambiguous} ambiguous, {stats.empty} empty"
    )
    if result.top1 is not None:
        typer.echo(f"accuracy: top1 {result.top1:.4f}")
    if result.puzzles is not None:
        typer.echo(
            f"puzzles:  {result.puzzles.rate:.4f} solved ({result.puzzles_blind} without a prefix)"
        )
    if result.elo is not None:
        typer.echo(f"elo:      {elo_line(result.elo)} over {result.elo.games} games")
    typer.echo(f"report:   {report.markdown}")

"""Tests for rukh.eval.karvonen: the public baseline goes through the harness unchanged.

Nothing here downloads anything. What is worth testing is the part that could silently make the
row wrong: that the checkpoint format loads (the ``_orig_mod.`` prefix of ``torch.compile``),
that the prompt is the author's spelling to the character, that what the model writes is read
the same way Qwen's answers are, and that a failed proposal is rescued *and* counted.
"""

TOY = GPTConfig(n_layer=2, n_head=2, n_embd=32, block_size=64, bias=False, vocab_size=32)


def _board(*sans: str) -> chess.Board:
    board = chess.Board()
    for san in sans:
        board.push_san(san)
    return board


def test_the_vocabulary_is_the_32_characters_of_meta_pkl() -> None:
    assert len(VOCAB) == 32 and len(STOI) == 32
    assert STOI[" "] == 0 and STOI[";"] == 15 and STOI["x"] == 31
    # Every SAN python-chess can write spells with it: captures, promotions, checks, castling.
    for san in ("exd5", "Nbd2", "O-O-O", "e8=Q+", "Qxf7#", "R1a3"):
        assert all(char in STOI for char in san), san


def test_a_compiled_checkpoint_loads_with_its_prefix_stripped(tmp_path: Path) -> None:
    torch.manual_seed(0)
    original = GPT(TOY)
    compiled = {f"_orig_mod.{key}": value for key, value in original.state_dict().items()}
    assert all(key.startswith("_orig_mod.") for key in compiled)
    torch.save(
        {"model_args": {**asdict(TOY), "dropout": 0.0}, "model": compiled}, tmp_path / "ckpt.pt"
    )

    loaded = load_karvonen(tmp_path / "ckpt.pt")
    assert set(strip_compiled(compiled)) == set(loaded.state_dict())
    weight = loaded.transformer.h[1].mlp.c_fc.weight
    assert torch.equal(weight, original.transformer.h[1].mlp.c_fc.weight)
    # The head is tied to the embedding, as nanoGPT ties it: one tensor, not two copies.
    assert loaded.lm_head.weight.data_ptr() == loaded.transformer.wte.weight.data_ptr()
    assert not loaded.training
    idx = torch.tensor([[STOI[c] for c in ";1.e4"]])
    assert loaded(idx).shape == (1, 5, 32)

@pytest.mark.parametrize(
    ("sans", "expected"),
    [
        ((), ";1."),
        (("e4",), ";1.e4"),
        (("e4", "e5"), ";1.e4 e5 2."),
        (("e4", "e5", "Nf3"), ";1.e4 e5 2.Nf3"),
        (("e4", "e5", "Nf3", "Nc6", "Bb5"), ";1.e4 e5 2.Nf3 Nc6 3.Bb5"),
    ],
    ids=["start", "black-to-move", "white-to-move", "three-plies", "five-plies"],
)
def test_the_transcript_is_the_authors_format_to_the_character(
    sans: tuple[str, ...], expected: str
) -> None:
    """``;`` first, no space after the move number, no trailing space when Black is to move."""
    assert transcript(_board(*sans)) == expected


def test_the_transcript_writes_san_in_the_position_it_belongs_to() -> None:
    """``Nbd2`` only reads as such where both knights reach d2; a replay from the root knows."""
    board = _board("Nf3", "d5", "d4", "Nf6", "Nbd2")
    assert transcript(board) == ";1.Nf3 d5 2.d4 Nf6 3.Nbd2"

class _Scripted(nn.Module):
    """A stand-in for the network that writes a fixed string, one character per call."""

    def __init__(self, text: str) -> None:
        super().__init__()
        self.cfg = TOY
        self.script = list(text)
        self.dummy = nn.Parameter(torch.zeros(1))

    def forward(self, idx: Tensor) -> Tensor:
        logits = torch.full((1, idx.shape[1], 32), -10.0)
        char = self.script.pop(0) if self.script else " "
        logits[0, -1, STOI[char]] = 10.0
        return logits


def test_writing_stops_at_the_space_that_ends_the_move() -> None:
    assert KarvonenPlayer(_Scripted("Nf3 Nc6 3.")).write(";1.e4 e5 2.") == "Nf3"
    # Black's move starts with the space the model writes itself; the second space ends it.
    assert KarvonenPlayer(_Scripted(" e5 2.Nf3")).write(";1.e4") == " e5"
    # A new game delimiter means the model thinks the game is over: nothing was written.
    assert KarvonenPlayer(_Scripted(";1.e4")).write(";1.e4 e5 2.") == ""
    # A cap, so a model that never writes a space cannot run for ever.
    assert KarvonenPlayer(_Scripted("a" * 40), max_new_chars=5).write(";1.") == "aaaaa"

def test_an_illegal_proposal_falls_back_and_is_counted(monkeypatch: pytest.MonkeyPatch) -> None:
    """Same rule as the decoder and Qwen: the game goes on, the count does not hide it.

    The queen is boxed in at the start, so ``Qh5`` is well formed and illegal; the rescue is the
    first legal move in the position's own order, and the model is not asked again.
    """
    player = KarvonenPlayer(_Scripted(""))
    monkeypatch.setattr(player, "write", lambda prompt: "Qh5")
    move, was_legal = player.choose(chess.Board())
    assert was_legal is False
    assert move in chess.Board().legal_moves
    assert player.stats.illegal == 1 and player.stats.asked == 1

    monkeypatch.setattr(player, "write", lambda prompt: "Nf9")
    _, was_legal = player.choose(chess.Board())
    assert was_legal is False and player.stats.unparseable == 1

    mate = chess.Board("7k/5KQ1/8/8/8/8/8/8 b - - 0 1")
    assert player.choose(mate) == (None, False)

def test_a_puzzle_prompt_is_rebuilt_from_the_games_own_moves() -> None:
    """The model cannot read a FEN, so the puzzle gets the game in front of it when there is one."""
    tok = UciTokenizer()
    header = [tok.bos_id, tok.vocab[elo_token(1800, "w")], tok.vocab[elo_token(1800, "b")]]
    history = header + [tok.vocab[uci] for uci in ("e2e4", "e7e5", "g1f3")]
    position = chess.Board(_board("e4", "e5", "Nf3").fen())  # a FEN board: no move stack
    assert not position.move_stack

    rebuilt = prefix_board(position, history, tok)
    assert transcript(rebuilt) == ";1.e4 e5 2.Nf3"
    # A prefix that does not lead to this position is not this game: the board is kept as is.
    elsewhere = chess.Board(_board("d4", "d5").fen())
    assert prefix_board(elsewhere, history, tok) is elsewhere

    source = karvonen_source(KarvonenPlayer(_Scripted(" Nc6")), tok)
    assert source(position, history) == chess.Move.from_uci("b8c6")
    assert source.blind == 0
    source(elsewhere, history)
    assert source.blind == 1

def test_the_row_is_marked_as_a_baseline_and_older_rows_are_not() -> None:
    row = row_of(_result())
    assert row.baseline is True
    assert row.stage == "karvonen-8l" and row.legality == pytest.approx(0.75)
    # Rows written before the flag existed carry no key and read as the course's own models.
    assert WebRow.model_validate({"stage": "tiny", "params": 1, "date": "2026-09-19"}).baseline is (
        False
    )


def test_the_benchmarks_table_marks_the_external_row_and_still_renders() -> None:
    rows = [
        row_of(_result()).model_dump(),
        {"stage": "tiny-greedy", "legality": 0.945, "date": "2026-09-20"},
    ]
    table = render_benchmarks(rows)
    assert "| `karvonen-8l` (externo) | 75.0 %" in table
    assert "| `tiny-greedy` | 94.5 %" in table
    assert "2 etapas medidas con la misma suite, 1 de ellas externas" in table