fix(local-encoder): token-budget fit + agent-session noise isolation for classify_zero_shot #98

Merged
alee merged 3 commits from fix/local-encoder-token-budget into main 2026-09-19 15:57:15 +00:00
3 changed files with 846 additions and 20 deletions

View File

@@ -1025,6 +1025,32 @@ The items below remain open.
`POST /outcome` is the answer for streamed traffic: it arrives afterwards,
so it works identically either way.
5. **`local_encoder` noise isolation — built for the confirmed shapes; residuals below.**
The raw task's fenced code blocks, `Tool result:`-shaped lines, and closed
`<system-reminder>` spans are now stripped by `_isolate_task_text` (pure,
stdlib-only, deterministic) before the fit + embed pass — `classify_zero_shot`
runs isolate → fit → prefix, so cleaning happens first and a noisy task often
fits the token window outright. Measured on the real `BAAI/bge-large-en-v1.5`
(2026-09-19): 8 clean one-sentence tasks scored 8/8, but the same instructions
wrapped in that noise scored 2/8 with the truncation fix already in place —
and tail-biased windowing ALONE also scored 2/8 on long noisy pairs, so the
fit does not subsume isolation. Post-isolation: 8/8 on short and long noisy
pairs, clean-vs-noisy pair consistency 8/8 + 8/8; `'[code]'`/`'[elided]'`
placeholder tokens measured worse than pure removal (7/8, 6/8 on long pairs)
and were rejected; a 20%-ratio floor guard measured harmful (3/8 — it reverts
exactly the short noisy inputs isolation exists to fix) in favor of an
absolute 24-char floor that only catches near-all-code inputs. Offline
regression: a noise tripwire (same instruction bare vs wrapped must classify
identically) fails against pre-isolation code and passes after. Still not
built: attention-masking de-weighting as an alternative to stripping (it
would preserve the noise tokens' presence without letting them dominate);
unclosed `<system-reminder>` tags and un-fenced diff hunks are left in place
(only closed-tag spans, fenced blocks, and marker-prefixed lines are
stripped); and a code-grounded instruction whose pasted snippet is the
subject can still land on a near-category — measured 3/4 on a 4-task
grounded set, the residual miss being description similarity
("refactor this helper" + code → `debugging`), not noise dominance.
## Gaming mode, and the backoff that used to do nothing
**The classifier circuit breaker did not break the circuit.**

View File

@@ -34,7 +34,10 @@ so the router itself has no UI dependency." A deployment that never selects
from __future__ import annotations
from typing import Any, Optional
import re
from typing import Any, Final
import logs
_MISSING_DEPENDENCY_MESSAGE = (
"classifier.mode is 'local_encoder' but the 'transformers'/'torch' "
@@ -87,6 +90,228 @@ _BGE_QUERY_INSTRUCTION = "Represent this sentence for searching relevant passage
# confidence being a genuine probability, not a monotonic score.
_SOFTMAX_TEMPERATURE = 0.10
# ── Task-text isolation (structural agent-session noise) ──────────────
#
# Production task text is raw user_content (dispatcher passes the RAW
# task, not the clamped/framed user_content): fenced code blocks, tool
# call/result stanzas, and system-reminder-style wrappers surround the
# prose instruction. Mean pooling weights every token equally, so when
# noise tokens outnumber the instruction words the pooled embedding —
# and therefore the verdict — describes the noise, not the intent.
# Observed on the real BAAI/bge-large-en-v1.5 (2026-09-19): 8 clean
# one-sentence tasks scored 8/8 correct, but the same instructions
# wrapped in a fenced code block + a "Tool result:" line + a
# <system-reminder> tag classified 2/8 — every non-summarization
# instruction flipped to a file/code-flavored category. Tail-biased
# windowing (the fit below) does NOT fix this: measured 2/8 on long
# noisy inputs too, because the window keeps code and instruction
# mixed. The helper below strips that structural noise BEFORE the fit +
# embed pass; cleaning first also means a noisy task often fits the
# window outright, so the fit sees prose, not a code wall.
#
# Deliberately conservative, pure, stdlib-only (``re``), deterministic,
# and free of false strip targets for ordinary prose: inline code spans
# (``like this``) are KEPT — instructions reference real file names and
# commands, which are signal, not noise.
# Fenced code blocks: strip the content outright. Empirically decided
# placeholder vs pure removal on the real model (2026-09-19, 8 noisy
# pairs x short+long): removal scored 8/8 on both, while '[code]' and
# '[elided]' placeholders scored 7/8 and 6/8 on long pairs — the
# placeholder token itself pulls the pooled embedding toward code-
# flavored categories and costs accuracy, so the winner is removal.
_FENCE_PAIR_RE: Final = re.compile(r"```.*?```", re.DOTALL)
_FENCE_OPEN_TAIL_RE: Final = re.compile(r"```.*\Z", re.DOTALL)
# System-reminder-style wrapper tags: agent harnesses inject session
# metadata inside custom tags; the content is context, not the task.
# Closed-tag spans only — an unclosed tag is left alone rather than
# guessing where it ends (the floor guard still protects the verdict).
_ISOLATE_WRAPPER_TAGS: Final[tuple[str, ...]] = ("system-reminder",)
_WRAPPER_TAG_RE: Final = re.compile(
r"<({})\b[^>]*>.*?</\s*\1\s*>".format("|".join(_ISOLATE_WRAPPER_TAGS)),
re.DOTALL | re.IGNORECASE,
)
# Tool-result lines: "Tool result: file read successfully, 12 lines" and
# the call/output variants seen in agent-session transcripts. Line-level
# and conservative: the line must START with the marker (case-
# insensitive), so prose merely mentioning a tool result survives.
_TOOL_LINE_RE: Final = re.compile(
r"^[ \t]*tool[\s_-]*(?:result|output|call)[\s_-]*:[^\n]*\n?",
re.IGNORECASE | re.MULTILINE,
)
# Floor guard: if stripping leaves (almost) nothing, the input was
# itself one big code block — e.g. a bare file dump with no instruction.
# Embedding a bare placeholder would score a meaningless stub, so fall
# back to the original text, whose code content is then the only signal
# available. Measured (2026-09-19, real model): code-GROUNDED tasks
# (instruction + pasted code) classify BETTER stripped even though
# stripping keeps well under 20% of the characters, so the guard keys on
# an absolute character floor only — a ratio guard would revert exactly
# the noise-wrapped short inputs the isolation exists to fix.
_ISOLATE_MIN_CHARS: Final[int] = 24
def _isolate_task_text(task: str) -> str:
"""Strip structural agent-session noise from *task* before embedding.
Removes fenced code block bodies, system-reminder-style wrapper tag
spans, and tool call/result lines — the shapes that dominate a
mean-pooled embedding when they outnumber the prose instruction.
Inline code spans and ordinary prose are kept. Pure string in,
string out; no I/O, no state.
If stripping leaves fewer than _ISOLATE_MIN_CHARS — the input was
essentially only noise (a bare code dump, no instruction) — the
original text is returned: the code content is then the only signal
there is, and a stripped stub would be a worse embedding target.
Idempotent: already-isolated text contains none of the stripped
shapes, so a second pass returns it unchanged.
"""
if not task:
return task
text = _FENCE_PAIR_RE.sub("", task)
text = _FENCE_OPEN_TAIL_RE.sub("", text)
text = _WRAPPER_TAG_RE.sub("", text)
text = _TOOL_LINE_RE.sub("", text)
if len(text) < _ISOLATE_MIN_CHARS:
return task
return text
# ── Token-budget fitting (pre-tokenization window selection) ──────────
#
# The final ``tokenizer(...)`` call in ``classify_zero_shot`` passes
# ``truncation=True`` with no ``max_length``, so HuggingFace truncates to
# the tokenizer's ``model_max_length`` — 512 tokens for the BERT-family
# BGE models — and, with the default ``truncation_side='right'``, KEEPS
# the head. The head of an agent-session prompt is thousands of chars of
# shared boilerplate while the task-specific content lives in the tail,
# so every long request scored the same truncated prefix and returned the
# same verdict (observed live on 2026-09-18: six distinct 6.4k-6.6k-char
# requests, all classified ``diff_checking`` at confidence 0.192363).
# The helpers below fit the task to the real token window BEFORE that
# final tokenize, choosing ourselves which tokens survive.
# Tail-biased head+tail split: when the task overflows the budget, the
# head slice keeps this share of the content budget and the tail gets the
# rest. A head slice survives at all because the raw task text can carry
# its instruction at either end ("Translate this: <doc>" vs "<doc> —
# translate this") — the same rationale as dispatcher.clamp_for_classifier,
# and it applies here directly because the local_encoder path receives the
# RAW task (dispatcher._classify_via_configured_mode), not the clamped/
# framed user_content the chat modes see. The head share stays small
# because the live failure mode was the opposite collapse: shared head
# boilerplate crowding out the tail.
_FIT_HEAD_FRACTION: Final[float] = 0.25
# Marker spliced between the kept head and tail, mirroring
# dispatcher.clamp_for_classifier's elision convention. Its own token
# cost is measured per tokenizer and reserved from the budget.
_FIT_ELISION_MARKER: Final[str] = "\n\n[... content elided ...]\n\n"
# decode() -> re-tokenize is not guaranteed token-identical for every
# input (unicode normalization, odd graphemes), so the fitted text is
# built a few tokens UNDER budget. If drift ever exceeds the margin, the
# final ``truncation=True`` clips a token or two off the tail — a
# bounded, graceful degradation, never a crash.
_FIT_SAFETY_MARGIN_TOKENS: Final[int] = 8
# model_max_length values some tokenizers report when the true window is
# unknown (HF's VERY_LARGE_INTEGER is int(1e30)). Unusable for budget
# math; falls back to the documented constant below.
_MODEL_MAX_LENGTH_SENTINEL: Final[int] = int(1e30)
# Conservative window used when model_max_length is missing or a
# sentinel: 512 is the BERT-family window the BGE models actually have.
_FALLBACK_MODEL_MAX_LENGTH: Final[int] = 512
# Set once the fallback above has been logged, so a misconfigured
# tokenizer logs one warning instead of one per request.
_budget_fallback_logged = False
def _special_tokens_allowance(tokenizer: Any) -> int:
"""Tokens the final tokenize adds around the text ([CLS]/[SEP] for
BERT-family tokenizers, including the BGE models)."""
add_fn = getattr(tokenizer, "num_special_tokens_to_add", None)
if add_fn is None:
return 2
return max(0, int(add_fn()))
def _usable_token_budget(tokenizer: Any) -> int:
"""Tokens available for task content under the model's real window.
usable = model_max_length - len(BGE instruction prefix tokens)
- special tokens added by the final tokenize.
A missing or sentinel ``model_max_length`` falls back to
_FALLBACK_MODEL_MAX_LENGTH (logged once) rather than raising — an
unusual tokenizer must degrade conservatively, not kill the request.
"""
global _budget_fallback_logged
max_len = getattr(tokenizer, "model_max_length", None)
if (
not isinstance(max_len, int)
or max_len <= 0
or max_len >= _MODEL_MAX_LENGTH_SENTINEL
):
if not _budget_fallback_logged:
_budget_fallback_logged = True
logs.warning(
"encoder_token_budget_fallback",
reported_model_max_length=str(max_len),
fallback=_FALLBACK_MODEL_MAX_LENGTH,
)
max_len = _FALLBACK_MODEL_MAX_LENGTH
prefix_ids = tokenizer(
_BGE_QUERY_INSTRUCTION, add_special_tokens=False,
)["input_ids"]
return max_len - len(list(prefix_ids)) - _special_tokens_allowance(tokenizer)
def _fit_task_to_token_budget(task: str, tokenizer: Any) -> str:
"""Fit *task* into the model's token window, biased toward the tail.
Tokenizes the bare task (no special tokens, no truncation) and, when
it overflows the usable budget, keeps a small head slice plus a
tail-majority window (see _FIT_HEAD_FRACTION for why a head slice
survives at all), spliced with the elision marker. The window is
selected on real token ids and decoded back to text — no
chars-per-token arithmetic anywhere.
Returns *task* unchanged when it already fits (including the empty
string). Never mutates tokenizer state — the cached tokenizer is
shared across requests — and the final ``truncation=True`` in
classify_zero_shot stays as a last-resort backstop.
"""
budget = _usable_token_budget(tokenizer)
if budget <= 0:
# Pathological window (smaller than prefix + specials): leave the
# task alone; the final tokenize does whatever it can.
return task
ids = list(tokenizer(task, add_special_tokens=False)["input_ids"])
if len(ids) <= budget:
return task
marker_tokens = len(list(
tokenizer(_FIT_ELISION_MARKER, add_special_tokens=False)["input_ids"]
))
content_budget = budget - marker_tokens - _FIT_SAFETY_MARGIN_TOKENS
if content_budget <= 0:
# Degenerate window: no room for the marker or the margin — keep
# a pure tail, which is where task-specific content lives.
return tokenizer.decode(ids[-budget:], skip_special_tokens=True)
head_n = int(content_budget * _FIT_HEAD_FRACTION)
tail_n = content_budget - head_n
head_text = tokenizer.decode(ids[:head_n], skip_special_tokens=True)
tail_text = tokenizer.decode(ids[-tail_n:], skip_special_tokens=True)
return head_text + _FIT_ELISION_MARKER + tail_text
# Cached loaded model + tokenizer, keyed by (model_id, device) so a config
# change to either picks up a fresh instance rather than silently reusing
# one built for a different model.
@@ -246,6 +471,15 @@ def classify_zero_shot(
Categories are matched under their natural-language description
(``_CATEGORY_DESCRIPTIONS``), never the raw config identifier — see that
map's docstring for why.
Long inputs are first fitted to the model's real token window (see
``_fit_task_to_token_budget``): the tokenizer's own silent truncation
would otherwise keep only the shared head of a long prompt and drop
the task-specific tail. Before that, structural agent-session noise
(fenced code blocks, tool call/result lines, system-reminder-style
tags) is stripped (see ``_isolate_task_text``): mean pooling weights
every token equally, so unwrapped noise outnumbers the instruction
and the verdict scores the noise, not the intent.
"""
import torch
@@ -257,8 +491,16 @@ def classify_zero_shot(
model, tokenizer, tensor_device, model_id, device, categories,
)
# Encode the input task with BGE's query instruction prefix
prefixed_task = _BGE_QUERY_INSTRUCTION + task
# Isolate then fit, in that order: stripping the noise first means
# the window sees prose and a noisy task often fits outright; the
# fit then protects whatever long prose survives. Without isolation
# the tokenizer below silently right-truncates to model_max_length
# (512 for the BGE models), which for long agent-session prompts
# scores only the shared head boilerplate and discards the
# task-specific tail: every long request returns the same verdict.
isolated_task = _isolate_task_text(task)
fitted_task = _fit_task_to_token_budget(isolated_task, tokenizer)
prefixed_task = _BGE_QUERY_INSTRUCTION + fitted_task
encoded = tokenizer(
prefixed_task,
padding=True,

View File

@@ -267,29 +267,64 @@ def _install_fake_modules():
class MockTokenizer:
"""Acts as transformers.AutoTokenizer.
Always returns fixed ``input_ids`` and unit ``attention_mask``,
both as ``MockTensor``.
Emulates the parts of the HF contract the encoder path relies
on: a known ``model_max_length`` (512, BERT/BGE family),
right-side truncation when called with ``truncation=True`` and
no ``max_length`` (HF's default ``truncation_side='right'``),
and plain Python id lists when called without ``return_tensors``
(the real BatchEncoding behaves the same). Ids are word lengths,
so token counts stay content-dependent.
"""
model_max_length = 512
truncation_side = "right"
@classmethod
def from_pretrained(cls, model_id, **kwargs):
return cls()
def __call__(self, text, **_kw):
# Return a single token with attention_mask = 1
def __call__(
self, text, add_special_tokens=True, truncation=False,
padding=True, return_tensors=None, max_length=None, **_kw,
):
ids = [len(word) for word in str(text).split()]
if truncation:
keep = self.model_max_length - (2 if add_special_tokens else 0)
ids = ids[:keep]
if add_special_tokens:
ids = [0] + ids + [1] # [CLS] ... [SEP]
input_ids = MockTensor([ids]) if return_tensors == "pt" else ids
return {
"input_ids": MockTensor([[1]]),
"attention_mask": MockTensor([[1]]),
"input_ids": input_ids,
"attention_mask": MockTensor([[1] * len(ids)]),
}
def num_special_tokens_to_add(self) -> int:
return 2
def decode(self, ids, skip_special_tokens=True):
kept = [i for i in ids if not (skip_special_tokens and i in (0, 1))]
return " ".join(str(i) for i in kept)
# ── Mock torch ─────────────────────────────────────────────────
def _flat_values(x):
"""Read a 1-D value list from a MockTensor of either orientation.
torch.tensor(flat_list) makes a (1, n) ROW, while softmax/argmax
results elsewhere are (n, 1) COLUMNS; the old readers only handled
the column shape, so a multi-category row silently collapsed to
its first value and every classify returned (category_0, 1.0).
"""
if isinstance(x, MockTensor):
if len(x._data) == 1:
return list(x._data[0])
return [row[0] for row in x._data]
return list(x)
def _softmax(x, dim=0):
"""Compute softmax over dim 0 using MockTensor values."""
if isinstance(x, MockTensor):
vals = [row[0] for row in x._data] if x._data else []
else:
vals = list(x)
vals = _flat_values(x)
max_v = max(vals)
exps = [math.exp(v - max_v) for v in vals]
total = sum(exps)
@@ -297,12 +332,8 @@ def _install_fake_modules():
return MockTensor([[p] for p in probs])
def _argmax(x, dim=None):
if isinstance(x, MockTensor):
vals = [row[0] for row in x._data] if x._data else []
else:
vals = list(x)
idx = vals.index(max(vals))
return MockTensor([idx])
vals = _flat_values(x)
return MockTensor([vals.index(max(vals))])
def _make_tensor(data):
if isinstance(data, MockTensor):
@@ -344,9 +375,11 @@ def _clear_caches():
"""All caches in local_encoder must be empty between tests."""
local_encoder._model_cache.clear()
local_encoder._desc_cache.clear()
local_encoder._budget_fallback_logged = False
yield
local_encoder._model_cache.clear()
local_encoder._desc_cache.clear()
local_encoder._budget_fallback_logged = False
@pytest.fixture
@@ -649,3 +682,528 @@ def test_desc_cache_extended_for_new_categories():
assert cached_after is not None
assert "coding_general" in cached_after
assert "debugging" in cached_after
# ── Token-budget fitting (the 2026-09-18 truncation regression) ────────
#
# classify_zero_shot's final tokenizer call used truncation=True without
# max_length, so HuggingFace clipped the prefixed task to the tokenizer's
# model_max_length from the RIGHT: for BGE (BERT family, 512) a 6000+
# char agent-session prompt became ~510 tokens of its shared head
# boilerplate, the task-specific tail never reached the model, and every
# long request returned the same verdict. Production evidence: six
# distinct requests (6425-6572 chars) all classified diff_checking at
# confidence 0.192363 for 20+ minutes.
#
# The fakes below emulate the HF contract faithfully enough that the
# regression test FAILS against the pre-fix code and PASSES after.
class SessionStyleTokenizer:
"""Fake HF tokenizer with the contract that broke production.
- known ``model_max_length`` (512, the BERT/BGE window);
- right-side truncation when called with ``truncation=True`` and no
``max_length`` (HF's default ``truncation_side='right'`` keeps the
head and silently drops the tail);
- plain Python id lists when called without ``return_tensors``;
- content-dependent ids: one id per distinct word, allocated in
first-seen order, with a ``decode`` that reconstructs text.
"""
def __init__(self, model_max_length: int = 512) -> None:
self.model_max_length = model_max_length
self.truncation_side = "right"
self._vocab: dict = {}
self._words: dict = {}
@classmethod
def from_pretrained(cls, model_id, **kwargs):
return cls()
def _ids_for(self, text: str) -> list:
ids = []
for word in text.split():
if word not in self._vocab:
next_id = len(self._vocab) + 2 # 0=[CLS], 1=[SEP]
self._vocab[word] = next_id
self._words[next_id] = word
ids.append(self._vocab[word])
return ids
def __call__(
self, text, add_special_tokens=True, truncation=False,
padding=True, return_tensors=None, max_length=None, **_kw,
):
ids = self._ids_for(str(text))
if truncation:
keep = self.model_max_length - (2 if add_special_tokens else 0)
ids = ids[:keep]
if add_special_tokens:
ids = [0] + ids + [1]
input_ids = MockTensor([ids]) if return_tensors == "pt" else ids
return {
"input_ids": input_ids,
"attention_mask": MockTensor([[1] * len(ids)]),
}
def num_special_tokens_to_add(self) -> int:
return 2
def decode(self, ids, skip_special_tokens=True):
words = [
self._words[i]
for i in ids
if not (skip_special_tokens and i in (0, 1))
]
return " ".join(words)
class ContentSignModel:
"""Embeds a token sequence as the SIGN of its last content token id.
Provably identical verdicts for inputs that truncate to the same
prefix (same last surviving token), and opposite sides of the
centroid pair for the regression test's two tails: with 100 alpha
words vs 101 beta words after the same boilerplate, the last content
token takes ids of opposite parity by construction. The seeded
[[1.0]] (coding_general) / [[-1.0]] (debugging) descriptions turn
that sign into a category flip.
"""
@classmethod
def from_pretrained(cls, model_id, **kwargs):
return cls()
def to(self, device_str):
return self
def __call__(self, **kwargs):
ids = kwargs["input_ids"]
if isinstance(ids, MockTensor):
ids = ids._data[0]
content = [i for i in ids if i not in (0, 1)]
last = float(content[-1]) if content else 0.0
sign = 1.0 if last % 2 == 0 else -1.0
return types.SimpleNamespace(last_hidden_state=MockTensor([[sign]]))
def _session_style_prompt(tail_words: list) -> str:
"""An agent-session-shaped prompt: huge shared head, varying tail."""
boilerplate = " ".join(
f"session guideline {i} assume tools available" for i in range(400)
)
return f"{boilerplate} {' '.join(tail_words)}"
def test_long_shared_prefix_tasks_do_not_collapse_to_identical_verdict():
"""THE regression tripwire: two long prompts sharing thousands of
chars of identical head boilerplate but differing only in tail
content must NOT produce the same (category, confidence).
Pre-fix, both prefixed tasks right-truncated to the same 510 shared
boilerplate tokens -> same embedding -> same verdict, exactly the
production signature. This fails against pre-fix code by
construction: the tail words sit beyond token 510 and are never seen
by the model. Post-fix the tail survives the fit and the verdicts
separate.
"""
tail_a = [f"alpha{i}" for i in range(100)]
tail_b = [f"beta{i}" for i in range(101)]
task_a = _session_style_prompt(tail_a)
task_b = _session_style_prompt(tail_b)
assert len(task_a) > 2000 and len(task_b) > 2000
assert task_a.startswith("session guideline 0") == task_b.startswith(
"session guideline 0"
)
assert task_a[: -len(" ".join(tail_a)) - 1] == (
task_b[: -len(" ".join(tail_b)) - 1]
)
sys.modules["transformers"].AutoTokenizer = SessionStyleTokenizer
sys.modules["transformers"].AutoModel = ContentSignModel
local_encoder._desc_cache[("m", "cpu")] = {
"coding_general": MockTensor([[1.0]]),
"debugging": MockTensor([[-1.0]]),
}
cat_a, conf_a = local_encoder.classify_zero_shot(
task_a, ["coding_general", "debugging"], model_id="m", device="cpu",
)
cat_b, conf_b = local_encoder.classify_zero_shot(
task_b, ["coding_general", "debugging"], model_id="m", device="cpu",
)
assert (cat_a, conf_a) != (cat_b, conf_b)
assert cat_a != cat_b
def test_fit_task_noop_when_under_budget():
"""A task within the usable budget passes through unchanged."""
tok = SessionStyleTokenizer(model_max_length=512)
short = "summarize this paragraph carefully"
assert local_encoder._fit_task_to_token_budget(short, tok) == short
def test_fit_task_over_budget_keeps_head_and_tail_drops_middle():
"""Over budget: the fitted text keeps a head slice and the tail,
drops the middle, splices the elision marker, and the re-tokenized
prefixed result fits the model window."""
tok = SessionStyleTokenizer(model_max_length=512)
task = (
" ".join(f"head{i}" for i in range(200))
+ " " + " ".join(f"mid{i}" for i in range(2000))
+ " " + " ".join(f"tail{i}" for i in range(200))
)
fitted = local_encoder._fit_task_to_token_budget(task, tok)
assert fitted != task
assert "head0" in fitted and "head121" in fitted
assert "tail199" in fitted and "mid1999" in fitted
assert "mid1000" not in fitted
assert "[... content elided ...]" in fitted
total = tok(local_encoder._BGE_QUERY_INSTRUCTION + fitted)
assert len(total["input_ids"]) <= tok.model_max_length
def test_fit_budget_subtracts_prefix_and_special_tokens():
"""The usable budget is model_max_length - prefix tokens - specials:
a task of exactly that many tokens is a no-op, one token more is
fitted, and the fitted result fits the real window."""
tok = SessionStyleTokenizer(model_max_length=512)
prefix_words = len(local_encoder._BGE_QUERY_INSTRUCTION.split())
budget = (
tok.model_max_length
- prefix_words
- tok.num_special_tokens_to_add()
)
exact = " ".join(f"w{i}" for i in range(budget))
assert local_encoder._fit_task_to_token_budget(exact, tok) == exact
over = exact + " overflow"
fitted = local_encoder._fit_task_to_token_budget(over, tok)
assert fitted != over
total = tok(local_encoder._BGE_QUERY_INSTRUCTION + fitted)
assert len(total["input_ids"]) <= tok.model_max_length
@pytest.mark.parametrize("bad_max_length", [None, int(1e30), 0, -5])
def test_fit_falls_back_when_model_max_length_unusable(bad_max_length):
"""A missing/sentinel/non-positive model_max_length falls back to the
conservative 512 window (logged once via the module flag) instead of
blowing up, and the fitted result still fits 512."""
tok = SessionStyleTokenizer(model_max_length=512)
tok.model_max_length = bad_max_length
task = " ".join(f"w{i}" for i in range(4000))
fitted = local_encoder._fit_task_to_token_budget(task, tok)
assert fitted != task
assert local_encoder._budget_fallback_logged
total = tok(local_encoder._BGE_QUERY_INSTRUCTION + fitted)
assert len(total["input_ids"]) <= local_encoder._FALLBACK_MODEL_MAX_LENGTH
def test_fit_degenerate_window_returns_task_unchanged():
"""A window smaller than prefix + specials leaves the task alone
rather than producing a nonsensical slice."""
tok = SessionStyleTokenizer(model_max_length=4) # < prefix(7) + 2
task = "a b c"
assert local_encoder._fit_task_to_token_budget(task, tok) == task
def test_classify_zero_shot_leaves_cached_tokenizer_state_untouched():
"""The fit must not mutate the shared cached tokenizer (no
truncation_side flips) — the cache is process-wide."""
sys.modules["transformers"].AutoTokenizer = SessionStyleTokenizer
sys.modules["transformers"].AutoModel = ContentSignModel
local_encoder.classify_zero_shot(
_session_style_prompt([f"alpha{i}" for i in range(100)]),
["coding_general", "debugging"],
model_id="m", device="cpu",
)
tok = local_encoder._model_cache[("m", "cpu")][1]
assert tok.truncation_side == "right"
assert tok.model_max_length == 512
# ── Task-text isolation (the 2026-09-19 noise-dominance regression) ────
#
# classify_zero_shot embedded the RAW task: a prose instruction wrapped
# in agent-session noise (fenced code block, "Tool result:" line,
# <system-reminder> tag) mean-pooled the noise tokens with equal weight,
# and the verdict scored the noise — measured on the real
# BAAI/bge-large-en-v1.5, 8 clean tasks scored 8/8 but the same
# instructions wrapped scored 2/8, on SHORT inputs far under any
# truncation limit (a second, distinct bug; the fit fix does not touch
# it — windowing alone measured 2/8 on long noisy pairs too).
#
# The fakes below reproduce that mechanism in miniature: tokens drawn
# from the code/structural vocabulary push the pooled embedding toward
# the code centroid, prose tokens toward the prose centroid, so
# noise-count dominance flips the verdict exactly like the real model.
_INSTRUCTION = (
"Please write a clear docstring and a short readme section "
"explaining the storage module."
)
_NOISE_REMINDER = (
"<system-reminder>\nsession context: agent harness active\n"
"</system-reminder>"
)
_NOISE_CODE = (
"```python\ndef parse_line(line):\n"
" tokens = line.strip().split(',')\n"
" if not tokens:\n"
" return []\n"
" return [int(t) for t in tokens if t.isdigit()]\n"
"FILES = {'a': 1, 'b': 2}\n```"
)
_NOISE_TOOL = "Tool result: file read successfully, 12 lines"
def _norm(word: str) -> str:
"""Normalize a whitespace token to its word characters (lowercased)."""
return "".join(ch for ch in word.lower() if ch.isalnum() or ch == "_")
# Every multi-char word of the noise fixture is a "code/structural"
# vocabulary entry — the fake tokenizer's stand-in for "this token
# embeds near code", mirroring how real code/tool tokens pull the
# pooled embedding toward code-flavored categories.
_CODE_VOCAB = frozenset(
w
for w in (
_norm(tok)
for tok in (_NOISE_REMINDER + " " + _NOISE_CODE + " " + _NOISE_TOOL).split()
)
if len(w) > 1
)
class NoiseDominantTokenizer:
"""Fake HF tokenizer whose ids encode word class.
Words in _CODE_VOCAB get ids >= 1000; prose words get ids from 2 up
(first-seen, like SessionStyleTokenizer). Otherwise the same HF
contract: known model_max_length, right-side truncation under
``truncation=True`` without ``max_length``, plain id lists without
``return_tensors``, and a decode that reconstructs text.
"""
def __init__(self, model_max_length: int = 512) -> None:
self.model_max_length = model_max_length
self.truncation_side = "right"
self._vocab: dict = {}
self._words: dict = {}
@classmethod
def from_pretrained(cls, model_id, **kwargs):
return cls()
def _ids_for(self, text: str) -> list:
ids = []
for word in str(text).split():
norm = _norm(word)
if not norm:
continue
if norm not in self._vocab:
if norm in _CODE_VOCAB:
next_id = 1000 + len(self._vocab)
else:
next_id = 2 + len(self._vocab)
self._vocab[norm] = next_id
self._words[next_id] = word
ids.append(self._vocab[norm])
return ids
def __call__(
self, text, add_special_tokens=True, truncation=False,
padding=True, return_tensors=None, max_length=None, **_kw,
):
ids = self._ids_for(text)
if truncation:
keep = self.model_max_length - (2 if add_special_tokens else 0)
ids = ids[:keep]
if add_special_tokens:
ids = [0] + ids + [1]
input_ids = MockTensor([ids]) if return_tensors == "pt" else ids
return {
"input_ids": input_ids,
"attention_mask": MockTensor([[1] * len(ids)]),
}
def num_special_tokens_to_add(self) -> int:
return 2
def decode(self, ids, skip_special_tokens=True):
words = [
self._words[i]
for i in ids
if not (skip_special_tokens and i in (0, 1))
]
return " ".join(words)
class NoiseDominantModel:
"""Embeds a token sequence as the MEAN of per-token class values.
+1.0 for code-vocabulary ids (>= 1000), -1.0 for prose ids: mean
pooling weights every token equally, so whichever class supplies
more tokens wins the pooled sign — the miniature of the real
failure, where code/tool noise outnumbers the instruction words and
the verdict follows the noise.
"""
@classmethod
def from_pretrained(cls, model_id, **kwargs):
return cls()
def to(self, device_str):
return self
def __call__(self, **kwargs):
ids = kwargs["input_ids"]
if isinstance(ids, MockTensor):
ids = ids._data[0]
content = [i for i in ids if i not in (0, 1)]
values = [1.0 if i >= 1000 else -1.0 for i in content]
mean = sum(values) / len(values) if values else 0.0
return types.SimpleNamespace(last_hidden_state=MockTensor([[mean]]))
def _noisy_instruction() -> str:
"""The clean instruction wrapped in the operator's repro noise:
system-reminder tag + fenced code block + tool-result line, the
instruction still clearly present at the END."""
return "\n\n".join([_NOISE_REMINDER, _NOISE_CODE, _NOISE_TOOL, _INSTRUCTION])
def test_noise_wrapped_task_classifies_same_as_clean_pair():
"""THE noise tripwire: the same prose instruction, bare and wrapped
in agent-session noise, must classify to the SAME category.
Fails against pre-isolation code by construction: the wrapped task's
code/structural tokens outnumber the prose tokens, so the pooled
mean goes positive (code centroid) while the bare instruction pools
negative (prose centroid) — a category flip, the miniature of the
real-model observation. Post-isolation the noise shapes are stripped
before embedding and both tasks pool to the prose centroid.
"""
clean = _INSTRUCTION
noisy = _noisy_instruction()
# Precondition: the fixture is short — windowing cannot apply, so
# this pair exercises isolation, not the fit.
assert len(noisy) < 1000
assert not (_CODE_VOCAB & {_norm(w) for w in clean.split()})
sys.modules["transformers"].AutoTokenizer = NoiseDominantTokenizer
sys.modules["transformers"].AutoModel = NoiseDominantModel
local_encoder._desc_cache[("m", "cpu")] = {
"docs_writing": MockTensor([[-1.0]]),
"coding_general": MockTensor([[1.0]]),
}
cat_clean, _ = local_encoder.classify_zero_shot(
clean, ["docs_writing", "coding_general"], model_id="m", device="cpu",
)
cat_noisy, _ = local_encoder.classify_zero_shot(
noisy, ["docs_writing", "coding_general"], model_id="m", device="cpu",
)
assert cat_clean == "docs_writing"
assert cat_noisy == cat_clean, (
f"noise flipped the verdict: clean={cat_clean}, noisy={cat_noisy}"
)
def test_isolate_strips_fenced_code_blocks():
"""A fenced code block is removed outright — body AND fences."""
task = "Summarize the setup steps.\n```python\nimport os\nos.getcwd()\n```"
out = local_encoder._isolate_task_text(task)
assert "import os" not in out
assert "os.getcwd()" not in out
assert "```" not in out
def test_isolate_removal_decision_no_placeholder():
"""The eval-decided placeholder-vs-removal winner is REMOVAL: the
stripped text must not carry any '[code]'-style placeholder — the
placeholder token itself pulled long noisy pairs to 6/8 vs removal's
8/8 on the real model."""
task = (
"Explain what this algorithm computes and when you would use it.\n"
"```\nfor i in range(3):\n print(i)\n```"
)
out = local_encoder._isolate_task_text(task)
assert "[code]" not in out
assert "[elided]" not in out
assert "for i in range(3)" not in out
def test_isolate_strips_tool_result_lines():
"""'Tool result:'-shaped lines are stripped line-level; result and
output/call variants share the marker shape."""
task = (
"Tool result: file read successfully, 12 lines\n"
"Tool output: exit code 0\n"
"Now summarize the findings above."
)
out = local_encoder._isolate_task_text(task)
assert "Tool result:" not in out
assert "Tool output:" not in out
assert "exit code 0" not in out
assert "summarize the findings" in out
def test_isolate_strips_system_reminder_content():
"""<system-reminder> spans are removed with their content; the tag
itself does not survive either."""
task = (
"<system-reminder>session id 42, tools enabled</system-reminder>\n"
"Translate the attached paragraph to French."
)
out = local_encoder._isolate_task_text(task)
assert "session id 42" not in out
assert "system-reminder" not in out
assert "Translate the attached paragraph" in out
def test_isolate_preserves_trailing_instruction():
"""The prose instruction survives verbatim when it trails the
noise — the operator's repro shape."""
assert _noisy_instruction().endswith(_INSTRUCTION)
out = local_encoder._isolate_task_text(_noisy_instruction())
assert out.strip().endswith(_INSTRUCTION)
def test_isolate_keeps_inline_code_spans():
"""Inline code spans are KEPT by default: instructions reference
real file names and commands, which are signal, not noise."""
task = "Rename `config_loader.py` to `settings.py` and fix the imports."
assert local_encoder._isolate_task_text(task) == task
def test_isolate_floor_guard_falls_back_on_all_code_task():
"""A task that is essentially only a code dump (nothing left above
the character floor) falls back to the original text rather than
embedding a stripped stub."""
bare_dump = "```python\n" + "\n".join(f"x{i} = {i}" for i in range(40)) + "\n```"
assert local_encoder._isolate_task_text(bare_dump) == bare_dump
def test_isolate_is_idempotent():
"""Applying isolation twice equals applying it once."""
once = local_encoder._isolate_task_text(_noisy_instruction())
assert local_encoder._isolate_task_text(once) == once
def test_isolate_unterminated_fence_keeps_leading_instruction():
"""An unterminated fence (truncated tool output) eats to the end of
the string, so a leading instruction survives and a trailing one
cannot — the floor guard then reverts rather than embed a stub."""
leading = (
"Summarize the parser module.\n\n```python\nimport csv\nrows = []\n"
)
out = local_encoder._isolate_task_text(leading)
assert "Summarize the parser module." in out
assert "import csv" not in out
trailing = "```python\nimport csv\n" + "x = 1\n" * 30 + "summarize it"
assert local_encoder._isolate_task_text(trailing) == trailing