tests/test_gemma4_e2b_token_loop.py

"""
test_gemma4_e2b_token_loop.py — Token-Loop Regression for Gemma 4 E2B
=====================================================================
Detects the "mathematics\\nmathematics\\nmathematics..." failure mode
from the 2026-06-08 E2B test (repetition collapse on narrow kurtosis).

This is a STRUCTURED unit test: the actual generation requires GPU,
so we test the *mechanisms* that prevent the loop:
  1. repetition_penalty in _px_config (= 1.15, set unconditionally for all models)
  2. Stability-break threshold (stability_cnt > 3, gated on phi_s > 0.9999)
  3. Anti-oscillation guards (per-layer-visit penalty + over-stable hub-jump)
  4. Generative output checks (when run with GPU)

Note (2026-06-18): the earlier `stability_cnt > 1` threshold and the dedicated
BOUNCE-BREAK block were removed in the 2026-06-11 patch.py rewrite. Token-loop
avoidance now comes from repetition_penalty=1.15 + no_repeat_ngram_size=3 plus
the is_gemma4 n_loops cap (8→4). The tests below were updated accordingly.

The GPU-dependent `test_token_loop_in_generation` is skipped if no
GPU is available — all other tests run on CPU only.

Usage:
  PYTHONPATH=. python -m pytest tests/test_gemma4_e2b_token_loop.py -v
  (with GPU) PYTHONPATH=. python tests/test_gemma4_e2b_token_loop.py
"""

import os
import sys
import unittest
from unittest.mock import MagicMock

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


# ═══════════════════════════════════════════════════════════════════════════════
# TestGemma4E2BTokenLoopMechanisms — verify the 4 mitigations
# ═══════════════════════════════════════════════════════════════════════════════

class TestGemma4E2BTokenLoopMechanisms(unittest.TestCase):
    """Each mechanism that prevents token loops is verified independently."""

    def test_repetition_penalty_in_config(self):
        """Mechanism 1: repetition_penalty=1.15 (set unconditionally in
        apply_px_patch for ALL models, not just gemma4)."""
        import inspect
        from px_patches.gemma4_2b_px import patch as patch_mod
        source = inspect.getsource(patch_mod)
        # Production sets repetition_penalty=1.15 unconditionally (patch.py:1008).
        self.assertIn('defaults["repetition_penalty"] = 1.15', source,
                      "patch.py must set repetition_penalty=1.15")
        self.assertIn('defaults["no_repeat_ngram_size"] = 3', source,
                      "patch.py must set no_repeat_ngram_size=3")

    def test_stability_break_threshold(self):
        """Mechanism 2: stability_cnt > 3 (gated on phi_s > 0.9999) triggers
        break, plus a tail safety guard stability_cnt > 5. The earlier > 1
        threshold was removed in the 2026-06-11 rewrite."""
        import inspect
        from px_patches.gemma4_2b_px import patch as patch_mod
        source = inspect.getsource(patch_mod)

        # Primary stability break, gated on near-perfect cosine stability:
        self.assertIn("if stability_cnt > 3:", source,
                      "stability break must fire at `stability_cnt > 3`")
        self.assertIn("phi_s > 0.9999", source,
                      "stability break must be gated on phi_s > 0.9999")
        # Tail-of-iteration safety guard:
        self.assertIn("if stability_cnt > 5: break", source,
                      "tail safety guard `stability_cnt > 5` must be present")
        # The old `> 1` one-line break must be gone:
        self.assertNotIn("stability_cnt > 1: h_exp = trans_out; break", source,
                         "Old `stability_cnt > 1` one-line break must be removed")

    def test_layer_anti_oscillation_guards_present(self):
        """Mechanism 3: the dedicated BOUNCE-BREAK block was removed
        (2026-06-11). Anti-oscillation is now handled by lighter in-loop
        guards: per-layer-visit penalty, over-stable hub-jump, and
        layer-visits bookkeeping, with per-step telemetry collection."""
        import inspect
        from px_patches.gemma4_2b_px import patch as patch_mod
        source = inspect.getsource(patch_mod)
        self.assertIn("layer_visits", source,
                      "layer-visits bookkeeping must be present")
        self.assertIn("dynamic_hub + 1", source,
                      "over-stable hub-jump guard must be present")
        self.assertIn("_px_current_telemetry_raw", source,
                      "per-step telemetry collection must be present")
        # The removed BOUNCE-BREAK marker must NOT have come back:
        self.assertNotIn("BOUNCE-BREAK", source,
                          "BOUNCE-BREAK block was removed 2026-06-11; "
                          "must not be reintroduced")

    def test_reduced_recur_range(self):
        """Mechanism 4: gemma4-e2b uses SCALE_DEFAULTS[1536] for 1B parity.
        Token-loop avoidance comes from repetition_penalty=1.15 +
        no_repeat_ngram_size=3 (set in patch.py) plus the is_gemma4 n_loops
        cap (8→4) inside apply_px_patch — not from a low gamma."""
        from config import MODEL_REGISTRY
        from px_patches.gemma4_2b_px.auto_tune import SCALE_DEFAULTS
        reg = MODEL_REGISTRY["gemma4-e2b-it"]
        kw = reg["patch_kwargs"]
        # No explicit override — defaults come from SCALE_DEFAULTS[1536]
        self.assertNotIn("recur_start", kw)
        self.assertNotIn("n_loops", kw)
        # SCALE_DEFAULTS[1536] must exist and have the 1B-parity values
        self.assertIn(1536, SCALE_DEFAULTS)
        sd = SCALE_DEFAULTS[1536]
        self.assertEqual(sd["recur_start"], 10)
        self.assertEqual(sd["recur_end"], 26)
        self.assertEqual(sd["n_loops"], 8)
        self.assertEqual(sd["gamma"], 0.12)


# ═══════════════════════════════════════════════════════════════════════════════
# TestGemma4E2BTokenRepetition — structural repetition detection helpers
# ═══════════════════════════════════════════════════════════════════════════════

class TestGemma4E2BTokenRepetition(unittest.TestCase):
    """Helper functions for detecting repetitive token patterns."""

    def _count_consecutive_repeats(self, tokens):
        """Return the max run of identical consecutive tokens."""
        if not tokens:
            return 0
        max_run = cur_run = 1
        for i in range(1, len(tokens)):
            if tokens[i] == tokens[i - 1]:
                cur_run += 1
                max_run = max(max_run, cur_run)
            else:
                cur_run = 1
        return max_run

    def _token_type_ratio(self, tokens):
        """Unique tokens / total tokens (1.0 = all unique, low = repetitive)."""
        if not tokens:
            return 1.0
        return len(set(tokens)) / len(tokens)

    def test_detects_three_consecutive_repeats(self):
        # The bug pattern: "math", "math", "math", ...
        tokens = ["the", "math", "math", "math", "induction"]
        repeats = self._count_consecutive_repeats(tokens)
        self.assertEqual(repeats, 3,
                         "Three consecutive 'math' tokens must be detected")

    def test_detects_long_run(self):
        # The actual 2026-06-08 output pattern
        tokens = ["mathematics"] * 5 + ["."]
        repeats = self._count_consecutive_repeats(tokens)
        self.assertEqual(repeats, 5)

    def test_token_type_ratio_low_for_loop(self):
        tokens = ["mathematics"] * 10
        ratio = self._token_type_ratio(tokens)
        self.assertLess(ratio, 0.2,
                        "Pure repetition must yield very low type ratio")

    def test_token_type_ratio_high_for_diverse(self):
        tokens = ["The", "concept", "of", "mathematical", "induction", "is", "fundamental", "."]
        ratio = self._token_type_ratio(tokens)
        self.assertGreater(ratio, 0.8,
                          "Diverse tokens must yield high type ratio")

    def test_no_repetition_in_normal_text(self):
        tokens = "Explain the concept of mathematical induction.".split()
        repeats = self._count_consecutive_repeats(tokens)
        self.assertEqual(repeats, 1, "Normal text has no consecutive repeats")


# ═══════════════════════════════════════════════════════════════════════════════
# TestGemma4E2BLoopPattern — the 2026-06-08 actual failure pattern
# ═══════════════════════════════════════════════════════════════════════════════

class TestGemma4E2BLoopPattern(unittest.TestCase):
    """Reproduces the exact failure pattern from 2026-06-08 and checks fixes."""

    def test_2026_06_08_failure_pattern(self):
        # The actual output: "Explain the concept of mathematical induction.\n
        #                    Explain what is the concept of mathematical.\n
        #                    Explain.\n  Explain why\n  the relationship between
        #                    the concept of mathematical induction.\n  Explain why
        #                    mathematics\n  mathematics\n  mathematics"
        response = (
            "Explain the concept of mathematical induction.\n"
            "Explain what is the concept of mathematical.\n"
            "Explain.\n  Explain why\n  the relationship between the concept of mathematical induction.\n"
            "  Explain why\n  mathematics\n  mathematics\n  mathematics"
        )
        # Word-level tokens
        words = response.split()
        # Count "mathematics" runs
        run_lengths = []
        i = 0
        while i < len(words):
            j = i
            while j < len(words) and words[j] == words[i]:
                j += 1
            run_lengths.append(j - i)
            i = j
        max_run = max(run_lengths) if run_lengths else 0
        # The failure: 3 consecutive "mathematics" at the end
        self.assertEqual(max_run, 3,
                         "The 2026-06-08 pattern has a 3-token run of 'mathematics'")

    def test_loop_should_have_been_broken(self):
        """Verify that with the current threshold (stability_cnt > 3, gated on
        phi > 0.9999), a sustained high-phi loop is eventually broken."""
        # Patch.py threshold: phi_s > 0.9999 sustained for stability_cnt > 3
        simulated_history = [0.99995, 0.99995, 0.99995, 0.99995, 0.99995]  # 5 high-phi steps
        stability_cnt = 0
        broke = False
        for phi in simulated_history:
            if phi > 0.9999:
                stability_cnt += 1
                if stability_cnt > 3:  # CURRENT THRESHOLD (was > 1 pre-2026-06-11)
                    broke = True
                    break
            else:
                stability_cnt = 0
        self.assertTrue(broke,
                        "Stability-break must fire once stability_cnt exceeds 3")
        self.assertEqual(stability_cnt, 4,
                         "stability_cnt must reach 4 to trigger break (> 3)")


# ═══════════════════════════════════════════════════════════════════════════════
# TestGemma4E2BGeneratedOutput — optional GPU test (skipped if no GPU)
# ═══════════════════════════════════════════════════════════════════════════════

class TestGemma4E2BGeneratedOutput(unittest.TestCase):
    """GPU test: actually generate text and check for loops. Skipped without GPU."""

    @unittest.skipUnless(
        torch.cuda.is_available() and os.environ.get("RUN_GPU_TESTS") == "1",
        "GPU not available or RUN_GPU_TESTS not set"
    )
    def test_token_loop_in_generation(self):
        """Actually generate from gemma4-e2b-it and verify no loops."""
        from model_manager import ModelManager
        manager = ModelManager()
        import asyncio
        loop = asyncio.new_event_loop()
        try:
            entry = loop.run_until_complete(
                manager.get_model("gemma4-e2b-it", px_subjective=True,
                                  px_config_preset="SUBJECTIVE")
            )
            model = entry["model"]
            tokenizer = entry["tokenizer"]
            prompt = "Explain the concept of mathematical induction."
            messages = [{"role": "user", "content": prompt}]
            input_text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

            with torch.no_grad():
                output_ids = model.generate(
                    **inputs, max_new_tokens=50, do_sample=True,
                    temperature=0.7, top_p=0.9,
                )

            generated = tokenizer.decode(
                output_ids[0][inputs["input_ids"].shape[1]:],
                skip_special_tokens=True
            )
            tokens = generated.split()

            # After token-loop fix: max 2 consecutive repeats
            max_run = 1
            cur = 1
            for i in range(1, len(tokens)):
                if tokens[i] == tokens[i - 1]:
                    cur += 1
                    max_run = max(max_run, cur)
                else:
                    cur = 1

            unique_ratio = len(set(tokens)) / max(len(tokens), 1)
            print(f"\nGenerated ({len(tokens)} tokens): {generated[:200]}")
            print(f"Max consecutive repeat: {max_run}, unique ratio: {unique_ratio:.2f}")

            # The 2026-06-08 fix should prevent long runs
            self.assertLessEqual(max_run, 2,
                                 f"Max consecutive repeat must be ≤ 2, got {max_run}")
            self.assertGreater(unique_ratio, 0.3,
                               f"Token type ratio must be > 0.3, got {unique_ratio:.2f}")
        finally:
            loop.close()
            try:
                loop.run_until_complete(manager.unload_model("gemma4-e2b-it"))
            except Exception:
                pass


if __name__ == "__main__":
    unittest.main(verbosity=2)