fix(local-encoder): token-budget fit + agent-session noise isolation for classify_zero_shot #98
26
CLAUDE.md
26
CLAUDE.md
@@ -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.**
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -648,4 +681,529 @@ def test_desc_cache_extended_for_new_categories():
|
||||
cached_after = local_encoder._desc_cache.get(("m", "cpu"))
|
||||
assert cached_after is not None
|
||||
assert "coding_general" in cached_after
|
||||
assert "debugging" 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
|
||||
Reference in New Issue
Block a user