27043046a6
_mean_pool(token_embeddings, attention_mask) reads list(_model_pooling_strategies.values())[-1] — the last-inserted value — instead of the strategy belonging to the model whose embeddings are being pooled. This works today because production loads one encoder per process and every test populates at most one dict entry, but breaks silently the moment two backbones share a process. Fix: thread model_id through _mean_pool and its call sites (_build_description_embeddings and _embed_task), resolving the strategy via _model_pooling_strategies.get(model_id, 'mean') instead of reading the global tail. Add test_mean_pool_dispatches_by_model_id_not_load_order which proves the failure mode is closed: two models with different strategies (cls, mean) are each pooled correctly regardless of which was inserted first.
1399 lines
52 KiB
Python
1399 lines
52 KiB
Python
"""local_encoder.py: the embedding+centroid classifier backing classifier.mode: local_encoder.
|
|
|
|
Never loads a real model. transformers/torch are not installed in this
|
|
environment (they are optional, per requirements-encoder.txt), and even
|
|
where they are, a unit test has no business paying multi-second model-load
|
|
time. Every test here stubs both packages in sys.modules.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
import local_encoder
|
|
|
|
|
|
# ── Mock Tensor ────────────────────────────────────────────────────────
|
|
|
|
|
|
class MockTensor:
|
|
"""Minimal tensor emulation for testing the embedding math in local_encoder.
|
|
|
|
Wraps a 2D list of floats (batch, features) internally and supports
|
|
the subset of torch operations used by mean-pool, L2-normalize, cosine
|
|
similarity, and softmax.
|
|
"""
|
|
|
|
def __init__(self, data):
|
|
# Normalise everything to a 2D list of floats.
|
|
if isinstance(data, (int, float)):
|
|
data = [[float(data)]]
|
|
elif isinstance(data, list):
|
|
if not data:
|
|
data = [[]]
|
|
elif not isinstance(data[0], list):
|
|
data = [[float(x) for x in data]]
|
|
else:
|
|
data = [[float(x) for x in row] for row in data]
|
|
self._data = data # list[list[float]]
|
|
|
|
# ── shape ──────────────────────────────────────────────────────
|
|
|
|
@property
|
|
def shape(self):
|
|
if not self._data or not self._data[0]:
|
|
return (0, 0)
|
|
return (len(self._data), len(self._data[0]))
|
|
|
|
def size(self, dim=None):
|
|
s = self.shape
|
|
if dim is None:
|
|
return s
|
|
return s[dim]
|
|
|
|
# ── identity / layout ──────────────────────────────────────────
|
|
|
|
def to(self, *_a, **_k):
|
|
return self
|
|
|
|
def float(self):
|
|
return self
|
|
|
|
def unsqueeze(self, dim):
|
|
"""Return self (for expand compatibility)."""
|
|
return self
|
|
|
|
def expand(self, *shape):
|
|
return self
|
|
|
|
@property
|
|
def T(self):
|
|
"""Transpose: (batch, features) -> (features, batch)."""
|
|
if not self._data or not self._data[0]:
|
|
return MockTensor([])
|
|
rows = len(self._data)
|
|
cols = len(self._data[0])
|
|
transposed = [[self._data[r][c] for r in range(rows)] for c in range(cols)]
|
|
return MockTensor(transposed)
|
|
|
|
def item(self) -> float:
|
|
"""Extract the first element as a scalar float."""
|
|
return self._data[0][0]
|
|
|
|
def __int__(self) -> int:
|
|
"""Convert to int (for int(torch.argmax(...)) compatibility)."""
|
|
return int(self._data[0][0])
|
|
|
|
def __float__(self) -> float:
|
|
"""Convert to float (for float(torch.argmax(...)) compatibility)."""
|
|
return self._data[0][0]
|
|
|
|
def __index__(self) -> int:
|
|
"""Support hex(), oct(), bin() and slice indices."""
|
|
return int(self._data[0][0])
|
|
|
|
def __repr__(self):
|
|
return f"MockTensor({self._data})"
|
|
|
|
def __getitem__(self, idx):
|
|
"""Support tensor indexing (e.g. probs[top_idx], emb[:, 0])."""
|
|
if isinstance(idx, tuple):
|
|
# Support tensor_embeddings[:, 0] for CLS pooling
|
|
if len(idx) == 2 and isinstance(idx[0], slice) and isinstance(idx[1], int):
|
|
return MockTensor([[row[idx[1]]] for row in self._data])
|
|
return self
|
|
if isinstance(idx, int):
|
|
if len(self._data) == 1 and len(self._data[0]) > 1:
|
|
# Row vector: return the element at column idx as a 0D-like tensor
|
|
return MockTensor([[self._data[0][idx]]])
|
|
# Column vector: return the row at idx
|
|
return MockTensor([self._data[idx]])
|
|
if isinstance(idx, MockTensor):
|
|
v = int(idx)
|
|
return self.__getitem__(v)
|
|
if isinstance(idx, slice):
|
|
return self
|
|
return MockTensor([[self._data[0][idx]]])
|
|
|
|
# ── element-wise arithmetic ────────────────────────────────────
|
|
|
|
def __mul__(self, other):
|
|
if isinstance(other, (int, float)):
|
|
return MockTensor([[x * other for x in row] for row in self._data])
|
|
if isinstance(other, MockTensor):
|
|
# same-shape element-wise multiply
|
|
d2 = other._data
|
|
return MockTensor(
|
|
[[a * b for a, b in zip(row, row2)] for row, row2 in zip(self._data, d2)]
|
|
)
|
|
return NotImplemented
|
|
|
|
def __rmul__(self, other):
|
|
if isinstance(other, (int, float)):
|
|
return MockTensor([[other * x for x in row] for row in self._data])
|
|
return NotImplemented
|
|
|
|
def __truediv__(self, other):
|
|
if isinstance(other, (int, float)):
|
|
return MockTensor([[x / other for x in row] for row in self._data])
|
|
if isinstance(other, MockTensor):
|
|
d2 = other._data
|
|
return MockTensor(
|
|
[[a / b for a, b in zip(row, row2)] for row, row2 in zip(self._data, d2)]
|
|
)
|
|
return NotImplemented
|
|
|
|
def __matmul__(self, other):
|
|
"""Dot product / matrix multiply.
|
|
|
|
For the test case: (1, n) @ (n, 1) -> (1, 1) scalar.
|
|
Also handles (1, n) @ (n, m) -> (1, m).
|
|
"""
|
|
if isinstance(other, MockTensor):
|
|
a = self._data
|
|
b = other._data
|
|
# a is (batch, features), b is (features, n_cats) or (features, 1)
|
|
if not a or not b:
|
|
return MockTensor(0.0)
|
|
n_features = len(a[0])
|
|
n_b_cols = len(b[0])
|
|
result_row = []
|
|
for j in range(n_b_cols):
|
|
total = sum(a[0][i] * b[i][j] for i in range(n_features))
|
|
result_row.append(total)
|
|
return MockTensor([result_row])
|
|
return NotImplemented
|
|
|
|
# ── reduction operations ───────────────────────────────────────
|
|
|
|
def sum(self, dim=None, keepdim=False):
|
|
if not self._data or not self._data[0]:
|
|
return MockTensor(0.0)
|
|
rows = len(self._data)
|
|
cols = len(self._data[0])
|
|
if dim is None:
|
|
total = sum(sum(row) for row in self._data)
|
|
if keepdim:
|
|
return MockTensor([[total]])
|
|
return MockTensor(total)
|
|
if dim == 0:
|
|
# sum over rows: produce (1, cols)
|
|
result = [sum(self._data[r][c] for r in range(rows)) for c in range(cols)]
|
|
if keepdim:
|
|
return MockTensor([result])
|
|
return MockTensor(result)
|
|
if dim == 1:
|
|
# sum over columns: produce (rows, 1)
|
|
result = [sum(row) for row in self._data]
|
|
if keepdim:
|
|
return MockTensor([[v] for v in result])
|
|
return MockTensor(result)
|
|
raise ValueError(f"MockTensor.sum: unsupported dim={dim}")
|
|
|
|
def clamp(self, min=None, max=None):
|
|
if min is not None:
|
|
builtin_max = __builtins__["max"] if isinstance(__builtins__, dict) else __builtins__.max
|
|
return MockTensor(
|
|
[[builtin_max(min, x) for x in row] for row in self._data]
|
|
)
|
|
return self
|
|
|
|
def norm(self, dim=None, keepdim=False):
|
|
"""L2 norm (Frobenius norm if dim is None)."""
|
|
if not self._data or not self._data[0]:
|
|
return MockTensor(0.0)
|
|
if dim == 1:
|
|
# per-row norm
|
|
norms = [
|
|
math.sqrt(sum(x * x for x in row))
|
|
for row in self._data
|
|
]
|
|
if keepdim:
|
|
return MockTensor([[v] for v in norms])
|
|
return MockTensor(norms)
|
|
# dim is None → Frobenius norm
|
|
total = math.sqrt(
|
|
sum(sum(x * x for x in row) for row in self._data)
|
|
)
|
|
if keepdim:
|
|
return MockTensor([[total]])
|
|
return MockTensor(total)
|
|
|
|
|
|
# ── Test helpers ───────────────────────────────────────────────────────
|
|
|
|
|
|
def _install_fake_modules():
|
|
"""Install stubs for ``transformers`` and ``torch`` in sys.modules.
|
|
|
|
The fake ``transformers`` module provides ``AutoModel`` and
|
|
``AutoTokenizer``. The default ``AutoModel.from_pretrained()`` returns a
|
|
model whose forward pass returns a ``last_hidden_state`` filled with 1.0,
|
|
and ``AutoTokenizer.from_pretrained()`` returns a tokenizer that produces
|
|
unit ``attention_mask`` tokens.
|
|
|
|
Individual tests may override ``from_pretrained`` to customise the model
|
|
or tokenizer behaviour.
|
|
"""
|
|
|
|
# ── Mock Model ────────────────────────────────────────────────
|
|
|
|
class MockModel:
|
|
"""Acts as transformers.AutoModel.
|
|
|
|
``from_pretrained()`` returns an instance whose ``__call__``
|
|
returns ``Namespace(last_hidden_state=MockTensor(...))``.
|
|
"""
|
|
|
|
_call_count = 0
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, model_id, **kwargs):
|
|
return cls()
|
|
|
|
def to(self, device_str):
|
|
return self
|
|
|
|
def __call__(self, **kwargs):
|
|
MockModel._call_count += 1
|
|
# Default: return all-1 embeddings so the math is
|
|
# deterministic (mean = 1, norm = sqrt(seq_len)).
|
|
attn = kwargs.get("attention_mask", MockTensor([[1]]))
|
|
seq_len = attn.shape[1] if attn.shape[0] > 0 else 1
|
|
batch = attn.shape[0] if attn.shape[0] > 0 else 1
|
|
hidden = MockTensor([[1.0] * seq_len for _ in range(batch)])
|
|
return types.SimpleNamespace(last_hidden_state=hidden)
|
|
|
|
# ── Mock Tokenizer ────────────────────────────────────────────
|
|
|
|
class MockTokenizer:
|
|
"""Acts as transformers.AutoTokenizer.
|
|
|
|
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, 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": 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."""
|
|
vals = _flat_values(x)
|
|
max_v = max(vals)
|
|
exps = [math.exp(v - max_v) for v in vals]
|
|
total = sum(exps)
|
|
probs = [e / total for e in exps]
|
|
return MockTensor([[p] for p in probs])
|
|
|
|
def _argmax(x, dim=None):
|
|
vals = _flat_values(x)
|
|
return MockTensor([vals.index(max(vals))])
|
|
|
|
def _make_tensor(data):
|
|
if isinstance(data, MockTensor):
|
|
return data
|
|
if isinstance(data, (int, float)):
|
|
return MockTensor([data])
|
|
if isinstance(data, (list, tuple)):
|
|
return MockTensor(data)
|
|
return MockTensor([0.0])
|
|
|
|
import contextlib
|
|
|
|
fake_torch = types.ModuleType("torch")
|
|
fake_torch.no_grad = contextlib.nullcontext
|
|
fake_torch.Tensor = MockTensor
|
|
fake_torch.tensor = _make_tensor
|
|
fake_torch.sum = lambda x, dim=None: x.sum(dim=dim)
|
|
fake_torch.softmax = _softmax
|
|
fake_torch.argmax = _argmax
|
|
fake_torch.__version__ = "0.0.0-mock"
|
|
|
|
fake_transformers = types.ModuleType("transformers")
|
|
fake_transformers.AutoModel = MockModel
|
|
fake_transformers.AutoTokenizer = MockTokenizer
|
|
|
|
sys.modules["transformers"] = fake_transformers
|
|
sys.modules["torch"] = fake_torch
|
|
|
|
# Reset call counters
|
|
MockModel._call_count = 0
|
|
MockModel.from_pretrained = classmethod(lambda cls, mid, **kw: cls())
|
|
|
|
|
|
# ── Fixtures ───────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
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._model_pooling_strategies.clear()
|
|
local_encoder._model_query_prefixes.clear()
|
|
local_encoder._budget_fallback_logged = False
|
|
yield
|
|
local_encoder._model_cache.clear()
|
|
local_encoder._desc_cache.clear()
|
|
local_encoder._model_pooling_strategies.clear()
|
|
local_encoder._model_query_prefixes.clear()
|
|
local_encoder._budget_fallback_logged = False
|
|
|
|
|
|
@pytest.fixture
|
|
def restore_sys_modules():
|
|
had_t = "transformers" in sys.modules
|
|
orig_t = sys.modules.get("transformers")
|
|
had_th = "torch" in sys.modules
|
|
orig_th = sys.modules.get("torch")
|
|
yield
|
|
if had_t:
|
|
sys.modules["transformers"] = orig_t
|
|
else:
|
|
sys.modules.pop("transformers", None)
|
|
if had_th:
|
|
sys.modules["torch"] = orig_th
|
|
else:
|
|
sys.modules.pop("torch", None)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def auto_install_fakes(restore_sys_modules):
|
|
"""Install fake transformers+torch for every test by default.
|
|
|
|
Tests that exercise the missing-dependency path must opt out by
|
|
clearing sys.modules within the test body.
|
|
"""
|
|
_install_fake_modules()
|
|
yield
|
|
|
|
|
|
# ── Tests ──────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_missing_dependency_raises_an_actionable_message(monkeypatch):
|
|
"""With transformers/torch removed from sys.modules, the lazy import
|
|
raises an ImportError with the actionable install message."""
|
|
# Remove the stubs
|
|
monkeypatch.delitem(sys.modules, "transformers", raising=False)
|
|
monkeypatch.delitem(sys.modules, "torch", raising=False)
|
|
|
|
import builtins
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *args, **kwargs):
|
|
if name == "transformers":
|
|
raise ImportError("No module named 'transformers'")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
|
|
with pytest.raises(ImportError, match="pip install -r requirements-encoder.txt"):
|
|
local_encoder.classify_zero_shot(
|
|
"do a thing", ["a", "b"], model_id="stub-model", device="cpu"
|
|
)
|
|
|
|
|
|
def test_classify_zero_shot_returns_top_label_and_score():
|
|
"""The category whose description embedding is closest to the input
|
|
embedding should be returned as the top label."""
|
|
# Pre-set description embeddings: cat_a at [1, 0], cat_b at [0, 1]
|
|
# (These are L2-normalised vectors.)
|
|
local_encoder._desc_cache[("stub-model", "cpu")] = {
|
|
"coding_general": MockTensor([[1.0, 0.0]]),
|
|
"debugging": MockTensor([[0.0, 1.0]]),
|
|
}
|
|
|
|
# The classify_zero_shot path will encode the input and compare.
|
|
# With the default mock model (all-1 hidden state in a single token),
|
|
# after mean-pool + L2-normalize the input embedding is [1.0].
|
|
# But the description embeddings are 2D [1.0, 0.0] and [0.0, 1.0].
|
|
# For the @ matmul to work, the input needs to be 2D too.
|
|
#
|
|
# We override the model here: the mock model returns a last_hidden_state
|
|
# that after mean-pool produces [0.9, 0.1] (biased toward cat_a).
|
|
class CustomModel:
|
|
_call_count = 0
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, model_id, **kwargs):
|
|
return cls()
|
|
|
|
def to(self, device_str):
|
|
return self
|
|
|
|
def __call__(self, **kwargs):
|
|
CustomModel._call_count += 1
|
|
# Return a (1, 2) embedding that will pool to [0.9, 0.1]
|
|
return types.SimpleNamespace(
|
|
last_hidden_state=MockTensor([[0.9, 0.1]])
|
|
)
|
|
|
|
sys.modules["transformers"].AutoModel = CustomModel
|
|
|
|
label, score = local_encoder.classify_zero_shot(
|
|
"fix this bug", ["coding_general", "debugging"],
|
|
model_id="stub-model", device="cpu",
|
|
)
|
|
assert label == "coding_general"
|
|
assert 0.0 <= score <= 1.0
|
|
|
|
|
|
def test_classify_zero_shot_sends_descriptions_not_raw_category_ids():
|
|
"""Descriptions from _CATEGORY_DESCRIPTIONS are used to build the
|
|
embedding cache, not raw category IDs."""
|
|
# The cache starts empty. After classify, it should be populated
|
|
# with the category IDs as keys.
|
|
assert local_encoder._desc_cache.get(("m", "cpu")) is None
|
|
|
|
local_encoder.classify_zero_shot(
|
|
"hello", ["general_chat"], model_id="m", device="cpu",
|
|
)
|
|
|
|
cached = local_encoder._desc_cache.get(("m", "cpu"))
|
|
assert cached is not None
|
|
assert "general_chat" in cached
|
|
|
|
|
|
def test_classify_zero_shot_falls_back_to_raw_string_for_unmapped_category():
|
|
"""A category not in _CATEGORY_DESCRIPTIONS uses the raw string as
|
|
its description and does not raise."""
|
|
local_encoder.classify_zero_shot(
|
|
"new feature", ["some_future_category"],
|
|
model_id="m", device="cpu",
|
|
)
|
|
# Should not raise. The cache must contain the category.
|
|
cached = local_encoder._desc_cache.get(("m", "cpu"))
|
|
assert cached is not None
|
|
assert "some_future_category" in cached
|
|
|
|
|
|
def test_model_is_built_once_and_reused():
|
|
"""AutoModel.from_pretrained must be called only once for repeated
|
|
classify calls with the same (model_id, device)."""
|
|
# Grab the original from_pretrained to count calls
|
|
from_pretrained_calls = []
|
|
|
|
RealAutoModel = sys.modules["transformers"].AutoModel
|
|
orig_fp = RealAutoModel.from_pretrained
|
|
|
|
@classmethod
|
|
def counting_fp(cls, model_id, **kwargs):
|
|
from_pretrained_calls.append(1)
|
|
return cls()
|
|
|
|
RealAutoModel.from_pretrained = counting_fp
|
|
|
|
local_encoder.classify_zero_shot("x", ["a"], model_id="same", device="cpu")
|
|
local_encoder.classify_zero_shot("y", ["a"], model_id="same", device="cpu")
|
|
|
|
assert len(from_pretrained_calls) == 1, "model was rebuilt instead of reused"
|
|
|
|
RealAutoModel.from_pretrained = orig_fp
|
|
|
|
|
|
def test_different_model_or_device_gets_its_own_cached_model():
|
|
"""Different (model_id, device) pairs get separate cached
|
|
(model, tokenizer)."""
|
|
from_pretrained_seen = []
|
|
|
|
RealAutoModel = sys.modules["transformers"].AutoModel
|
|
|
|
@classmethod
|
|
def recording_fp(cls, model_id, **kwargs):
|
|
from_pretrained_seen.append(model_id)
|
|
return cls()
|
|
|
|
RealAutoModel.from_pretrained = recording_fp
|
|
|
|
local_encoder.classify_zero_shot("x", ["a"], model_id="model-a", device="cpu")
|
|
local_encoder.classify_zero_shot("x", ["a"], model_id="model-b", device="cpu")
|
|
local_encoder.classify_zero_shot("x", ["a"], model_id="model-a", device="cpu")
|
|
|
|
assert from_pretrained_seen == ["model-a", "model-b"], (
|
|
"a config change should get a fresh model"
|
|
)
|
|
|
|
|
|
def test_device_resolves_correctly():
|
|
"""verify_that the model is placed on the right device string."""
|
|
# We can inspect the model.to() call via a custom wrapper
|
|
device_seen = []
|
|
|
|
class DeviceCheckingModel:
|
|
@classmethod
|
|
def from_pretrained(cls, model_id, **kwargs):
|
|
return cls()
|
|
|
|
def to(self, device_str):
|
|
device_seen.append(device_str)
|
|
return self
|
|
|
|
def __call__(self, **kwargs):
|
|
return types.SimpleNamespace(
|
|
last_hidden_state=MockTensor([[1.0]])
|
|
)
|
|
|
|
sys.modules["transformers"].AutoModel = DeviceCheckingModel
|
|
|
|
# CPU
|
|
local_encoder.ensure_available("m", "cpu")
|
|
assert "cpu" in device_seen, "cpu device should go to 'cpu'"
|
|
|
|
# CUDA
|
|
device_seen.clear()
|
|
local_encoder.ensure_available("m", "cuda")
|
|
assert "cuda:0" in device_seen, "cuda device should go to 'cuda:0'"
|
|
|
|
|
|
def test_ensure_available_forces_the_load():
|
|
"""ensure_available must trigger model loading."""
|
|
loaded = []
|
|
|
|
RealAutoModel = sys.modules["transformers"].AutoModel
|
|
|
|
@classmethod
|
|
def recording_fp(cls, model_id, **kwargs):
|
|
loaded.append(model_id)
|
|
return cls()
|
|
|
|
RealAutoModel.from_pretrained = recording_fp
|
|
|
|
local_encoder.ensure_available("test-model", "cpu")
|
|
assert loaded == ["test-model"]
|
|
|
|
|
|
def test_ensure_available_propagates_missing_dependency(monkeypatch):
|
|
"""Missing transformers raises ImportError from ensure_available."""
|
|
monkeypatch.delitem(sys.modules, "transformers", raising=False)
|
|
|
|
import builtins
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *args, **kwargs):
|
|
if name == "transformers":
|
|
raise ImportError("No module named 'transformers'")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
|
|
# Restore torch so the code path reaches the transformers import
|
|
with pytest.raises(ImportError, match="pip install -r requirements-encoder.txt"):
|
|
local_encoder.ensure_available("stub-model", "cpu")
|
|
|
|
|
|
def test_a_real_failure_other_than_missing_dependency_propagates_unmodified():
|
|
"""A RuntimeError from from_pretrained must pass through unchanged."""
|
|
|
|
class BlowsUpModel:
|
|
@classmethod
|
|
def from_pretrained(cls, model_id, **kwargs):
|
|
raise RuntimeError("model checkpoint not found")
|
|
|
|
def to(self, device_str):
|
|
return self
|
|
|
|
sys.modules["transformers"].AutoModel = BlowsUpModel
|
|
|
|
with pytest.raises(RuntimeError, match="model checkpoint not found"):
|
|
local_encoder.classify_zero_shot(
|
|
"x", ["a"], model_id="bad-model", device="cpu"
|
|
)
|
|
|
|
|
|
def test_confidence_is_in_range_zero_to_one():
|
|
"""Confidence from classify_zero_shot must always be in [0, 1]."""
|
|
label, score = local_encoder.classify_zero_shot(
|
|
"some text", ["a", "b"], model_id="m", device="cpu",
|
|
)
|
|
assert isinstance(label, str)
|
|
assert 0.0 <= score <= 1.0, f"confidence {score} is outside [0, 1]"
|
|
|
|
|
|
def test_desc_cache_is_populated_on_first_call():
|
|
"""After the first classify call, the description cache is populated
|
|
for the (model_id, device) pair."""
|
|
assert local_encoder._desc_cache.get(("m", "cpu")) is None
|
|
|
|
local_encoder.classify_zero_shot("hello", ["a"], model_id="m", device="cpu")
|
|
|
|
cached = local_encoder._desc_cache.get(("m", "cpu"))
|
|
assert cached is not None
|
|
assert "a" in cached
|
|
|
|
|
|
def test_desc_cache_extended_for_new_categories():
|
|
"""Calling with a new category extends the existing cache rather than
|
|
rebuilding from scratch."""
|
|
local_encoder.classify_zero_shot(
|
|
"hello", ["coding_general"], model_id="m", device="cpu",
|
|
)
|
|
cached_before = local_encoder._desc_cache.get(("m", "cpu"))
|
|
assert cached_before is not None
|
|
assert "coding_general" in cached_before
|
|
assert "debugging" not in cached_before
|
|
|
|
local_encoder.classify_zero_shot(
|
|
"fix bug", ["coding_general", "debugging"], model_id="m", device="cpu",
|
|
)
|
|
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
|
|
|
|
|
|
# ── 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
|
|
|
|
|
|
# ── Per-model pooling strategy + query prefix (Todo 2) ──────────────────
|
|
#
|
|
# Tests for the CLS/mean pooling dispatch and per-model query prefix
|
|
# reading. The pooling strategy is stored in _model_pooling_strategies at
|
|
# model-load time; _mean_pool dispatches to _pool_embeddings based on it.
|
|
# The query prefix is stored in _model_query_prefixes and used in
|
|
# classify_zero_shot.
|
|
|
|
|
|
class _SpyTokenizer:
|
|
"""Fake tokenizer that records the last text passed to ``__call__``."""
|
|
last_text: str = ""
|
|
model_max_length = 512
|
|
truncation_side = "right"
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, model_id, **kwargs):
|
|
return cls()
|
|
|
|
def __call__(
|
|
self, text, add_special_tokens=True, truncation=False,
|
|
padding=True, return_tensors=None, max_length=None, **_kw,
|
|
):
|
|
_SpyTokenizer.last_text = str(text)
|
|
ids = [len(w) for w in str(text).split()]
|
|
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):
|
|
return " ".join(str(i) for i in ids)
|
|
|
|
|
|
class _DummyForwardModel:
|
|
"""Minimal model stub that returns a known last_hidden_state."""
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, model_id, **kwargs):
|
|
return cls()
|
|
|
|
def to(self, device_str):
|
|
return self
|
|
|
|
def __call__(self, **kwargs):
|
|
return types.SimpleNamespace(
|
|
last_hidden_state=MockTensor([[1.0, 2.0, 3.0]])
|
|
)
|
|
|
|
|
|
def test_cls_pooling_returns_first_token():
|
|
"""CLS pooling returns the first token of each sequence."""
|
|
local_encoder._model_pooling_strategies["cls-model"] = "cls"
|
|
embeddings = MockTensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
|
|
mask = MockTensor([[1, 1, 1], [1, 1, 1]])
|
|
result = local_encoder._mean_pool(embeddings, mask, "cls-model")
|
|
# CLS: first token from each row → [[1.0], [4.0]]
|
|
assert result._data == [[1.0], [4.0]]
|
|
|
|
|
|
def test_mean_pool_fallback_when_strategy_dict_empty():
|
|
"""When no strategy is recorded, _mean_pool defaults to mean pooling."""
|
|
local_encoder._model_pooling_strategies.clear()
|
|
embeddings = MockTensor([[2.0, 4.0, 6.0]])
|
|
mask = MockTensor([[1, 1, 1]])
|
|
result = local_encoder._mean_pool(embeddings, mask, "unknown-model")
|
|
# Mean of [2.0, 4.0, 6.0] = 4.0
|
|
assert abs(result._data[0][0] - 4.0) < 1e-9
|
|
|
|
|
|
def test_mean_pool_fallback_when_strategy_is_mean():
|
|
"""When strategy is explicitly 'mean', mean pooling is used."""
|
|
local_encoder._model_pooling_strategies["mean-model"] = "mean"
|
|
embeddings = MockTensor([[2.0, 4.0, 6.0]])
|
|
mask = MockTensor([[1, 1, 1]])
|
|
result = local_encoder._mean_pool(embeddings, mask, "mean-model")
|
|
assert abs(result._data[0][0] - 4.0) < 1e-9
|
|
|
|
|
|
def test_query_prefix_uses_per_model_prefix_when_set():
|
|
"""classify_zero_shot applies the prefix from _model_query_prefixes."""
|
|
sys.modules["transformers"].AutoTokenizer = _SpyTokenizer
|
|
sys.modules["transformers"].AutoModel = _DummyForwardModel
|
|
local_encoder._model_cache.clear()
|
|
|
|
# First classify call loads the model; _read_query_prefix returns ""
|
|
# (the mock transformers has no real utils.hub.cached_file), so the
|
|
# stored prefix is empty at this point.
|
|
local_encoder.classify_zero_shot(
|
|
"dummy", ["a"], model_id="prefix-model", device="cpu",
|
|
)
|
|
# Now overwrite the prefix — the model is cached so the next call
|
|
# will skip _load_model_tokenizer and use this value.
|
|
local_encoder._model_query_prefixes["prefix-model"] = "CUSTOM_PREFIX: "
|
|
_SpyTokenizer.last_text = ""
|
|
|
|
local_encoder.classify_zero_shot(
|
|
"test task", ["a"], model_id="prefix-model", device="cpu",
|
|
)
|
|
assert _SpyTokenizer.last_text.startswith("CUSTOM_PREFIX: ")
|
|
|
|
|
|
def test_query_prefix_falls_back_to_bge_instruction():
|
|
"""When _model_query_prefixes has no prefix for a model,
|
|
_BGE_QUERY_INSTRUCTION is used as the fallback."""
|
|
local_encoder._model_query_prefixes.clear()
|
|
sys.modules["transformers"].AutoTokenizer = _SpyTokenizer
|
|
sys.modules["transformers"].AutoModel = _DummyForwardModel
|
|
local_encoder._model_cache.clear()
|
|
|
|
_SpyTokenizer.last_text = ""
|
|
local_encoder.classify_zero_shot(
|
|
"test task", ["a"], model_id="no-prefix-model", device="cpu",
|
|
)
|
|
assert _SpyTokenizer.last_text.startswith(
|
|
local_encoder._BGE_QUERY_INSTRUCTION
|
|
)
|
|
|
|
|
|
def test_mean_pool_requires_model_id():
|
|
"""_mean_pool requires model_id — callers always have it in scope."""
|
|
local_encoder._model_pooling_strategies.clear()
|
|
embeddings = MockTensor([[3.0, 6.0, 9.0]])
|
|
mask = MockTensor([[1, 1, 1]])
|
|
result = local_encoder._mean_pool(embeddings, mask, "some-model")
|
|
assert abs(result._data[0][0] - 6.0) < 1e-9
|
|
|
|
|
|
def test_mean_pool_dispatches_by_model_id_not_load_order():
|
|
"""Pooling strategy follows model_id, not insertion order.
|
|
|
|
Two models with different strategies loaded into
|
|
_model_pooling_strategies: model A = cls (inserted first),
|
|
model B = mean (inserted second). Pooling model A must use
|
|
cls, pooling model B must use mean, regardless of what was
|
|
inserted last.
|
|
"""
|
|
local_encoder._model_pooling_strategies.clear()
|
|
# Load model A (cls) first, model B (mean) second
|
|
local_encoder._model_pooling_strategies["model-a"] = "cls"
|
|
local_encoder._model_pooling_strategies["model-b"] = "mean"
|
|
embeddings = MockTensor([[10.0, 20.0]])
|
|
mask = MockTensor([[1, 1]])
|
|
result_a = local_encoder._mean_pool(embeddings, mask, "model-a")
|
|
result_b = local_encoder._mean_pool(embeddings, mask, "model-b")
|
|
# Model A uses CLS → first token = 10.0
|
|
assert abs(result_a._data[0][0] - 10.0) < 1e-9
|
|
# Model B uses mean → mean of [10.0, 20.0] = 15.0
|
|
assert abs(result_b._data[0][0] - 15.0) < 1e-9
|
|
# Reverse order: load B first, A second
|
|
local_encoder._model_pooling_strategies.clear()
|
|
local_encoder._model_pooling_strategies["model-b"] = "mean"
|
|
local_encoder._model_pooling_strategies["model-a"] = "cls"
|
|
result_a = local_encoder._mean_pool(embeddings, mask, "model-a")
|
|
result_b = local_encoder._mean_pool(embeddings, mask, "model-b")
|
|
assert abs(result_a._data[0][0] - 10.0) < 1e-9
|
|
assert abs(result_b._data[0][0] - 15.0) < 1e-9
|
|
|
|
|
|
def test_pool_embeddings_cls_direct():
|
|
"""_pool_embeddings with strategy='cls' returns first token."""
|
|
embeddings = MockTensor([[10.0, 20.0], [30.0, 40.0]])
|
|
mask = MockTensor([[1, 1], [1, 1]])
|
|
result = local_encoder._pool_embeddings(embeddings, mask, "cls")
|
|
assert result._data == [[10.0], [30.0]]
|
|
|
|
|
|
def test_pool_embeddings_mean_direct():
|
|
"""_pool_embeddings with strategy='mean' returns attention-weighted mean."""
|
|
embeddings = MockTensor([[1.0, 2.0, 3.0]])
|
|
mask = MockTensor([[1, 1, 1]])
|
|
result = local_encoder._pool_embeddings(embeddings, mask, "mean")
|
|
assert abs(result._data[0][0] - 2.0) < 1e-9 |