El encoder a mano: el mismo bloque con `causal=False`

"""``PositionEncoder``: a hand-written bidirectional transformer over positions.

The brother of ``MoveDecoder``, and deliberately built from the same ``rukh.models.layers``
blocks: the whole difference is the attention mask. The decoder answers "what comes next?", so
a position may only look left; the encoder answers "what is this position?", so every token
sees every other one, including the ones after it. That is why the test that proves the decoder
right (a future token never changes a past logit) has to fail here, and ``test_encoder`` asserts
the opposite.

Two input schemes share the class (``EncoderConfig.input``):

``moves``
    the game so far in the P1 UCI vocabulary, the very tokens the decoder was trained on, so
    the masked-move head can be tied to the embedding and a pretrained decoder's intuition is
    directly comparable.
``squares``
    the 69 tokens of ``rukh.models.squares``, the position itself rather than its history.

Padding is expressed as a key mask, not as a special attention: ``attention_mask`` is ``True``
on real tokens and ``False`` on ``<pad>``, and the padded keys are switched off for every
query. The rows of padded queries are still computed (nobody reads them) and at least one real
token per row is required, otherwise softmax would see an entirely masked row and return NaN.
"""

from __future__ import annotations

import math
from typing import Literal

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

from rukh.models import layers
from rukh.models.config import EncoderConfig
from rukh.models.layers import rope_tables

PAD_ID = 0
"""``<pad>`` is id 0 in both schemes."""

MMM_IGNORE_INDEX = -100
"""Label of a position the masked-move loss must skip.

Not the decoder's ``0``. ``0`` is ``<pad>`` in **both** schemes, so using it to mean "nothing to
predict" would overload one id with two jobs: the loss could no longer tell a position it must
skip from a position where the right answer happens to be ``<pad>``. The decoder gets away with
it because its target is the next token of a packed stream and ``<pad>`` is never a target there;
here the distinction has to be explicit, and ``-100`` is outside every vocabulary, so "not
predicted" and "predict ``<pad>``" stay different things.
"""

def _tracing() -> bool:
    """``True`` while ``torch.export`` or ``torch.compile`` is capturing a graph.

    The padding check below reads a tensor to decide whether to raise, and a data-dependent
    branch like that is precisely what a graph capture cannot represent. Under ``torch.export``
    it would send the exporter back to the deprecated TorchScript tracer (D-027); under
    ``torch.compile`` it is worse in a quieter way — Dynamo cannot prove the condition, so it
    graph-breaks around it on **every** masked step of the MMM loop, which is a synchronisation
    point plus two half-graphs per step for a check that has already passed. Eager behaviour is
    unchanged: outside a capture both calls are ``False`` and the ValueError is raised as before.
    """
    return torch.compiler.is_exporting() or torch.compiler.is_compiling()

class PositionEncoder(nn.Module):
    """Bidirectional transformer encoder: ``forward`` returns the hidden states ``(B, T, d)``."""

    def __init__(self, cfg: EncoderConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or EncoderConfig()
        self.tokens = nn.Embedding(self.cfg.tokens, self.cfg.d_model)
        self.positions = (
            nn.Embedding(self.cfg.seq, self.cfg.d_model) if self.cfg.pos == "learned" else None
        )
        self.drop = nn.Dropout(self.cfg.dropout)
        self.blocks = nn.ModuleList(
            layers.Block(
                self.cfg.d_model, self.cfg.n_head, self.cfg.ff, self.cfg.dropout, causal=False
            )
            for _ in range(self.cfg.n_layer)
        )
        self.ln_f = nn.LayerNorm(self.cfg.d_model)
        self.mlm_head = nn.Linear(self.cfg.d_model, self.cfg.tokens, bias=False)
        if self.cfg.tie_embeddings and self.cfg.input == "moves":
            # Tied only for ``moves``: it is the P1 vocabulary, where the embedding and the head
            # describe the same 2 030 moves. The 47 square tokens are too few to be worth tying.
            self.mlm_head.weight = self.tokens.weight
        if self.cfg.pos == "rope":
            cos, sin = rope_tables(self.cfg.seq, self.cfg.head_dim, torch.device("cpu"))
            self.register_buffer("rope_cos", cos, persistent=False)
            self.register_buffer("rope_sin", sin, persistent=False)
        self.apply(layers.init_weights)
        scale = 0.02 / math.sqrt(2 * self.cfg.n_layer)
        for name, param in self.named_parameters():
            if name.endswith("proj.weight"):
                nn.init.normal_(param, mean=0.0, std=scale)

    def num_params(self, non_embedding: bool = True) -> int:
        """Parameter count; ``non_embedding`` drops the learned position table (nanoGPT rule)."""
        total = sum(p.numel() for p in self.parameters())
        if non_embedding and self.positions is not None:
            total -= self.positions.weight.numel()
        return total

    @staticmethod
    def padding_mask(idx: Tensor) -> Tensor:
        """``True`` where ``idx`` is a real token, ``False`` on ``<pad>``; shape ``(B, T)``."""
        return idx != PAD_ID

    @staticmethod
    def _key_mask(attention_mask: Tensor | None) -> Tensor | None:
        """``(B, T)`` into the ``(B, 1, 1, T)`` boolean key mask SDPA expects."""
        if attention_mask is None:
            return None
        mask = attention_mask.bool()
        if mask.dim() != 2:
            raise ValueError(f"attention_mask must be (B, T), got {tuple(attention_mask.shape)}")
        if not _tracing() and not bool(mask.any(dim=-1).all()):
            # An entirely masked row would make softmax return NaN, so it is refused here rather
            # than debugged three layers down. The check reads a tensor, which is exactly what
            # `torch.export` cannot trace (a data-dependent guard), and skipping it under the
            # exporter is what keeps the encoder on the modern exporter instead of the
            # deprecated tracer: the exported graph is a pure function of its input either way.
            raise ValueError("every sequence needs at least one unmasked token")
        return mask[:, None, None, :]

    def forward(self, idx: Tensor, attention_mask: Tensor | None = None) -> Tensor:
        """Hidden states ``(B, T, d_model)``; ``attention_mask`` is ``True`` on real tokens."""
        seq = idx.shape[-1]
        if seq > self.cfg.seq:
            raise ValueError(f"sequence of {seq} tokens is longer than block {self.cfg.seq}")
        key_mask = self._key_mask(attention_mask)
        x = self.tokens(idx)
        cos = sin = None
        if self.positions is not None:
            steps = torch.arange(seq, device=idx.device)
            x = x + self.positions(steps)
        else:
            cos, sin = self.rope_cos, self.rope_sin
        x = self.drop(x)
        for block in self.blocks:
            x = block(x, cos, sin, key_mask)
        return self.ln_f(x)

    def masked_lm(
        self, idx: Tensor, labels: Tensor | None = None, attention_mask: Tensor | None = None
    ) -> tuple[Tensor, Tensor | None]:
        """Masked-move logits ``(B, T, V)`` and, with ``labels``, the scalar loss.

        ``labels`` is ``MMM_IGNORE_INDEX`` wherever nothing was masked; see ``rukh.train.mmm``.
        """
        logits = self.mlm_head(self(idx, attention_mask))
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]),
                labels.reshape(-1),
                ignore_index=MMM_IGNORE_INDEX,
            )
        return logits, loss

    def pool(
        self,
        idx: Tensor,
        how: Literal["cls", "mean"] = "mean",
        attention_mask: Tensor | None = None,
    ) -> Tensor:
        """One vector per sequence, ``(B, d_model)``: the position's representation.

        ``cls`` takes position 0 (``<bos>`` for ``moves``, ``<cls>`` for ``squares``) and
        ``mean`` averages the real tokens only. Without an explicit ``attention_mask`` the
        padding is read off ``idx`` itself, so ``pool(idx, "mean")`` never averages ``<pad>``
        into the representation, which is the whole point of the operation.
        """
        if how not in ("cls", "mean"):
            raise ValueError(f"how must be 'cls' or 'mean', got {how!r}")
        mask = self.padding_mask(idx) if attention_mask is None else attention_mask.bool()
        hidden = self(idx, mask)
        if how == "cls":
            return hidden[:, 0]
        weights = mask.unsqueeze(-1).to(hidden.dtype)
        return (hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1.0)

"""Tests for rukh.models.encoder: shapes, bidirectionality, padding, pooling and the preset."""

from __future__ import annotations

import pytest
import torch

from rukh.models import EncoderConfig, PositionEncoder
from rukh.models.encoder import MMM_IGNORE_INDEX
from rukh.models.squares import SQUARE_TOKENS, SQUARE_VOCAB_SIZE, fen_to_tokens

pytestmark = pytest.mark.unit

TOY = EncoderConfig(vocab_size=64, n_layer=2, n_head=4, d_model=32, block=16, dropout=0.0)
START = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq -"


def toy_model(seed: int = 0, **overrides: object) -> PositionEncoder:
    """A deterministic two-layer encoder in eval mode."""
    torch.manual_seed(seed)
    return PositionEncoder(TOY.model_copy(update=overrides)).eval()


def toy_batch(batch: int = 2, seq: int = 8, seed: int = 1) -> torch.Tensor:
    generator = torch.Generator().manual_seed(seed)
    return torch.randint(1, TOY.vocab_size, (batch, seq), generator=generator)

def test_forward_returns_one_hidden_state_per_token() -> None:
    hidden = toy_model()(toy_batch())
    assert hidden.shape == (2, 8, TOY.d_model)


@pytest.mark.parametrize("pos", ["learned", "rope"])
def test_a_later_token_does_change_the_earlier_outputs(pos: str) -> None:
    """The inverse of the decoder's causality test: this model is bidirectional on purpose."""
    model = toy_model(pos=pos)
    idx = toy_batch(seq=12)
    cut = 5
    changed = idx.clone()
    changed[:, cut + 1 :] = (changed[:, cut + 1 :] + 7) % (TOY.vocab_size - 1) + 1
    assert not torch.equal(idx, changed)
    with torch.no_grad():
        base, other = model(idx), model(changed)
    assert not torch.allclose(base[:, : cut + 1], other[:, : cut + 1], atol=1e-5)

def test_padding_does_not_leak_into_the_real_tokens() -> None:
    model = toy_model()
    short = toy_batch(seq=6, seed=2)
    padded = torch.cat([short, torch.zeros((2, 4), dtype=torch.long)], dim=1)
    with torch.no_grad():
        plain = model(short, model.padding_mask(short))
        with_padding = model(padded, model.padding_mask(padded))
    assert torch.allclose(plain, with_padding[:, :6], atol=1e-5)
    # Without the mask the padding is just another token and does change the answer.
    with torch.no_grad():
        unmasked = model(padded)
    assert not torch.allclose(plain, unmasked[:, :6], atol=1e-5)


def test_an_entirely_masked_sequence_is_refused() -> None:
    model = toy_model()
    idx = torch.zeros((2, 6), dtype=torch.long)
    with pytest.raises(ValueError, match="at least one unmasked token"):
        model(idx, model.padding_mask(idx))
    with pytest.raises(ValueError, match=r"\(B, T\)"):
        model(toy_batch(), torch.ones((2, 1, 8), dtype=torch.bool))

def test_mean_pooling_ignores_the_padding() -> None:
    model = toy_model()
    short = toy_batch(seq=6, seed=3)
    padded = torch.cat([short, torch.zeros((2, 5), dtype=torch.long)], dim=1)
    with torch.no_grad():
        assert torch.allclose(model.pool(short, "mean"), model.pool(padded, "mean"), atol=1e-5)
        assert torch.allclose(model.pool(short, "cls"), model.pool(padded, "cls"), atol=1e-5)
        hidden = model(short, model.padding_mask(short))
        assert torch.allclose(model.pool(short, "mean"), hidden.mean(dim=1), atol=1e-5)
        assert torch.allclose(model.pool(short, "cls"), hidden[:, 0], atol=1e-6)
    assert model.pool(short).shape == (2, TOY.d_model)  # mean is the default
    with pytest.raises(ValueError, match="cls"):
        model.pool(short, "first")  # type: ignore[arg-type]

def test_masked_lm_scores_only_the_masked_positions() -> None:
    model = toy_model()
    idx = toy_batch(seq=8, seed=4)
    labels = torch.full_like(idx, MMM_IGNORE_INDEX)
    labels[:, 2] = idx[:, 2]
    logits, loss = model.masked_lm(idx, labels)
    assert logits.shape == (2, 8, TOY.vocab_size)
    assert loss is not None and loss.ndim == 0 and torch.isfinite(loss)
    assert model.masked_lm(idx)[1] is None

    # Moving a label to another position changes the loss; ignoring everything gives no loss.
    elsewhere = torch.full_like(idx, MMM_IGNORE_INDEX)
    elsewhere[:, 5] = idx[:, 5]
    assert not torch.allclose(loss, model.masked_lm(idx, elsewhere)[1], atol=1e-6)
    assert torch.isnan(model.masked_lm(idx, torch.full_like(idx, MMM_IGNORE_INDEX))[1])
    # ``<pad>`` (0) is a legitimate label here, unlike in the decoder.
    pad_label = torch.full_like(idx, MMM_IGNORE_INDEX)
    pad_label[:, 1] = 0
    assert torch.isfinite(model.masked_lm(idx, pad_label)[1])

def test_the_masked_move_head_is_tied_only_for_moves() -> None:
    moves = toy_model()
    assert moves.mlm_head.weight is moves.tokens.weight
    untied = toy_model(tie_embeddings=False)
    assert untied.mlm_head.weight is not untied.tokens.weight
    torch.manual_seed(0)
    squares = PositionEncoder(EncoderConfig(input="squares", n_layer=2, n_head=4, d_model=32))
    assert squares.mlm_head.weight is not squares.tokens.weight
    assert squares.tokens.num_embeddings == SQUARE_VOCAB_SIZE


def test_the_squares_scheme_has_69_fixed_positions() -> None:
    cfg = EncoderConfig(input="squares", n_layer=2, n_head=4, d_model=32, dropout=0.0)
    assert cfg.seq == SQUARE_TOKENS and cfg.tokens == SQUARE_VOCAB_SIZE
    torch.manual_seed(0)
    model = PositionEncoder(cfg).eval()
    assert model.positions is not None and model.positions.num_embeddings == SQUARE_TOKENS
    idx = torch.tensor([fen_to_tokens(START), fen_to_tokens(f"{START} 30 40")])
    with torch.no_grad():
        assert model(idx).shape == (2, SQUARE_TOKENS, 32)
        assert model.pool(idx, "cls").shape == (2, 32)
    with pytest.raises(ValueError, match="longer than block"):
        model(torch.zeros((1, SQUARE_TOKENS + 1), dtype=torch.long))

def test_same_seed_gives_the_same_weights_and_hidden_states() -> None:
    idx = toy_batch()
    with torch.no_grad():
        a, b, c = toy_model(seed=3)(idx), toy_model(seed=3)(idx), toy_model(seed=4)(idx)
    assert torch.equal(a, b)
    assert not torch.allclose(a, c, atol=1e-4)


def test_preset_size(capsys: pytest.CaptureFixture[str]) -> None:
    counts = {
        scheme: PositionEncoder(EncoderConfig(input=scheme)).num_params(non_embedding=False)
        for scheme in ("moves", "squares")
    }
    with capsys.disabled():
        print()
        for scheme, total in counts.items():
            print(f"encoder preset ({scheme:<7}) {total:>12,} parameters")
    # 8 layers of d=384 are 14,195,712 parameters, plus 779,520 of embedding (2 030 moves) and
    # 76,800 of position table: 15,052,800. The plan's "18M-25M" band was an estimate its own
    # preset (n_layer=8, n_head=6, d_model=384) cannot reach; the preset is what is binding.
    assert 14_000_000 <= counts["moves"] <= 16_000_000
    assert counts["squares"] < counts["moves"]  # 47 square tokens against 2 030 moves

def test_config_rejects_unknown_keys_and_bad_shapes() -> None:
    with pytest.raises(ValueError):
        EncoderConfig(n_layers=3)  # type: ignore[call-arg]
    with pytest.raises(ValueError, match="not divisible"):
        EncoderConfig(d_model=100, n_head=8)
    with pytest.raises(ValueError, match="even head dimension"):
        EncoderConfig(d_model=12, n_head=4, pos="rope")
    with pytest.raises(ValueError, match="dropout"):
        EncoderConfig(dropout=1.0)
    with pytest.raises(ValueError, match="block"):
        EncoderConfig(block=0)
    assert EncoderConfig(d_model=64, n_head=4).ff == 256
    assert EncoderConfig(d_model=64, n_head=4, d_ff=128).ff == 128


def test_a_few_steps_of_gradient_descent_reduce_the_loss() -> None:
    torch.manual_seed(5)
    model = PositionEncoder(TOY)
    idx = toy_batch(seq=8, seed=6)
    labels = torch.full_like(idx, MMM_IGNORE_INDEX)
    labels[:, ::2] = idx[:, ::2]
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)
    losses = []
    for _ in range(10):
        _, loss = model.masked_lm(idx, labels)
        assert loss is not None
        losses.append(loss.detach().item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    assert losses[-1] < losses[0]