tests/test_tokenizer.py

"""Phase 1 — tests unitaires du tokenizer BPE de Cortex-0.

Rapides (CPU, quelques secondes) : un **tokenizer jouet** est entraîné une fois
pour toute la session sur un petit corpus français embarqué. Le tokenizer final
(32k, corpus ~200 Mo) est produit par `notebooks/01_tokenizer.ipynb` sur Colab.

Deux catégories de tests :
  * invariants qui ne dépendent PAS du corpus (tokens spéciaux, round-trip,
    accents, emojis, scripts inconnus, byte_fallback) → testés sur le tokenizer jouet ;
  * invariants du tokenizer FINAL (vocab == 32000) → ignorés si l'artefact réel
    n'est pas présent (avant le run Colab).
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from src.tokenizer import load_tokenizer
from src.tokenizer.special_tokens import (
    ALL_SPECIAL_TOKENS,
    BASE_SPECIAL_TOKENS,
    MODALITY_SPECIAL_TOKENS,
    N_SPECIAL_TOKENS,
    REASONING_SPECIAL_TOKENS,
    SPECIAL_TOKEN_IDS,
    SPECIAL_TOKENS_MAP_FILE,
    TOKENIZER_CONFIG_FILE,
    TOKENIZER_DIR,
    VOCAB_SIZE,
)
from src.tokenizer.train_tokenizer import (
    assert_special_token_ids,
    save_artifacts,
    train_tokenizer,
)

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

# 20 phrases françaises variées (accents, ponctuation, majuscules, chiffres).
FRENCH_SENTENCES = [
    "Bonjour, comment allez-vous aujourd'hui ?",
    "Le café coûte trois euros à Paris.",
    "L'élève écrit une phrase sur le tableau noir.",
    "Où est la forêt de Brocéliande ?",
    "Le garçon mange une pomme rouge et juteuse.",
    "Les enfants jouent dans le jardin après l'école.",
    "Il fait très beau à Montréal en été.",
    "La bibliothèque ouvre à neuf heures précises.",
    "Ce problème de mathématiques est vraiment difficile.",
    "Le chat dort paisiblement sur le canapé du salon.",
    "J'ai acheté deux baguettes et un croissant.",
    "La France compte de nombreuses régions pittoresques.",
    "Elle répond toujours avec beaucoup de patience.",
    "N'oubliez pas de fermer la porte à clé !",
    "Le train pour Lyon part dans quinze minutes.",
    "Nous irons à la mer dès que possible.",
    "Cette recette nécessite un four préchauffé.",
    "La langue française est riche en nuances.",
    "Il pleut des cordes depuis ce matin.",
    "Merci beaucoup pour votre aide précieuse.",
]

TOY_CORPUS = FRENCH_SENTENCES * 60 + [
    "Le soleil brille, les oiseaux chantent, tout va bien.",
    "Une phrase française supplémentaire pour l'entraînement.",
]

ACCENTS = "é è ê à ç ù î ô ë ï ü ÿ œ æ"
EMOJIS = "🦾 🇬🇳 🚀 🙂 ❤️ 🎉"
UNKNOWN_SCRIPTS = "こんにちは世界 مرحبا بالعالم Здравствуй мир"


@pytest.fixture(scope="session")
def toy_dir(tmp_path_factory) -> Path:
    """Entraîne un tokenizer jouet (vocab ~1500) une fois par session."""
    out = tmp_path_factory.mktemp("toy_tokenizer")
    tok = train_tokenizer(TOY_CORPUS, vocab_size=1500, min_frequency=1, show_progress=False)
    save_artifacts(tok, out)
    return out


@pytest.fixture(scope="session")
def toy(toy_dir):
    return load_tokenizer(toy_dir)


# ---------------------------------------------------------------------------
# Source de vérité des tokens spéciaux
# ---------------------------------------------------------------------------


def test_special_token_lists_are_disjoint_and_unique():
    assert len(set(ALL_SPECIAL_TOKENS)) == N_SPECIAL_TOKENS
    assert set(BASE_SPECIAL_TOKENS).isdisjoint(REASONING_SPECIAL_TOKENS)
    assert set(BASE_SPECIAL_TOKENS).isdisjoint(MODALITY_SPECIAL_TOKENS)
    assert set(REASONING_SPECIAL_TOKENS).isdisjoint(MODALITY_SPECIAL_TOKENS)


def test_special_token_counts():
    assert len(BASE_SPECIAL_TOKENS) == 4
    assert len(REASONING_SPECIAL_TOKENS) == 6
    assert len(MODALITY_SPECIAL_TOKENS) == 9
    assert N_SPECIAL_TOKENS == 19


def test_special_token_ids_are_contiguous():
    assert sorted(SPECIAL_TOKEN_IDS.values()) == list(range(N_SPECIAL_TOKENS))


def test_required_special_tokens_present():
    required = [
        "<|pad|>", "<|bos|>", "<|eos|>", "<|unk|>",
        "<|think|>", "<|plan|>", "<|tool_call|>", "<|tool_result|>", "<|critic|>", "<|moe_route|>",
        "<|image_start|>", "<|image_end|>", "<|image_pad|>",
        "<|audio_start|>", "<|audio_end|>", "<|audio_pad|>",
        "<|video_start|>", "<|video_end|>", "<|video_pad|>",
    ]
    for token in required:
        assert token in SPECIAL_TOKEN_IDS, f"token spécial manquant : {token}"


# ---------------------------------------------------------------------------
# Artefacts sur disque
# ---------------------------------------------------------------------------


def test_three_artifacts_written(toy_dir):
    for name in ("tokenizer.json", SPECIAL_TOKENS_MAP_FILE, TOKENIZER_CONFIG_FILE):
        assert (toy_dir / name).is_file(), f"{name} absent"


def test_special_tokens_map_contains_all_required(toy_dir):
    data = json.loads((toy_dir / SPECIAL_TOKENS_MAP_FILE).read_text(encoding="utf-8"))
    present = set(data.get("additional_special_tokens", []))
    present |= {data.get("pad_token"), data.get("bos_token"), data.get("eos_token"), data.get("unk_token")}
    for token in ALL_SPECIAL_TOKENS:
        assert token in present, f"{token} absent de special_tokens_map.json"


def test_special_token_ids_match_expected(toy):
    assert_special_token_ids(toy.raw)
    for token, expected in SPECIAL_TOKEN_IDS.items():
        assert toy.token_to_id(token) == expected, token


def test_special_token_ids_stable_across_two_runs(tmp_path):
    a = train_tokenizer(TOY_CORPUS, vocab_size=1500, min_frequency=1, show_progress=False)
    b = train_tokenizer(TOY_CORPUS, vocab_size=1500, min_frequency=1, show_progress=False)
    assert a.get_vocab() == b.get_vocab()


# ---------------------------------------------------------------------------
# Round-trip et cas limites
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("sentence", FRENCH_SENTENCES)
def test_roundtrip_french_sentences(toy, sentence):
    assert toy.decode(toy.encode(sentence)) == sentence


def test_roundtrip_accents(toy):
    assert toy.decode(toy.encode(ACCENTS)) == ACCENTS


def test_roundtrip_emojis(toy):
    assert toy.decode(toy.encode(EMOJIS)) == EMOJIS


def test_roundtrip_unknown_scripts(toy):
    assert toy.decode(toy.encode(UNKNOWN_SCRIPTS)) == UNKNOWN_SCRIPTS


def test_roundtrip_empty_and_whitespace(toy):
    assert toy.decode(toy.encode("")) == ""
    assert toy.decode(toy.encode("   ")) == "   "


def test_roundtrip_very_long(toy):
    long_text = " ".join(FRENCH_SENTENCES * 20)
    assert toy.decode(toy.encode(long_text)) == long_text


def test_special_tokens_recognized_when_in_text(toy):
    text = "Bonjour <|bos|> le monde <|eos|>"
    ids = toy.encode(text)
    assert toy.token_to_id("<|bos|>") in ids
    assert toy.token_to_id("<|eos|>") in ids


def test_encode_returns_ints(toy):
    ids = toy.encode("Un test.")
    assert ids and all(isinstance(i, int) for i in ids)
    assert all(0 <= i < toy.vocab_size for i in ids)


def test_decode_skips_special_tokens(toy):
    ids = [toy.token_to_id("<|bos|>")] + toy.encode("salut") + [toy.token_to_id("<|eos|>")]
    decoded = toy.decode(ids, skip_special_tokens=True)
    assert "<|bos|>" not in decoded and "<|eos|>" not in decoded
    assert "salut" in decoded


def test_byte_fallback_is_enabled(toy):
    # Un caractère jamais vu (CJK) doit survivre au round-trip -> repli octet.
    assert toy.decode(toy.encode("漢字第")) == "漢字第"


# ---------------------------------------------------------------------------
# Entraînement
# ---------------------------------------------------------------------------


def test_trainer_respects_vocab_size_cap():
    tok = train_tokenizer(TOY_CORPUS, vocab_size=800, min_frequency=1, show_progress=False)
    assert tok.get_vocab_size() <= 800


def test_min_frequency_reduces_vocab():
    low = train_tokenizer(TOY_CORPUS, vocab_size=2000, min_frequency=1, show_progress=False)
    high = train_tokenizer(TOY_CORPUS, vocab_size=2000, min_frequency=20, show_progress=False)
    assert high.get_vocab_size() <= low.get_vocab_size()


# ---------------------------------------------------------------------------
# Invariants du tokenizer FINAL (nécessite le run Colab)
# ---------------------------------------------------------------------------


@pytest.mark.skipif(
    not (ROOT / TOKENIZER_DIR / "tokenizer.json").is_file(),
    reason="tokenizer final absent : généré par notebooks/01_tokenizer.ipynb (Colab)",
)
def test_final_tokenizer_vocab_size_is_exactly_32000():
    final = load_tokenizer()
    assert final.vocab_size == VOCAB_SIZE
    for token, expected in SPECIAL_TOKEN_IDS.items():
        assert final.token_to_id(token) == expected