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.
This commit is contained in:
Shivam 2026-05-25 20:16:51 +05:30 committed by GitHub
parent 2ed7b44ade
commit 4e95d4c71f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 209 additions and 18 deletions

View File

@ -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?

View File

@ -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(

View File

@ -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":

View File

@ -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
hands tensor ids (instead of decoded strings) gets an actionable
error rather than silently misbehaving.
- Whitespace-mode tokens must be strings; tokeniser-aware mode accepts
integer token ids through the dedicated ``*_tokenized`` helpers.
- ``_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 0.0
counts: dict[tuple[str, ...], int] = {}
for i in range(len(tok) - n + 1):
gram = tok[i : i + n]
counts[gram] = counts.get(gram, 0) + 1
if not counts:
return 0.0
repeating = sum(1 for c in counts.values() if c > 1)
return repeating / len(counts)
return _score_repetition(tok, ngram_n=n)
def score_trajectory_repetition_tokenized(
token_ids: object,
*,
ngram_n: object = 2,
) -> float:
"""Per-trajectory repetition score over tokenizer ids.
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]

View File

@ -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