tests/test_speak.py

# tests/test_speak.py
"""Plan 6, tâche 1 — speak() : la boucle de décodage mécanique anti-attracteur.

Transforms (znorm, rep-pen, top-k, top-p) vérifiés sur des tenseurs CONNUS ;
puis la boucle speak() sur un tiny CteCore réel (et un tiny FractusCore pour
le duck-typing) : déterminisme par seed, anti-répétition à rep_pen élevé
(aucun token émis 3 fois sur 12 pas), domination du bias fort (δ=+50 →
premier choix, 10 seeds), EOS append-then-break strippé au décodage.
"""
from collections import Counter

import pytest

torch = pytest.importorskip("torch", reason="speak nécessite torch (substrat sans torch)")

from fractus_vorax.model.cte_core import CteCore, CteCoreConfig
from fractus_vorax.model.speak import (
    apply_rep_penalty,
    apply_top_k,
    apply_top_p,
    speak,
    znorm_logits,
)


def _tiny(**overrides) -> CteCoreConfig:
    return CteCoreConfig(vocab_size=128, **overrides)


class _TinyCharTok:
    """Tokenizer char-level sur le vocab tiny (128) — injectif sur les ids < 128."""

    vocab_size = 128

    def encode(self, text: str) -> list[int]:
        return [ord(c) % 128 for c in text]

    def decode(self, ids: list[int]) -> str:
        return "".join(chr(int(i) % 128) for i in ids)


# ---------------------------------------------------------------------------
# Transforms — tenseurs connus
# ---------------------------------------------------------------------------

def test_znorm_known_tensor():
    x = torch.tensor([1.0, 2.0, 3.0])  # mean 2, std non biaisé 1
    assert torch.allclose(znorm_logits(x, scale=3.0), torch.tensor([-3.0, 0.0, 3.0]))
    # défaut = 3.0 (l'échelle empirique qui libère le lexique sur le 1B)
    assert torch.allclose(znorm_logits(x), znorm_logits(x, scale=3.0))
    assert torch.allclose(znorm_logits(x, scale=1.0), torch.tensor([-1.0, 0.0, 1.0]))
    # pur : l'entrée n'est jamais mutée
    assert torch.equal(x, torch.tensor([1.0, 2.0, 3.0]))


def test_znorm_crushes_attractor_scale():
    """Propriété clé : std de sortie == scale, quelle que soit l'échelle
    d'entrée (état calme std ~26 ou état-attracteur à des centaines)."""
    calm = torch.randn(2048) * 26.0
    attractor = torch.randn(2048) * 400.0
    for src in (calm, attractor):
        out = znorm_logits(src, scale=3.0)
        assert abs(float(out.std()) - 3.0) < 1e-3
        assert abs(float(out.mean())) < 1e-3


def test_apply_rep_penalty_known_tensor():
    lg = torch.tensor([6.0, 3.0, -3.0])
    out = apply_rep_penalty(lg, {0}, penalty=6.0)
    assert torch.allclose(out, torch.tensor([1.0, 3.0, -3.0]))
    out2 = apply_rep_penalty(lg, {0, 1}, penalty=3.0)
    assert torch.allclose(out2, torch.tensor([2.0, 1.0, -3.0]))
    # division SIMPLE (spec du plan) : un logit négatif pénalisé se rapproche
    # de 0 — comportement documenté, sans effet pratique (les tokens émis post-
    # znorm vivent dans le haut de la distribution, positifs)
    out3 = apply_rep_penalty(lg, {2}, penalty=6.0)
    assert torch.allclose(out3, torch.tensor([6.0, 3.0, -0.5]))
    # pur ; emitted vide → inchangé
    assert torch.equal(lg, torch.tensor([6.0, 3.0, -3.0]))
    assert torch.equal(apply_rep_penalty(lg, set()), lg)
    with pytest.raises(ValueError):
        apply_rep_penalty(lg, {0}, penalty=0.0)


def test_apply_top_k_known_tensor():
    lg = torch.tensor([1.0, 5.0, 3.0, 2.0])
    out = apply_top_k(lg, 2)
    expected = torch.tensor([float("-inf"), 5.0, 3.0, float("-inf")])
    assert torch.equal(out, expected)
    # k = taille du vocab → tout survit (identité)
    assert torch.equal(apply_top_k(lg, 4), lg)
    with pytest.raises(ValueError):
        apply_top_k(lg, 0)


def test_apply_top_p_known_tensor():
    lg = torch.log(torch.tensor([0.6, 0.3, 0.1]))
    # p=0.85 : cumul EXCLUSIF = [0, 0.6, 0.9] → les 2 premiers survivent
    out = apply_top_p(lg, 0.85)
    assert torch.isfinite(out[0]) and torch.isfinite(out[1])
    assert out[2] == float("-inf")
    # p minuscule : le top-1 SURVIT TOUJOURS (jamais de masque vide)
    out1 = apply_top_p(lg, 0.01)
    assert torch.isfinite(out1[0])
    assert out1[1] == float("-inf") and out1[2] == float("-inf")
    # p >= 1 → identité
    assert torch.equal(apply_top_p(lg, 1.0), lg)
    with pytest.raises(ValueError):
        apply_top_p(lg, 0.0)


# ---------------------------------------------------------------------------
# speak() — tiny CteCore réel
# ---------------------------------------------------------------------------

def test_speak_deterministic_same_seed():
    torch.manual_seed(0)
    core = CteCore(_tiny())
    tok = _TinyCharTok()
    text_a, diags_a = speak(core, tok, "hello", max_new_tokens=8, seed=42)
    text_b, diags_b = speak(core, tok, "hello", max_new_tokens=8, seed=42)
    assert text_a == text_b
    # rejeu EXACT y compris diagnostics : reset_states(1) au départ rend la
    # boucle indépendante des états porteurs laissés par l'appel précédent
    assert diags_a == diags_b
    assert isinstance(text_a, str) and text_a != ""
    # structure des diagnostics par pas
    assert 1 <= len(diags_a) <= 8
    d = diags_a[0]
    assert set(d.keys()) == {"logit_std_avant", "top3", "chosen"}
    assert isinstance(d["logit_std_avant"], float) and d["logit_std_avant"] > 0.0
    assert len(d["top3"]) == 3
    assert all(isinstance(s, str) and isinstance(v, float) for s, v in d["top3"])
    assert isinstance(d["chosen"], str)


def test_speak_high_rep_pen_never_triple():
    """rep_pen élevé (12, haut de la fourchette empirique 6-12) : aucun token
    n'est émis 3 fois sur 12 pas — l'attracteur de répétition est brisé."""
    torch.manual_seed(0)
    core = CteCore(_tiny())
    tok = _TinyCharTok()
    text, diags = speak(
        core, tok, "the meaning of", max_new_tokens=12, rep_pen=12.0, seed=0
    )
    assert len(diags) == 12
    counts = Counter(d["chosen"] for d in diags)
    assert max(counts.values()) <= 2, f"attracteur non brisé: {counts}"


def test_speak_strong_bias_dominates_first_pick():
    """δ=+50 sur un id : le premier choix EST cet id, 10 seeds sur 10."""
    torch.manual_seed(0)
    core = CteCore(_tiny())
    tok = _TinyCharTok()
    for seed in range(10):
        text, diags = speak(
            core, tok, "hello", max_new_tokens=1, bias_tokens={42: 50.0}, seed=seed
        )
        assert diags[0]["chosen"] == tok.decode([42])


def test_speak_eos_append_then_break_stripped():
    """Le bias force 42 = EOS au premier pas → append-then-break : 1 pas de
    diagnostic, EOS strippé au décodage → texte vide (sémantique P4)."""
    torch.manual_seed(0)
    core = CteCore(_tiny())
    tok = _TinyCharTok()
    text, diags = speak(
        core, tok, "hello", max_new_tokens=6,
        eos_token_id=42, bias_tokens={42: 50.0}, seed=0,
    )
    assert len(diags) == 1
    assert text == ""


def test_speak_eos_from_tokenizer():
    """Le tokenizer peut exposer eos_token_id (Gpt2BpeTokenizer le fait) :
    sans paramètre explicite, c'est lui qui clot la boucle."""

    class _TokWithEos(_TinyCharTok):
        eos_token_id = 42

    torch.manual_seed(0)
    core = CteCore(_tiny())
    text, diags = speak(
        core, _TokWithEos(), "hi", max_new_tokens=6, bias_tokens={42: 50.0}, seed=0
    )
    assert len(diags) == 1
    assert text == ""


def test_speak_rejects_bad_arguments():
    torch.manual_seed(0)
    core = CteCore(_tiny())
    tok = _TinyCharTok()
    with pytest.raises(ValueError):
        speak(core, tok, "", max_new_tokens=4)  # prompt vide → aucun token
    with pytest.raises(ValueError):
        speak(core, tok, "hello", temperature=0.0)  # greedy = core.generate
    with pytest.raises(ValueError):
        speak(core, tok, "hello", bias_tokens={999: 1.0})  # id hors vocab


def test_speak_duck_types_fractus_core():
    """Le contrat est le duck-typing : FractusCore (stateless, sans
    reset_states) parle aussi, de façon déterministe."""
    from fractus_vorax.model.fractus_core import FractusCore, FractusCoreConfig

    torch.manual_seed(0)
    cfg = FractusCoreConfig(
        vocab_size=128, d_model=64, n_layers=2, n_experts=4, top_k=2,
        rank=16, max_seq_len=64,
    )
    core = FractusCore(cfg)
    tok = _TinyCharTok()
    text_a, diags_a = speak(core, tok, "hello", max_new_tokens=6, seed=3)
    text_b, diags_b = speak(core, tok, "hello", max_new_tokens=6, seed=3)
    assert text_a == text_b and diags_a == diags_b
    assert len(diags_a) == 6
    assert all(set(d.keys()) == {"logit_std_avant", "top3", "chosen"} for d in diags_a)