feat(echo-trap): add tokenizer-aware repetition scoring (#242)

Closes #241.

Adds opt-in tokenizer-aware n-gram path for the v0.70.0 Part F echo-trap detector. The existing whitespace `score_echo_signal` is unchanged; callers opt in via the new `score_trajectory_repetition_tokenized` / `score_echo_signal_tokenized` helpers or the `--echo-trap-tokenizer-aware` train flag.

Acceptance criterion from #241 verified: synthetic case where decoded strings differ by punctuation but token-id sequence repeats — whitespace path returns OK, tokenizer-aware path returns TRAP.

README.md

@@ -131,9 +131,12 @@ soup iterative-dpo \
 soup train --config grpo.yaml \
     --echo-trap-enabled \
     --echo-trap-threshold 0.6 \
-    --echo-trap-halt
+    --echo-trap-halt \
+    --echo-trap-tokenizer-aware
 ```
 
+`--echo-trap-tokenizer-aware` switches echo-trap n-grams from whitespace tokens to the active tokenizer's integer ids. This catches subword repetition that punctuation-heavy decoded text can hide, but the score becomes tokenizer-specific rather than vocabulary-agnostic.
+
 Every detector composes with v0.34 `soup why` (anomaly explainer), v0.32 spike recovery, and v0.53.11 #127 `GRPOStabilityCallback` so a single training run can have InfoRM + echo-trap + spike-recovery + ref-model regen all active simultaneously without duplicating trajectory / state collection. Live trainer-callback wiring for all 6 Parts lands in v0.70.1 (`build_reward_hack_callback`, `build_uld_projection`, `build_minillm_callback`, `build_rl_checkpoint_callback`, `run_iterative_dpo`, `build_echo_trap_callback`); today every CLI / config flag is validated at schema-load so misconfigured runs fail loudly at config-load time.
 
 ## Why Soup?

soup_cli/commands/train.py

@@ -155,6 +155,14 @@ def train(
             "Default deny (v0.36.0). Only enable if you trust the source."
         ),
     ),
+    echo_trap_tokenizer_aware: bool = typer.Option(
+        False,
+        "--echo-trap-tokenizer-aware",
+        help=(
+            "Use tokenizer-id n-grams for echo-trap scoring. Requires "
+            "training.echo_trap_enabled=true on grpo/ppo."
+        ),
+    ),
     profile_run: bool = typer.Option(
         False,
         "--profile",
@@ -230,6 +238,17 @@ def train(
     console.print(f"[dim]Loading config from {config_path}...[/]")
     cfg = load_config(config_path)
 
+    # --- Echo-trap tokenizer-aware shortcut ---
+    if echo_trap_tokenizer_aware:
+        if not cfg.training.echo_trap_enabled:
+            console.print(
+                "[red]--echo-trap-tokenizer-aware requires "
+                "training.echo_trap_enabled=true[/]"
+            )
+            raise typer.Exit(1)
+        cfg.training.echo_trap_tokenizer_aware = True
+        console.print("[green]Echo-trap tokenizer-aware scoring enabled[/]")
+
     # --- --push-as / --hf-resume validation ---
     if push_as:
         from soup_cli.utils.hf import validate_repo_id
@@ -440,6 +459,8 @@ def train(
                     script_args.append("--wandb")
                 if tensorboard:
                     script_args.append("--tensorboard")
+                if echo_trap_tokenizer_aware:
+                    script_args.append("--echo-trap-tokenizer-aware")
                 if yes:
                     script_args.append("--yes")
                 argv = build_accelerate_argv(

soup_cli/config/schema.py

@@ -1366,6 +1366,15 @@ class TrainingConfig(BaseModel):
             "echo_trap_enabled=True. (v0.70.0)"
         ),
     )
+    echo_trap_tokenizer_aware: bool = Field(
+        default=False,
+        description=(
+            "Use tokenizer-id n-grams for echo-trap scoring instead of "
+            "whitespace tokens. More sensitive to subword repetition but "
+            "bound to the active tokenizer vocabulary. Requires "
+            "echo_trap_enabled=True. (v0.70.x)"
+        ),
+    )
 
     # ---- v0.70.0 Part D — Mid-epoch RL checkpoint ------------------------
     rl_checkpoint_save_every_steps: Optional[int] = Field(
@@ -1520,6 +1529,7 @@ class TrainingConfig(BaseModel):
     @field_validator(
         "echo_trap_enabled",
         "echo_trap_halt",
+        "echo_trap_tokenizer_aware",
         mode="before",
     )
     @classmethod
@@ -4035,15 +4045,23 @@ class SoupConfig(BaseModel):
         ``echo_trap_enabled`` is a silent no-op footgun — reject.
         """
         tcfg = self.training
-        if not tcfg.echo_trap_enabled and not tcfg.echo_trap_halt:
+        if (
+            not tcfg.echo_trap_enabled
+            and not tcfg.echo_trap_halt
+            and not tcfg.echo_trap_tokenizer_aware
+        ):
             return self
-        if not tcfg.echo_trap_enabled and tcfg.echo_trap_halt:
+        if not tcfg.echo_trap_enabled and (
+            tcfg.echo_trap_halt or tcfg.echo_trap_tokenizer_aware
+        ):
             raise ValueError(
-                "echo_trap_halt=True requires echo_trap_enabled=True"
+                "echo_trap_halt / echo_trap_tokenizer_aware require "
+                "echo_trap_enabled=True"
             )
         if self.task not in ("grpo", "ppo"):
             raise ValueError(
-                "echo_trap_enabled / echo_trap_halt are only valid on "
+                "echo_trap_enabled / echo_trap_halt / "
+                "echo_trap_tokenizer_aware are only valid on "
                 f"task in {{'grpo', 'ppo'}}; got task={self.task!r}"
             )
         if self.backend == "mlx":

soup_cli/utils/echo_trap.py

@@ -14,9 +14,8 @@ trajectory collection.
 Security:
 - Pure-Python math (no torch import at module top).
 - Bool / NaN / Inf / range rejection on every numeric input.
-- Tokens must be strings; non-str rejected loudly so a caller that
+- Whitespace-mode tokens must be strings; tokeniser-aware mode accepts
-  hands tensor ids (instead of decoded strings) gets an actionable
+  integer token ids through the dedicated ``*_tokenized`` helpers.
-  error rather than silently misbehaving.
 - ``_MAX_BATCH_TRAJECTORIES = 100_000`` DoS cap (matches v0.55 /
   v0.65 / v0.66 cap policy).
 - ``_MAX_NGRAM_N = 32`` keeps the n-gram counter bounded.
@@ -75,6 +74,41 @@ def _check_tokens(tokens: object) -> tuple[str, ...]:
     return tuple(iterator)
 
 
+def _check_token_ids(token_ids: object) -> tuple[int, ...]:
+    if isinstance(token_ids, (str, bytes)):
+        raise TypeError("token_ids must be a sequence of ints, not str/bytes")
+    try:
+        iterator = list(token_ids)  # type: ignore[arg-type]
+    except TypeError as exc:
+        raise TypeError(
+            f"token_ids must be iterable, got {type(token_ids).__name__}"
+        ) from exc
+    if len(iterator) > _MAX_TRAJECTORY_TOKENS:
+        raise ValueError(
+            f"trajectory has {len(iterator)} token ids, exceeds "
+            f"{_MAX_TRAJECTORY_TOKENS} cap"
+        )
+    for idx, token_id in enumerate(iterator):
+        if isinstance(token_id, bool) or not isinstance(token_id, int):
+            raise TypeError(
+                f"token_ids[{idx}] must be int, got {type(token_id).__name__}"
+            )
+    return tuple(iterator)
+
+
+def _score_repetition(units: Sequence[object], *, ngram_n: int) -> float:
+    if len(units) < ngram_n:
+        return 0.0
+    counts: dict[tuple[object, ...], int] = {}
+    for i in range(len(units) - ngram_n + 1):
+        gram = tuple(units[i : i + ngram_n])
+        counts[gram] = counts.get(gram, 0) + 1
+    if not counts:
+        return 0.0
+    repeating = sum(1 for count in counts.values() if count > 1)
+    return repeating / len(counts)
+
+
 def score_trajectory_repetition(tokens: object, *, ngram_n: object = 2) -> float:
     """Per-trajectory repetition score.
 
@@ -88,16 +122,24 @@ def score_trajectory_repetition(tokens: object, *, ngram_n: object = 2) -> float
     """
     n = _check_ngram_n(ngram_n)
     tok = _check_tokens(tokens)
-    if len(tok) < n:
+    return _score_repetition(tok, ngram_n=n)
-        return 0.0
+
-    counts: dict[tuple[str, ...], int] = {}
+
-    for i in range(len(tok) - n + 1):
+def score_trajectory_repetition_tokenized(
-        gram = tok[i : i + n]
+    token_ids: object,
-        counts[gram] = counts.get(gram, 0) + 1
+    *,
-    if not counts:
+    ngram_n: object = 2,
-        return 0.0
+) -> float:
-    repeating = sum(1 for c in counts.values() if c > 1)
+    """Per-trajectory repetition score over tokenizer ids.
-    return repeating / len(counts)
+
+    This mirrors :func:`score_trajectory_repetition`, but operates on
+    integer token ids before decoding/whitespace splitting can hide
+    subword repetition. It is intentionally separate so the existing
+    string-token API keeps rejecting accidental tensor-id input.
+    """
+    n = _check_ngram_n(ngram_n)
+    ids = _check_token_ids(token_ids)
+    return _score_repetition(ids, ngram_n=n)
 
 
 def score_echo_signal(
@@ -135,6 +177,41 @@ def score_echo_signal(
     return sum(scores) / len(scores)
 
 
+def score_echo_signal_tokenized(
+    trajectories: object,
+    *,
+    ngram_n: object = 2,
+) -> float:
+    """Mean repetition score across token-id trajectories.
+
+    Use this when the caller has access to the trainer tokenizer and can
+    pass ``tokenizer.encode(text)`` output rather than decoded strings.
+    """
+    n = _check_ngram_n(ngram_n)
+    if isinstance(trajectories, (str, bytes)):
+        raise TypeError(
+            "trajectories must be a sequence of token-id sequences, not str/bytes"
+        )
+    try:
+        batch = list(trajectories)  # type: ignore[arg-type]
+    except TypeError as exc:
+        raise TypeError(
+            f"trajectories must be iterable, got "
+            f"{type(trajectories).__name__}"
+        ) from exc
+    if len(batch) > _MAX_BATCH_TRAJECTORIES:
+        raise ValueError(
+            f"batch has {len(batch)} trajectories, exceeds "
+            f"{_MAX_BATCH_TRAJECTORIES} cap"
+        )
+    if not batch:
+        return 0.0
+    scores: list[float] = []
+    for traj in batch:
+        scores.append(score_trajectory_repetition_tokenized(traj, ngram_n=n))
+    return sum(scores) / len(scores)
+
+
 def classify_echo_signal(signal: object) -> str:
     """Map a signal in ``[0, 1]`` to OK / WARN / TRAP.
 
@@ -222,6 +299,7 @@ def build_echo_trap_callback(
     threshold: float,
     halt_on_trap: bool = True,
     ngram_n: int = 2,
+    tokenizer_aware: bool = False,
 ):
     """Live HF Trainer callback for echo-trap detection.
 
@@ -244,6 +322,11 @@ def build_echo_trap_callback(
         raise TypeError(
             f"halt_on_trap must be bool, got {type(halt_on_trap).__name__}"
         )
+    if not isinstance(tokenizer_aware, bool):
+        raise TypeError(
+            "tokenizer_aware must be bool, got "
+            f"{type(tokenizer_aware).__name__}"
+        )
     _check_ngram_n(ngram_n)
     raise NotImplementedError(
         f"Live echo-trap HF Trainer callback (threshold={fv}) is deferred "
@@ -260,10 +343,14 @@ __all__ = [
     "build_echo_trap_callback",
     "classify_echo_signal",
     "score_echo_signal",
+    "score_echo_signal_tokenized",
     "score_trajectory_repetition",
+    "score_trajectory_repetition_tokenized",
 ]
 
 
 # Type aliases retained for the v0.70.1 wiring.
 TrajectoryTokens = Sequence[str]
 TrajectoryBatch = Iterable[TrajectoryTokens]
+TokenIdTrajectory = Sequence[int]
+TokenIdTrajectoryBatch = Iterable[TokenIdTrajectory]

tests/test_v0700_part_f.py

@@ -19,7 +19,9 @@ class TestEchoTrapPublicSurface:
         from soup_cli.utils import echo_trap
 
         assert hasattr(echo_trap, "score_trajectory_repetition")
+        assert hasattr(echo_trap, "score_trajectory_repetition_tokenized")
         assert hasattr(echo_trap, "score_echo_signal")
+        assert hasattr(echo_trap, "score_echo_signal_tokenized")
         assert hasattr(echo_trap, "classify_echo_signal")
         assert hasattr(echo_trap, "EchoTrapReport")
         assert hasattr(echo_trap, "build_echo_trap_callback")
@@ -151,6 +153,38 @@ class TestScoreEchoSignal:
             score_echo_signal(big, ngram_n=2)
 
 
+class TestTokenizedEchoSignal:
+    def test_tokenized_repetition_catches_subword_echo_trap(self):
+        from soup_cli.utils.echo_trap import (
+            classify_echo_signal,
+            score_echo_signal,
+            score_echo_signal_tokenized,
+        )
+
+        decoded_tokens = ["ha,", "ha.", "ha!", "ha?", "ha;", "ha:"]
+        repeated_token_ids = [101, 202, 101, 202, 101, 202, 101, 202]
+
+        whitespace_score = score_echo_signal([decoded_tokens], ngram_n=2)
+        tokenized_score = score_echo_signal_tokenized([repeated_token_ids], ngram_n=2)
+
+        assert classify_echo_signal(whitespace_score) == "OK"
+        assert classify_echo_signal(tokenized_score) == "TRAP"
+
+    def test_tokenized_trajectory_rejects_non_int_ids(self):
+        from soup_cli.utils.echo_trap import score_trajectory_repetition_tokenized
+
+        with pytest.raises(TypeError, match="token_ids"):
+            score_trajectory_repetition_tokenized([1, "2", 3], ngram_n=2)
+        with pytest.raises(TypeError, match="token_ids"):
+            score_trajectory_repetition_tokenized([1, True, 3], ngram_n=2)
+
+    def test_tokenized_batch_rejects_str(self):
+        from soup_cli.utils.echo_trap import score_echo_signal_tokenized
+
+        with pytest.raises(TypeError, match="token-id"):
+            score_echo_signal_tokenized("not ids", ngram_n=2)
+
+
 class TestClassifyEchoSignal:
     """OK / WARN / TRAP taxonomy (mirrors v0.26 / v0.56 / v0.70 Part A).
 
@@ -310,6 +344,15 @@ class TestBuildEchoTrapCallbackDeferred:
         with pytest.raises(TypeError, match="halt"):
             build_echo_trap_callback(threshold=0.5, halt_on_trap="yes")  # type: ignore[arg-type]
 
+    def test_tokenizer_aware_must_be_bool(self):
+        from soup_cli.utils.echo_trap import build_echo_trap_callback
+
+        with pytest.raises(TypeError, match="tokenizer_aware"):
+            build_echo_trap_callback(
+                threshold=0.5,
+                tokenizer_aware="yes",  # type: ignore[arg-type]
+            )
+
 
 # ---------------------------------------------------------------------------
 # Schema integration — TrainingConfig + SoupConfig
@@ -324,6 +367,7 @@ class TestSchemaTrainingConfig:
         assert tcfg.echo_trap_enabled is False
         assert tcfg.echo_trap_threshold == 0.6
         assert tcfg.echo_trap_halt is False
+        assert tcfg.echo_trap_tokenizer_aware is False
 
     def test_threshold_bounds(self):
         from pydantic import ValidationError
@@ -352,6 +396,7 @@ data:
 training:
   echo_trap_enabled: true
   echo_trap_threshold: 0.55
+  echo_trap_tokenizer_aware: true
 """
 
     def test_grpo_accepted(self):
@@ -359,6 +404,7 @@ training:
 
         cfg = load_config_from_string(self._yaml("grpo"))
         assert cfg.training.echo_trap_enabled is True
+        assert cfg.training.echo_trap_tokenizer_aware is True
 
     def test_ppo_accepted(self):
         from soup_cli.config.loader import load_config_from_string
@@ -388,6 +434,22 @@ training:
 """
             )
 
+    def test_tokenizer_aware_without_enabled_rejected(self):
+        from soup_cli.config.loader import load_config_from_string
+
+        with pytest.raises(ValueError, match="echo_trap_enabled"):
+            load_config_from_string(
+                """
+base: meta-llama/Llama-3.1-8B
+task: grpo
+data:
+  train: ./data/train.jsonl
+  format: chatml
+training:
+  echo_trap_tokenizer_aware: true
+"""
+            )
+
 
 # ---------------------------------------------------------------------------
 # Source wiring guards