QLoRA con `transformers`: el decoder entra en el ecosistema

"""``MoveDecoder`` dressed as a ``PreTrainedModel``, so the Hugging Face stack works on it.

The decoder is deliberately plain PyTorch: it is the model of this course and reading it should
not require knowing anybody's framework. But the rest of the ecosystem -- ``peft``, ``trl``'s
trainers, ``AutoModel``, the Hub's inference widgets -- speaks ``PreTrainedModel``, and a wrapper
is a hundred lines whereas re-implementing LoRA, DPO and GRPO is not.

The wrapper **owns** a ``MoveDecoder``; it does not reimplement one. So there is exactly one
definition of the forward pass in this repository, and the test that the two give identical logits
is a test that the wrapping is faithful rather than a second implementation drifting from the
first. The weights are the same objects, which is also why ``from_decoder`` costs nothing.

What this buys, concretely, is the check M4 owes the reader: the hand-written LoRA of
``rukh.models.lora`` and ``peft``'s, on the same weights with the same rank and the same targets,
have to produce the same loss curve. Without the wrapper the two could not be put side by side at
all, and "we implemented LoRA" would rest on the author's word.
"""

class RukhConfig(PretrainedConfig):
    """``DecoderConfig`` in the shape ``transformers`` expects to serialise.

    The fields are the decoder's, spelled the same way, so a ``config.json`` written here and the
    YAML a run was configured with say the same thing. ``transformers`` insists on a few names of
    its own (``vocab_size``, ``model_type``); everything else is left alone rather than renamed
    into ``hidden_size``/``num_hidden_layers``, because a reader who has the decoder open should
    recognise what they are looking at.
    """

    model_type = MODEL_TYPE

    def __init__(
        self,
        vocab_size: int = 2030,
        n_layer: int = 12,
        n_head: int = 8,
        d_model: int = 512,
        block: int = 200,
        dropout: float = 0.0,
        pos: str = "learned",
        tie_embeddings: bool = True,
        **kwargs: Any,
    ) -> None:
        self.vocab_size = vocab_size
        self.n_layer = n_layer
        self.n_head = n_head
        self.d_model = d_model
        self.block = block
        self.dropout = dropout
        self.pos = pos
        self.tie_embeddings = tie_embeddings
        kwargs.setdefault("tie_word_embeddings", tie_embeddings)
        super().__init__(**kwargs)

    @classmethod
    def from_decoder(cls, cfg: DecoderConfig, **kwargs: Any) -> RukhConfig:
        return cls(
            vocab_size=cfg.vocab_size,
            n_layer=cfg.n_layer,
            n_head=cfg.n_head,
            d_model=cfg.d_model,
            block=cfg.block,
            dropout=cfg.dropout,
            pos=cfg.pos,
            tie_embeddings=cfg.tie_embeddings,
            **kwargs,
        )

    def decoder_config(self) -> DecoderConfig:
        """The ``DecoderConfig`` this config describes.

        Not called ``decoder``: ``transformers`` reads a config attribute of that name as the
        decoder half of an encoder-decoder pair and tries to call ``to_dict`` on it, which turns
        a bound method into a crash the first time anything asks for a generation config.
        """
        return DecoderConfig(
            vocab_size=self.vocab_size,
            n_layer=self.n_layer,
            n_head=self.n_head,
            d_model=self.d_model,
            block=self.block,
            dropout=self.dropout,
            pos=self.pos,  # type: ignore[arg-type]
            tie_embeddings=self.tie_embeddings,
        )

class RukhForCausalLM(PreTrainedModel, GenerationMixin):
    """A ``PreTrainedModel`` whose whole body is one ``MoveDecoder``.

    ``GenerationMixin`` comes along because ``prepare_inputs_for_generation`` is defined and a
    class that answers it without inheriting the mixin is a model that looks generative and is
    not. Nothing in this project generates through ``.generate()`` -- the sampler of
    ``rukh.infer`` does it with the legality mask, which ``transformers`` knows nothing about --
    but a wrapper that half-implements an interface is worse than one that implements it.
    """

    config_class = RukhConfig
    base_model_prefix = "decoder"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Block"]
    _tied_weights_keys = {"decoder.lm_head.weight": "decoder.tokens.weight"}
    """Which name is the copy and which is the storage: the head shares its tensor with the
    embedding table, because ``MoveDecoder`` ties them in its constructor. ``safetensors`` stores
    tensors and not aliases, so ``save_pretrained`` refuses a state dict with two names for one
    tensor unless it is told. This is the same fact ``rukh.train.checkpoint.TIED_SOURCES`` records
    for the project's own checkpoint format."""

    def __init__(self, config: RukhConfig) -> None:
        super().__init__(config)
        self.decoder = MoveDecoder(config.decoder_config())
        # Registers the tied-weight map and the other static properties on the top model. It does
        # not touch the weights in this version of `transformers`, so calling it after building
        # the decoder is safe; skipping it leaves `all_tied_weights_keys` unset and every save or
        # load fails on an attribute the framework expects every model to have.
        self.post_init()

    @classmethod
    def from_decoder(cls, decoder: MoveDecoder, **kwargs: Any) -> RukhForCausalLM:
        """Wrap an existing decoder **without copying its weights**.

        The tensors are shared, so training the wrapper trains the decoder and a checkpoint taken
        from either is the same checkpoint. That is the property that makes the peft comparison
        meaningful: both sides start from literally the same numbers.
        """
        model = cls(RukhConfig.from_decoder(decoder.cfg, **kwargs))
        model.decoder = decoder
        return model

    def get_input_embeddings(self) -> nn.Module:
        return self.decoder.tokens

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.decoder.tokens = value  # type: ignore[assignment]

    def get_output_embeddings(self) -> nn.Module:
        return self.decoder.lm_head

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,  # noqa: ARG002 - the decoder is always causal
        return_dict: bool | None = None,  # noqa: ARG002 - a dataclass is the only output shape
        **kwargs: Any,
    ) -> CausalLMOutput:
        """Logits, and the decoder's own loss when ``labels`` are given.

        ``labels`` are **already shifted** here, unlike most causal models in ``transformers``.
        The packed stream of this project stores ``(x, y)`` pairs where ``y`` is ``x`` moved one
        step, so shifting again inside the model would teach it to predict the move after next.
        The loss is the decoder's, ``ignore_index=0``, so a run through the wrapper and a run
        through the training loop optimise the same number.
        """
        logits, loss = self.decoder(input_ids, labels)
        return CausalLMOutput(loss=loss, logits=logits)

    def prepare_inputs_for_generation(
        self, input_ids: torch.Tensor, **kwargs: Any
    ) -> dict[str, Any]:
        """Crop to the context window; there is no KV cache to carry between steps."""
        return {"input_ids": input_ids[:, -self.config.block :]}

    @property
    def ignore_index(self) -> int:
        """The id the loss skips: ``<pad>``, id 0, in every scheme of this project."""
        return IGNORE_INDEX

def test_the_hand_written_lora_and_peft_optimise_the_same_number() -> None:
    """The check that makes the chapter honest.

    Both sides start from the same base weights and are given the same `A` (copied across, so the
    comparison is of the maths and not of two random draws), then take the same optimiser steps on
    the same batch. If the update rule differed at all -- the scaling, which side `A` multiplies,
    whether `B` starts at zero -- the losses would separate immediately.
    """
    x, y = _batch()
    r, alpha, steps = 4, 8, 6

    ours = _decoder()
    apply_lora(ours, LoraConfig(r=r, alpha=alpha, targets=("attn_out",)))
    mine = [p for p in ours.parameters() if p.requires_grad]

    theirs = RukhForCausalLM.from_decoder(_decoder())
    adapted = peft.get_peft_model(theirs, _peft_config(r, alpha, ["attn.proj"]))

    # Copy our `A` into theirs so both start from the same point; `B` is zero on both sides.
    ours_by_layer = {name: module for name, module in lora_modules(ours)}
    with torch.no_grad():
        for name, module in adapted.named_modules():
            if not hasattr(module, "lora_A") or "default" not in getattr(module, "lora_A", {}):
                continue
            key = name.replace("base_model.model.decoder.", "").replace(".base_layer", "")
            module.lora_A["default"].weight.copy_(ours_by_layer[key].a[0])

    mine_opt = torch.optim.SGD(mine, lr=0.5)
    their_opt = torch.optim.SGD([p for p in adapted.parameters() if p.requires_grad], lr=0.5)

    my_curve: list[float] = []
    their_curve: list[float] = []
    for _ in range(steps):
        mine_opt.zero_grad()
        loss = ours(x, y)[1]
        assert loss is not None
        loss.backward()
        mine_opt.step()
        my_curve.append(float(loss.detach()))

        their_opt.zero_grad()
        other = adapted(x, labels=y).loss
        other.backward()
        their_opt.step()
        their_curve.append(float(other.detach()))

    assert my_curve[-1] < my_curve[0]  # it is training at all
    for step, (a, b) in enumerate(zip(my_curve, their_curve, strict=True)):
        assert a == pytest.approx(b, abs=1e-5), f"step {step}: {a} vs {b}"


def test_peft_leaves_the_base_weights_frozen_exactly_as_ours_does() -> None:
    theirs = RukhForCausalLM.from_decoder(_decoder())
    adapted = peft.get_peft_model(theirs, _peft_config(4, 8, ["attn.proj"]))
    trainable = {name for name, p in adapted.named_parameters() if p.requires_grad}
    assert trainable
    assert all("lora_" in name for name in trainable)

[project.optional-dependencies]
cpu = ["torch>=2.11,<2.12"]
cu128 = ["torch>=2.11,<2.12"]
# The Hugging Face training stack, needed only by M4: the `PreTrainedModel` wrapper around
# `MoveDecoder`, `peft` as the reference LoRA the hand-written one is checked against, and the
# QLoRA fine-tune of Qwen3 on PGN. It is an extra and not a core dependency because the CPU CI
# job has a three-minute budget and nothing in P0-P3 imports any of it.
hf = ["transformers>=5.17", "peft>=0.21", "trl>=1.13", "accelerate>=1.15", "bitsandbytes>=0.50"]

"""QLoRA of a general language model on the same games, so the comparison is fair.

M4 owes the reader a number for a question the whole project raises by existing: why build a
115 M decoder for chess when a general model five times its size can be fine-tuned in an
afternoon? The only answer worth printing is a measurement, and a measurement is only worth
printing when the two sides differ in one thing. So:

- the **games** are the same, drawn from the same month and held out against the same validation
  split (``rukh.data.pgn_text``);
- the **metrics** are the same, because the fine-tuned Qwen is plugged into the harness of P2 as
  another ``MoveSource`` (``rukh.eval.qwen_source``) and plays the same Stockfish ladder;
- the **representation** is what differs, and that is the point. The decoder reads one token per
  legal move; Qwen reads ``1. e4 e5 2. Nf3`` cut into whatever pieces its BPE makes of it, and
  nothing in it forbids writing a move that does not exist.

This file is the training half, and it is deliberately somebody else's code: ``transformers`` for
the model, ``peft`` for the adapter, ``trl`` for the loop. The project has its own LoRA and its own
loop and does not need either here -- what it needs is the experience of the ecosystem path,
which is what a reader will actually reach for, and a run that is comparable to the published
QLoRA recipes rather than to ours.

Four-bit is attempted and not required. The memory it saves is the reason QLoRA exists, and at
0.6 B parameters on a 32 GB card there is nothing to save: the honest thing is to try it, report
whether it worked, and fall back to bf16 rather than pretend a technique was demonstrated when
the dependency would not build.
"""

from __future__ import annotations

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

from pydantic import Field

from rukh.config import BaseConfig
from rukh.paths import resolve

if TYPE_CHECKING:
    from collections.abc import Mapping

log = logging.getLogger(__name__)

__all__ = ["QwenConfig", "QwenReport", "four_bit_config", "train_qwen"]

def four_bit_config() -> tuple[Any | None, str | None]:
    """``(BitsAndBytesConfig, None)`` when 4-bit is available, ``(None, reason)`` when it is not.

    NF4 with double quantisation and a bf16 compute dtype: the recipe of the QLoRA paper, so the
    run is comparable with the published ones rather than with a variant of our own.
    """
    try:
        import bitsandbytes  # noqa: F401
        import torch
        from transformers import BitsAndBytesConfig
    except ImportError as exc:
        return None, f"bitsandbytes is not installed ({exc})"
    return (
        BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
        ),
        None,
    )

def _dataset(cfg: QwenConfig, split: str) -> Any:
    from datasets import load_dataset

    path = resolve(cfg.data_dir) / f"{split}.jsonl"
    if not path.is_file():
        raise FileNotFoundError(f"{path} not found; run `rukh data pgn-text` first")
    data = load_dataset("json", data_files={split: path.as_posix()})[split]
    if split == "train" and cfg.n_train is not None:
        data = data.select(range(min(cfg.n_train, len(data))))
    return data


def _count_params(model: Any) -> tuple[int, int]:
    """``(trainable, total)``, counting a 4-bit weight as the parameters it stands for.

    ``bitsandbytes`` stores two 4-bit values per byte, so ``numel()`` on a quantised weight
    returns half the parameter count and the "0.1 % trainable" line would come out twice as
    flattering as the truth.
    """
    trainable = total = 0
    for param in model.parameters():
        count = param.numel()
        if param.__class__.__name__ == "Params4bit":
            count *= 2
        total += count
        if param.requires_grad:
            trainable += count
    return trainable, total

def train_qwen(cfg: QwenConfig, device: str | None = None) -> QwenReport:
    """Fine-tune ``cfg.model`` on PGN text with LoRA, in 4 bits when that is possible."""
    import torch
    from peft import LoraConfig as PeftLoraConfig
    from peft import get_peft_model, prepare_model_for_kbit_training
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from trl import SFTConfig, SFTTrainer

    where = device or ("cuda" if torch.cuda.is_available() else "cpu")
    quantization, reason = (None, "not requested")
    if cfg.four_bit:
        quantization, reason = four_bit_config()
    if quantization is None and cfg.four_bit:
        log.warning("4-bit unavailable, training in bf16 instead: %s", reason)

    tokenizer = AutoTokenizer.from_pretrained(cfg.model)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        cfg.model,
        dtype=torch.bfloat16,
        quantization_config=quantization,
        device_map={"": where} if quantization is not None else None,
    )
    if quantization is not None:
        model = prepare_model_for_kbit_training(model)
    elif where == "cuda":
        model = model.to(where)
    model.config.use_cache = False

    weights_memory = torch.cuda.memory_allocated() / 1e6 if where == "cuda" else None
    adapted = get_peft_model(
        model,
        PeftLoraConfig(
            r=cfg.r,
            lora_alpha=cfg.alpha,
            lora_dropout=cfg.dropout,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=list(cfg.targets),
        ),
    )
    trainable, total = _count_params(adapted)
    log.info(
        "%s: %s trainable of %s (%.3f %%), %s",
        cfg.model,
        f"{trainable:,}",
        f"{total:,}",
        100 * trainable / max(total, 1),
        "4-bit" if quantization is not None else "bf16",
    )

    out_dir = resolve(cfg.out_dir) / cfg.run_name
    args = SFTConfig(
        output_dir=out_dir.as_posix(),
        dataset_text_field="text",
        max_length=cfg.max_length,
        packing=False,
        per_device_train_batch_size=cfg.batch_size,
        gradient_accumulation_steps=cfg.grad_accum,
        learning_rate=cfg.lr,
        max_steps=cfg.max_steps,
        warmup_steps=cfg.warmup,
        lr_scheduler_type="cosine",
        logging_steps=25,
        save_strategy="no",
        bf16=True,
        seed=cfg.seed,
        report_to=["mlflow"],
        run_name=cfg.run_name,
    )
    trainer = SFTTrainer(
        model=adapted,
        args=args,
        train_dataset=_dataset(cfg, "train"),
        processing_class=tokenizer,
    )
    if where == "cuda":
        torch.cuda.reset_peak_memory_stats()
    result = trainer.train()
    peak = torch.cuda.max_memory_allocated() / 1e6 if where == "cuda" else None

    out_dir.mkdir(parents=True, exist_ok=True)
    adapted.save_pretrained(out_dir.as_posix())
    tokenizer.save_pretrained(out_dir.as_posix())
    report = QwenReport(
        model=cfg.model,
        adapter_dir=out_dir.as_posix(),
        four_bit=quantization is not None,
        fallback_reason=None if quantization is not None else reason,
        trainable_params=trainable,
        total_params=total,
        train_samples=len(trainer.train_dataset),
        steps=cfg.max_steps,
        final_loss=float(result.training_loss) if result.training_loss is not None else None,
        weights_memory_mb=weights_memory,
        peak_memory_mb=peak,
    )
    (out_dir / "run.json").write_text(report.model_dump_json(indent=2) + "\n", encoding="utf-8")
    return report

def load_qwen_adapter(adapter_dir: Path | str, device: str | None = None) -> tuple[Any, Any]:
    """``(model, tokenizer)`` of a finished run, merged and in eval mode.

    Merged on purpose: the evaluation plays hundreds of games move by move, and a merged model
    runs the base matrices instead of the base plus two extra matmuls per adapted projection.
    """
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    directory = Path(adapter_dir)
    config: Mapping[str, Any] = json.loads(
        (directory / "adapter_config.json").read_text(encoding="utf-8")
    )
    where = device or ("cuda" if torch.cuda.is_available() else "cpu")
    base = AutoModelForCausalLM.from_pretrained(
        str(config["base_model_name_or_path"]), dtype=torch.bfloat16
    )
    model = PeftModel.from_pretrained(base, directory.as_posix())
    model = model.merge_and_unload().to(where).eval()
    tokenizer = AutoTokenizer.from_pretrained(directory.as_posix())
    return model, tokenizer

# Qwen3-0.6B fine-tuned on the same games the decoder trains on, written as PGN text.
#
# The point is the comparison, so the budget is chosen to be *generous* to Qwen rather than fair
# to us: 1 500 steps of 32 sequences of up to 512 tokens is about 24 M tokens of chess for a model
# that already speaks, against the 1.68 B the decoder needed from scratch. If it still loses by a
# wide margin, nobody can say it was starved.
#
# `r = 16` on all four attention projections is the usual QLoRA recipe, twice the rank the
# project's own adapters use, for the same reason: this side gets the benefit of the doubt.
model: Qwen/Qwen3-0.6B
data_dir: data/pgn-text
out_dir: checkpoints
run_name: qwen3-pgn-qlora
max_length: 512
batch_size: 8
grad_accum: 4
lr: 2.0e-4
max_steps: 1500
warmup: 100
r: 16
alpha: 32
dropout: 0.05
targets: [q_proj, k_proj, v_proj, o_proj]
four_bit: true
n_train: 100000
seed: 42

# Two steps over a hundred games: not a run, a rehearsal.
#
# The real Qwen fine-tune is half an hour of GPU and everything that can fail in it fails in the
# first thirty seconds -- the model download, the quantisation backend, the `trl` API, the tokenizer
# padding, the save. Finding that out at minute twenty-five is the expensive way.

"""A fine-tuned general model plugged into the harness of P2, so the comparison is like for like.

The decoder cannot write a move that does not exist: its vocabulary *is* the set of moves, and an
"illegal" move is a legal move played in the wrong position. A language model writing SAN can fail
in three ways instead of one, and telling them apart is most of what this comparison teaches:

- **unparseable** -- what it wrote is not a move at all (``Nf9``, ``O-O-O-O``, a sentence);
- **illegal** -- a well-formed move that this position does not allow;
- **ambiguous** -- a move that two pieces could make, which SAN is supposed to disambiguate and
  the model did not.

All three are counted separately and none is silently repaired. A game still has to finish, so a
failed proposal falls back to a legal move, exactly as the decoder's does, and the fallback is
recorded rather than hidden -- same rule as ``illegal_proposals`` in ``rukh.infer.game``.

The prompt is the PGN text of ``rukh.data.pgn_text``, which is the format the model was fine-tuned
on, with the same rating header the decoder gets as ``<wXXXX> <bXXXX>``. Generation is greedy and
six tokens long: a move is one to four of Qwen's tokens and nothing longer is going to help.
"""

from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING, Any

import chess
from pydantic import BaseModel, ConfigDict

from rukh.data.pgn_text import prompt_for

if TYPE_CHECKING:
    pass

log = logging.getLogger(__name__)

__all__ = ["QwenPlayer", "QwenStats", "parse_san_answer", "qwen_source"]

MAX_NEW_TOKENS = 6
DECODER_PLIES = 196
"""Plies the decoder can carry (``block`` 200 minus its three header tokens and the move it is
about to make). The text model gets the same number, not the 120 its fine-tuning corpus was cut
at: a harness that let one side play longer games than the other would be measuring the harness.
Past 120 plies the prompt is longer than anything Qwen was fine-tuned on, and that is a property
of the model rather than of the measurement -- it belongs in the result, not in the setup."""
MOVE_NUMBER = re.compile(r"^\d+\.+")
"""``12.`` or ``12...`` at the head of the answer: bookkeeping, not a move."""
TRAILING = "!?"
"""Annotation glyphs a commentator adds; ``+`` and ``#`` are part of SAN and stay."""

class QwenStats(BaseModel):
    """How the answers came out, by failure mode rather than as one number."""

    model_config = ConfigDict(extra="forbid")

    asked: int = 0
    legal: int = 0
    unparseable: int = 0
    """The answer was not a move in any position."""
    illegal: int = 0
    """A well-formed move that this position does not allow."""
    ambiguous: int = 0
    """SAN that two pieces could satisfy; the model failed to disambiguate."""
    empty: int = 0
    """The model produced no non-whitespace token at all."""

    @property
    def legal_rate(self) -> float:
        return self.legal / self.asked if self.asked else 0.0

    def record(self, outcome: str) -> None:
        self.asked += 1
        setattr(self, outcome, getattr(self, outcome) + 1)

def parse_san_answer(board: chess.Board, answer: str) -> tuple[chess.Move | None, str]:
    """``(move, outcome)`` for the model's continuation, with the outcome named.

    Only the first token is considered. A model that writes ``e4 e5 2. Nf3`` has answered the
    question asked and then kept going; scoring the whole continuation would be scoring a
    different question, and taking the last move would be scoring luck.
    """
    text = answer.strip()
    if not text:
        return None, "empty"
    first = text.split()[0]
    stripped = MOVE_NUMBER.sub("", first)
    if not stripped:  # the answer was just a move number; the move is the next token
        parts = text.split()
        if len(parts) < 2:
            return None, "empty"
        stripped = parts[1]
    stripped = stripped.rstrip(TRAILING)
    if not stripped:
        return None, "unparseable"
    try:
        return board.parse_san(stripped), "legal"
    except chess.IllegalMoveError:
        return None, "illegal"
    except chess.AmbiguousMoveError:
        return None, "ambiguous"
    except (chess.InvalidMoveError, ValueError):
        return None, "unparseable"

class QwenPlayer:
    """A ``rukh.infer.Player`` backed by a causal language model that writes SAN."""

    def __init__(
        self,
        model: Any,
        tokenizer: Any,
        stats: QwenStats | None = None,
        max_new_tokens: int = MAX_NEW_TOKENS,
        max_plies: int = DECODER_PLIES,
    ) -> None:
        self.model, self.tokenizer = model, tokenizer
        self.stats = stats if stats is not None else QwenStats()
        self.max_new_tokens = max_new_tokens
        self.max_plies = max_plies
        self.white_elo = self.black_elo = 1800

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

    def start(self, white_elo: int, black_elo: int) -> None:
        self.white_elo, self.black_elo = white_elo, black_elo

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

    def answer(self, board: chess.Board) -> str:
        """The raw continuation the model writes for this position."""
        import torch

        prompt = prompt_for(board, self.white_elo, self.black_elo)
        encoded = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        with torch.no_grad():
            out = self.model.generate(
                **encoded,
                max_new_tokens=self.max_new_tokens,
                do_sample=False,
                pad_token_id=self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
            )
        written = out[0, encoded["input_ids"].shape[1] :]
        return str(self.tokenizer.decode(written, skip_special_tokens=True))

    def propose(self, board: chess.Board) -> tuple[chess.Move | None, str]:
        """What the model wrote, parsed, with the outcome named. No mask, no rescue."""
        move, outcome = parse_san_answer(board, self.answer(board))
        self.stats.record(outcome)
        return move, outcome

    def choose(self, board: chess.Board) -> tuple[chess.Move | None, bool]:
        """The move to play, and whether the model's own answer was usable.

        The fallback is the first legal move in the position's own order, not a second sample:
        the model already answered, and asking again until it says something legal would measure
        a different model from the one the legality rate describes.
        """
        move, outcome = self.propose(board)
        if move is not None:
            return move, True
        legal = list(board.legal_moves)
        return (legal[0] if legal else None), False


def qwen_source(player: QwenPlayer) -> Any:
    """The ``MoveSource`` the puzzle suite expects, backed by the same player.

    Puzzles are scored on the move the model *chooses*, so the fallback applies here too and a
    position where it wrote nonsense counts as a failed puzzle rather than as a missing one.
    """

    def choose(board: chess.Board, history: list[int]) -> chess.Move | None:  # noqa: ARG001
        move, _ = player.choose(board)
        return move

    return choose