tests/test_phase0_params.py

"""Phase 0 — vérifie que les tailles de modèle annoncées dans les docs sont vraies.

Un chiffre de paramètres écrit dans la documentation peut dériver : ce test
recalcule les valeurs à partir de `configs/model.yaml` et échoue en cas d'écart.
"""

from pathlib import Path

import pytest

yaml = pytest.importorskip("yaml")

ROOT = Path(__file__).resolve().parents[1]


@pytest.fixture(scope="module")
def cfg():
    return yaml.safe_load((ROOT / "configs" / "model.yaml").read_text(encoding="utf-8"))


def count_params(vocab: int, n_embd: int, n_layer: int, tie: bool, ffn_mult: int = 4) -> int:
    """Compte les paramètres d'un transformer decoder-only pre-norm sans biais."""
    embed = vocab * n_embd
    per_layer = 4 * n_embd * n_embd + 2 * n_embd * (ffn_mult * n_embd) + 2 * n_embd
    total = embed + n_layer * per_layer + n_embd
    if not tie:
        total += vocab * n_embd
    return total


def test_default_profile_matches_documented_size(cfg):
    tm, tk = cfg["text_model"], cfg["tokenizer"]
    n = count_params(tk["vocab_size"], tm["n_embd"], tm["n_layer"], tm["tie_embeddings"])
    assert round(n / 1e6, 1) == 22.9, f"profil défaut = {n/1e6:.1f}M, attendu 22.9M"


def test_50m_profile_is_actually_about_50m(cfg):
    tm, tk = cfg["text_model_50m"], cfg["tokenizer"]
    n = count_params(tk["vocab_size"], tm["n_embd"], tm["n_layer"], tm["tie_embeddings"])
    assert 48e6 <= n <= 53e6, f"profil 50m = {n/1e6:.1f}M, hors de la cible ~50M"


def test_docs_state_the_real_default_size(cfg):
    arch = (ROOT / "docs" / "architecture.md").read_text(encoding="utf-8")
    assert "22.9 M" in arch, "architecture.md doit annoncer la vraie taille du profil défaut"
    assert "51.6 M" in arch, "architecture.md doit annoncer la vraie taille du profil 50M"