diff --git a/packages/sie_server/bundles/default.yaml b/packages/sie_server/bundles/default.yaml index e4f4a81c9..e495c0339 100644 --- a/packages/sie_server/bundles/default.yaml +++ b/packages/sie_server/bundles/default.yaml @@ -62,7 +62,8 @@ deps: # colbert_rotary_flash, jina_flash_cross_encoder einops: '>=0.8' # sentence_transformer; native causal CrossEncoder module format requires 5.4.1 - sentence-transformers: '>=5.4.1,<6' + # The local-code trust boundary is validated against this exact release. + sentence-transformers: '==5.7.0' # bge_m3_flag FlagEmbedding: '>=1.3' # grounding_dino. EXACT, not a range. `uv run --with-requirements` resolves the diff --git a/packages/sie_server/bundles/transformers5.yaml b/packages/sie_server/bundles/transformers5.yaml index 3beca954a..6ac39b7d0 100644 --- a/packages/sie_server/bundles/transformers5.yaml +++ b/packages/sie_server/bundles/transformers5.yaml @@ -28,9 +28,9 @@ deps: # installs already resolved to the latest 5.x, so the floor is the # guarantee, not a version change in practice. transformers: '>=5.17,<6' - # Multimodal SparseEncoder (image documents) landed on the 5.6 line - # (st_sparse_vision adapter). - sentence-transformers: '>=5.6,<6' + # Includes multimodal SparseEncoder (image documents), with the local-code + # trust boundary validated against this exact release. + sentence-transformers: '==5.7.0' # Pin torch — flash-attn's unpinned torch dep can pull in 2.10+ which breaks torchvision torch: '>=2.9,<2.10' # Keep the compiled torchvision operators paired with torch 2.9.x. A broad diff --git a/packages/sie_server/pyproject.toml b/packages/sie_server/pyproject.toml index 71aabc98a..b8d9c8d7d 100644 --- a/packages/sie_server/pyproject.toml +++ b/packages/sie_server/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ # the base venv can't silently jump to a transformers 6 major. "transformers>=4.57,<6", "tokenizers>=0.20,<1", - "sentence-transformers>=5.0,<6", + "sentence-transformers==5.7.0", "safetensors>=0.4,<1", # Artifact resolution (snapshot_download / hf_hub_download / scan_cache_dir) # is imported directly by the adapters and the disk cache, and `requests` by diff --git a/packages/sie_server/src/sie_server/adapters/_sentence_transformer_trust.py b/packages/sie_server/src/sie_server/adapters/_sentence_transformer_trust.py new file mode 100644 index 000000000..e5d12b423 --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/_sentence_transformer_trust.py @@ -0,0 +1,172 @@ +"""Load-only compatibility boundary for Sentence Transformers 5.x custom code. + +The supported 5.x releases implicitly trust local model directories. Enforce an +explicit opt-in at the actual class-resolution sites without changing library +globals, native model classes, or their inference methods. +""" + +from __future__ import annotations + +import inspect +from collections.abc import Mapping +from types import FunctionType, MethodType, ModuleType +from typing import Any + +import sentence_transformers +from sentence_transformers import CrossEncoder, SentenceTransformer, SparseEncoder + +_SUPPORTED_VERSIONS = frozenset({"5.4.1", "5.6.0", "5.7.0"}) +_ROOT_CLASSES = (SentenceTransformer, SparseEncoder, CrossEncoder) +_RESOLVER = "_load_module_class_from_ref" +_REFUSAL = "Custom Sentence Transformers code requires trust_remote_code=True." +_UNSUPPORTED = "Unsupported Sentence Transformers loading implementation." +_NESTED_IMPORTS = { + "5.4.1": ("import_from_string", "import_from_string"), + "5.6.0": ("import_module_class", "import_from_string"), + "5.7.0": ("import_module_class", "import_module_class"), +} +_KWARG_GROUPS = ( + "model_kwargs", + "processor_kwargs", + "config_kwargs", + "tokenizer_kwargs", + "automodel_args", + "model_args", + "tokenizer_args", + "config_args", +) + + +def _native_references(exports: tuple[ModuleType, ...], base: type, aliases: dict[str, str]) -> dict[str, type]: + """Build a closed name map from installed, trusted library exports. + + Model-supplied names never reach an import function, even for malformed + names resembling a native namespace. Legacy aliases map to the same class. + """ + references: dict[str, type] = {} + for namespace in exports: + for name in namespace.__all__: + value = getattr(namespace, name) + if not inspect.isclass(value) or not issubclass(value, base) or inspect.isabstract(value): + continue + references[f"{namespace.__name__}.{name}"] = value + references[f"{value.__module__}.{value.__name__}"] = value + for old, new in aliases.items(): + for reference, value in tuple(references.items()): + if reference.startswith(new + "."): + references[old + reference[len(new) :]] = value + # Historical import_from_string also accepts a module alias whose + # final component is the exported class name (models.Pooling). + if new == value.__module__ and old.rsplit(".", 1)[-1] == value.__name__: + references[old] = value + return references + + +class _NativeLoadBoundary: + """Private state used only while one native root constructor executes.""" + + def __init__(self) -> None: + # This private layout is optional across the public ST5 range. Import + # it only after selecting a supported untrusted path, never for True. + from sentence_transformers.base import modules as base_modules + from sentence_transformers.cross_encoder import modules as cross_modules + from sentence_transformers.sentence_transformer import modules as sentence_modules + from sentence_transformers.sentence_transformer.modules import tokenizer as tokenizers + from sentence_transformers.sparse_encoder import modules as sparse_modules + from sentence_transformers.util.deprecated_import import DEPRECATED_MODULE_PATHS + + self.router = base_modules.Router + self.word_embeddings = sentence_modules.WordEmbeddings + self.modules = _native_references( + (base_modules, sentence_modules, sparse_modules, cross_modules), + base_modules.Module, + DEPRECATED_MODULE_PATHS, + ) + self.tokenizers = _native_references((tokenizers,), tokenizers.WordTokenizer, DEPRECATED_MODULE_PATHS) + self.facades: dict[type, type] = {} + self.import_names = _NESTED_IMPORTS[sentence_transformers.__version__] + + def resolve_module(self, class_ref: object, *_args: object, **_kwargs: object) -> type: + if not isinstance(class_ref, str) or class_ref not in self.modules: + raise ValueError(_REFUSAL) + native = self.modules[class_ref] + if issubclass(native, self.router) or native is self.word_embeddings: + return self._facade(native) + return native + + def resolve_tokenizer(self, class_ref: object, *_args: object, **_kwargs: object) -> type: + if not isinstance(class_ref, str) or class_ref not in self.tokenizers: + raise ValueError(_REFUSAL) + return self.tokenizers[class_ref] + + def resolve_root(self, _model: object, class_ref: object, *_args: object, **_kwargs: object) -> type: + return self.resolve_module(class_ref) + + def _facade(self, native: type) -> type: + if native in self.facades: + return self.facades[native] + method = getattr(native, "load", None) + if not isinstance(method, MethodType): + raise ValueError(_UNSUPPORTED) + loader = method.__func__ + import_name = self.import_names[1 if native is self.word_embeddings else 0] + if not isinstance(loader, FunctionType) or import_name not in loader.__code__.co_names: + raise ValueError(_UNSUPPORTED) + # Copy globals for this function only. The original library function, + # class and globals remain unchanged, including during concurrent loads. + globals_copy = dict(loader.__globals__) + globals_copy[import_name] = self.resolve_tokenizer if native is self.word_embeddings else self.resolve_module + guarded = FunctionType(loader.__code__, globals_copy, loader.__name__, loader.__defaults__, loader.__closure__) + guarded.__kwdefaults__ = loader.__kwdefaults__ + guarded.__annotations__ = loader.__annotations__ + # Bind cls to the ORIGINAL native class. The facade is never a model + # instance; only its load entry point is presented to upstream code. + facade = type(native.__name__, (native,), {"load": staticmethod(MethodType(guarded, native))}) + self.facades[native] = facade + return facade + + +def load_sentence_transformer[T: (SentenceTransformer, SparseEncoder, CrossEncoder)]( + model_class: type[T], + model_name_or_path: str, + *, + trust_remote_code: bool | None = False, + **kwargs: Any, +) -> T: + """Construct the native model, enforcing explicit trust for custom classes. + + This compatibility helper intentionally supports only the inspected 5.x + releases. An unknown version fails before an untrusted model is loaded. + Explicit True preserves the original constructor path exactly. + """ + if trust_remote_code is True: + return model_class(model_name_or_path, trust_remote_code=True, **kwargs) + if trust_remote_code is not False and trust_remote_code is not None: + raise ValueError(_REFUSAL) + if sentence_transformers.__version__ not in _SUPPORTED_VERSIONS or model_class not in _ROOT_CLASSES: + raise ValueError(_UNSUPPORTED) + # The upstream decorators rename these aliases after this boundary. Check + # every supported spelling before that happens, including canonical/alias + # collisions. Do not mutate the caller's dictionaries. + for name in _KWARG_GROUPS: + nested = kwargs.get(name) + if nested is not None and not isinstance(nested, Mapping): + raise ValueError(_REFUSAL) + if isinstance(nested, Mapping) and "trust_remote_code" in nested: + nested_trust = nested["trust_remote_code"] + if nested_trust is not False and nested_trust is not None: + raise ValueError(_REFUSAL) + kwargs[name] = {**nested, "trust_remote_code": False} + + boundary = _NativeLoadBoundary() + model = model_class.__new__(model_class) + # nn.Module is not initialized yet; use object assignment deliberately. + # Native nn.Module initialization preserves unrelated instance attributes. + object.__setattr__(model, _RESOLVER, MethodType(boundary.resolve_root, model)) + try: + model_class.__init__(model, model_name_or_path, trust_remote_code=False, **kwargs) + return model + finally: + # Remove even if construction/nn.Module initialization failed. Never + # leave a bound method cycle or an inference-time compatibility wrapper. + object.__getattribute__(model, "__dict__").pop(_RESOLVER, None) diff --git a/packages/sie_server/src/sie_server/adapters/cross_encoder/__init__.py b/packages/sie_server/src/sie_server/adapters/cross_encoder/__init__.py index 1be489341..9d7a2d4d5 100644 --- a/packages/sie_server/src/sie_server/adapters/cross_encoder/__init__.py +++ b/packages/sie_server/src/sie_server/adapters/cross_encoder/__init__.py @@ -25,6 +25,7 @@ from sentence_transformers import CrossEncoder from sie_server.adapters._base_adapter import BaseAdapter +from sie_server.adapters._sentence_transformer_trust import load_sentence_transformer from sie_server.adapters._spec import AdapterSpec from sie_server.adapters._types import ERR_NOT_LOADED, ERR_REQUIRES_TEXT, AttnImplementation, ComputePrecision from sie_server.adapters.cross_encoder._cpu_alignment import normalize_bert_cpu_storage @@ -127,7 +128,8 @@ def load(self, device: str) -> None: model_kwargs["dtype"] = torch.bfloat16 # float32 -> CrossEncoder default, no dtype kwarg needed - model = CrossEncoder( + model = load_sentence_transformer( + CrossEncoder, self._model_name_or_path, device=device, trust_remote_code=self._trust_remote_code, diff --git a/packages/sie_server/src/sie_server/adapters/cross_encoder/_cpu_alignment.py b/packages/sie_server/src/sie_server/adapters/cross_encoder/_cpu_alignment.py index 0ba764177..eabd75577 100644 --- a/packages/sie_server/src/sie_server/adapters/cross_encoder/_cpu_alignment.py +++ b/packages/sie_server/src/sie_server/adapters/cross_encoder/_cpu_alignment.py @@ -10,7 +10,7 @@ from transformers import BertForSequenceClassification -def normalize_bert_cpu_storage(model: torch.nn.Module) -> None: +def normalize_bert_cpu_storage(model: torch.nn.Module | None) -> None: """Copy unsafe standard BERT float32 storage once, retaining tensor objects. The caller restricts this optional compatibility path to Darwin arm64 CPU. diff --git a/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py b/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py index 7e6c355b4..e95818e46 100644 --- a/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py +++ b/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py @@ -113,12 +113,12 @@ def runtime_config(self) -> LoraRuntimeConfig: def bind(self, outer: CrossEncoder) -> None: model = outer.model - encoder_prefix = _ENCODER_PREFIXES.get(type(model)) - if encoder_prefix is None or getattr(model, "is_quantized", False): + if model is None or type(model) not in _ENCODER_PREFIXES or getattr(model, "is_quantized", False): raise ValueError( "CrossEncoder LoRA requires a standard, unquantized BertForSequenceClassification " "or XLMRobertaForSequenceClassification" ) + encoder_prefix = _ENCODER_PREFIXES[type(model)] if getattr(model, "peft_config", None) or getattr(model, "_hf_peft_config_loaded", False): raise ValueError("The cross-encoder must start from a base checkpoint without attached adapters") self._outer, self._model = outer, model diff --git a/packages/sie_server/src/sie_server/adapters/native_causal_cross_encoder/adapter.py b/packages/sie_server/src/sie_server/adapters/native_causal_cross_encoder/adapter.py index 77573b623..9ec2a09e9 100644 --- a/packages/sie_server/src/sie_server/adapters/native_causal_cross_encoder/adapter.py +++ b/packages/sie_server/src/sie_server/adapters/native_causal_cross_encoder/adapter.py @@ -125,10 +125,13 @@ def load(self, device: str) -> None: from sentence_transformers.cross_encoder.modules import LogitScore from sentence_transformers.util import batch_to_device + from sie_server.adapters._sentence_transformer_trust import load_sentence_transformer + self._require_saved_modules() precision = self._compute_precision or ("bfloat16" if device.startswith("cuda") else "float32") dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}[precision] - model = CrossEncoder( + model = load_sentence_transformer( + CrossEncoder, self._model_name_or_path, revision=self._revision, device=device, diff --git a/packages/sie_server/src/sie_server/adapters/sentence_transformer/__init__.py b/packages/sie_server/src/sie_server/adapters/sentence_transformer/__init__.py index 226151f52..16a6112ff 100644 --- a/packages/sie_server/src/sie_server/adapters/sentence_transformer/__init__.py +++ b/packages/sie_server/src/sie_server/adapters/sentence_transformer/__init__.py @@ -15,6 +15,7 @@ from sentence_transformers import SentenceTransformer, SparseEncoder from sie_server.adapters._base_adapter import BaseAdapter +from sie_server.adapters._sentence_transformer_trust import load_sentence_transformer from sie_server.adapters._spec import AdapterSpec from sie_server.adapters._types import ERR_NOT_LOADED, ERR_REQUIRES_TEXT, ComputePrecision from sie_server.core.inference_output import EncodeOutput, SparseVector @@ -88,7 +89,7 @@ def load(self, device: str) -> None: # Omit revision unless pinned so the library default (main) is unchanged. if self._revision is not None: load_kwargs["revision"] = self._revision - self._model = SentenceTransformer(self._model_name_or_path, **load_kwargs) + self._model = load_sentence_transformer(SentenceTransformer, self._model_name_or_path, **load_kwargs) if self._max_seq_length is not None: self._model.max_seq_length = self._max_seq_length @@ -249,7 +250,7 @@ def load(self, device: str) -> None: # Omit revision unless pinned so the library default (main) is unchanged. if self._revision is not None: load_kwargs["revision"] = self._revision - self._model = SparseEncoder(self._model_name_or_path, **load_kwargs) + self._model = load_sentence_transformer(SparseEncoder, self._model_name_or_path, **load_kwargs) if self._max_seq_length is not None: self._model.max_seq_length = self._max_seq_length diff --git a/packages/sie_server/src/sie_server/adapters/st_sparse_vision/adapter.py b/packages/sie_server/src/sie_server/adapters/st_sparse_vision/adapter.py index 853eeb337..064e0b08e 100644 --- a/packages/sie_server/src/sie_server/adapters/st_sparse_vision/adapter.py +++ b/packages/sie_server/src/sie_server/adapters/st_sparse_vision/adapter.py @@ -4,9 +4,9 @@ documents in addition to text (visual-SPLADE family, e.g. naver/v-splade-quality): text or page-image in, sparse vocab-sized vector out. -Requires the ``transformers5`` bundle (sentence-transformers >=5.6 / -transformers >=5.14); the multimodal ``SparseEncoder`` routing does not exist -on the sentence-transformers 5.0 line the default bundle pins. +Requires the ``transformers5`` bundle for the checkpoint's Transformers >=5.14 +modeling support. Both this bundle and the default bundle pin +sentence-transformers 5.7.0, which includes multimodal ``SparseEncoder`` routing. """ from __future__ import annotations @@ -21,6 +21,7 @@ from sentence_transformers import SparseEncoder from sie_server.adapters._base_adapter import BaseAdapter +from sie_server.adapters._sentence_transformer_trust import load_sentence_transformer from sie_server.adapters._spec import AdapterSpec from sie_server.adapters._types import ERR_NOT_LOADED, ComputePrecision from sie_server.adapters._utils import validate_output_types @@ -106,7 +107,7 @@ def load(self, device: str) -> None: load_kwargs["revision"] = self._revision if self._trust_remote_code: self._register_remote_code_dir() - self._model = SparseEncoder(self._model_name_or_path, **load_kwargs) + self._model = load_sentence_transformer(SparseEncoder, self._model_name_or_path, **load_kwargs) if self._max_seq_length is not None: self._model.max_seq_length = self._max_seq_length diff --git a/packages/sie_server/tests/adapters/test_cross_encoder_cpu_alignment.py b/packages/sie_server/tests/adapters/test_cross_encoder_cpu_alignment.py index d091623dc..5cda3f00d 100644 --- a/packages/sie_server/tests/adapters/test_cross_encoder_cpu_alignment.py +++ b/packages/sie_server/tests/adapters/test_cross_encoder_cpu_alignment.py @@ -225,7 +225,7 @@ def test_off_platform_load_does_not_invoke_alignment(monkeypatch, system, machin cache_calls = [] monkeypatch.setattr(cross_encoder.sys, "platform", system) monkeypatch.setattr(cross_encoder.platform, "machine", lambda: machine) - monkeypatch.setattr(cross_encoder, "CrossEncoder", lambda *args, **kwargs: outer) + monkeypatch.setattr(cross_encoder, "load_sentence_transformer", lambda *args, **kwargs: outer) # Both the model and its accelerator are simulated, including cache cleanup. monkeypatch.setattr(torch.cuda, "empty_cache", lambda: cache_calls.append("cuda")) monkeypatch.setattr(torch.mps, "empty_cache", lambda: cache_calls.append("mps")) @@ -250,7 +250,7 @@ def test_failed_normalization_never_publishes_outer(model, monkeypatch): outer = SimpleNamespace(model=model) monkeypatch.setattr(cross_encoder.sys, "platform", "darwin") monkeypatch.setattr(cross_encoder.platform, "machine", lambda: "arm64") - monkeypatch.setattr(cross_encoder, "CrossEncoder", lambda *args, **kwargs: outer) + monkeypatch.setattr(cross_encoder, "load_sentence_transformer", lambda *args, **kwargs: outer) adapter = cross_encoder.CrossEncoderAdapter("unused") with pytest.raises(ValueError, match="storage overlaps"): adapter.load("cpu") @@ -263,7 +263,9 @@ def test_failed_copy_never_publishes_outer(model, monkeypatch): before = pointers(model) monkeypatch.setattr(cross_encoder.sys, "platform", "darwin") monkeypatch.setattr(cross_encoder.platform, "machine", lambda: "arm64") - monkeypatch.setattr(cross_encoder, "CrossEncoder", lambda *args, **kwargs: SimpleNamespace(model=model)) + monkeypatch.setattr( + cross_encoder, "load_sentence_transformer", lambda *args, **kwargs: SimpleNamespace(model=model) + ) def fail_copy(*args, **kwargs): raise MemoryError("controlled load copy failure") @@ -305,7 +307,7 @@ def factory(*args, **kwargs): monkeypatch.setattr(cross_encoder.sys, "platform", "darwin") monkeypatch.setattr(cross_encoder.platform, "machine", lambda: "arm64") - monkeypatch.setattr(cross_encoder, "CrossEncoder", factory) + monkeypatch.setattr(cross_encoder, "load_sentence_transformer", factory) adapter = cross_encoder.CrossEncoderAdapter("unused") def attempt(): diff --git a/packages/sie_server/tests/adapters/test_native_causal_cross_encoder.py b/packages/sie_server/tests/adapters/test_native_causal_cross_encoder.py index 3b83bb56f..79040c8f1 100644 --- a/packages/sie_server/tests/adapters/test_native_causal_cross_encoder.py +++ b/packages/sie_server/tests/adapters/test_native_causal_cross_encoder.py @@ -13,7 +13,6 @@ import numpy as np import pytest -import sentence_transformers import torch from jinja2.sandbox import ImmutableSandboxedEnvironment from sentence_transformers.base import modules as native_modules @@ -141,7 +140,10 @@ def construct(*args: Any, **kwargs: Any) -> ToyCrossEncoder: return model loader = Mock(side_effect=construct) - monkeypatch.setattr(sentence_transformers, "CrossEncoder", loader) + monkeypatch.setattr( + "sie_server.adapters._sentence_transformer_trust.load_sentence_transformer", + lambda model_class, *args, **kwargs: loader(*args, **kwargs), + ) monkeypatch.setattr(native_modules, "Transformer", ToyTransformer) modules_file = tmp_path / "modules.json" modules_file.write_text( diff --git a/packages/sie_server/tests/adapters/test_sentence_transformer_trust.py b/packages/sie_server/tests/adapters/test_sentence_transformer_trust.py new file mode 100644 index 000000000..b21cd9d2d --- /dev/null +++ b/packages/sie_server/tests/adapters/test_sentence_transformer_trust.py @@ -0,0 +1,689 @@ +"""Real local-model controls for the load-only ST custom-code boundary.""" + +from __future__ import annotations + +import gc +import importlib +import json +import sys +import threading +import tomllib +import types +import uuid +import weakref +from collections import UserDict +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np +import pytest +import sentence_transformers +import torch +import yaml +from packaging.requirements import Requirement +from sentence_transformers import CrossEncoder, SentenceTransformer, SparseEncoder +from sentence_transformers.base.modules import Dense, Module, Router +from sentence_transformers.sentence_transformer.modules import Pooling, WordEmbeddings +from sentence_transformers.sentence_transformer.modules.tokenizer import WhitespaceTokenizer +from sie_server.adapters import _sentence_transformer_trust as trust +from sie_server.adapters.cross_encoder import CrossEncoderAdapter +from sie_server.adapters.cross_encoder.runtime_lora import CrossEncoderLoraAdapter +from sie_server.adapters.native_causal_cross_encoder.adapter import NativeCausalCrossEncoderAdapter +from sie_server.adapters.sentence_transformer import SentenceTransformerDenseAdapter, SentenceTransformerSparseAdapter +from sie_server.adapters.st_sparse_vision.adapter import SparseEncoderVisionAdapter +from transformers import ( + BertConfig, + BertForMaskedLM, + BertForSequenceClassification, + BertModel, + BertTokenizerFast, + dynamic_module_utils, +) + + +class _ControlledAutoImportMarker(BaseException): + """Survive AutoProcessor's broad Exception fallback to prove code execution.""" + + +def test_dependency_pins_match_supported_loading_boundary() -> None: + root = Path(__file__).resolve().parents[4] + server = root / "packages" / "sie_server" + project = tomllib.loads((server / "pyproject.toml").read_text()) + requirement = next( + Requirement(value) + for value in project["project"]["dependencies"] + if Requirement(value).name == "sentence-transformers" + ) + (specifier,) = requirement.specifier + assert specifier.operator == "==" + assert specifier.version in trust._SUPPORTED_VERSIONS + for bundle in ("default", "transformers5"): + config = yaml.safe_load((server / "bundles" / f"{bundle}.yaml").read_text()) + assert config["deps"]["sentence-transformers"] == str(requirement.specifier) + lock = tomllib.loads((root / "uv.lock").read_text()) + assert [package["version"] for package in lock["package"] if package["name"] == "sentence-transformers"] == [ + specifier.version + ] + + +@pytest.fixture(autouse=True) +def _isolated_dynamic_module_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + cache = str(tmp_path / "hf-modules") + monkeypatch.setenv("HF_MODULES_CACHE", cache) + monkeypatch.setattr(dynamic_module_utils, "HF_MODULES_CACHE", cache) + # Transformers appends this cache to sys.path when loading custom code. + monkeypatch.setattr(sys, "path", list(sys.path)) + control = types.ModuleType("_sie_st_trust_auto_import_control") + control.__dict__["ControlledAutoImportMarker"] = _ControlledAutoImportMarker + monkeypatch.setitem(sys.modules, control.__name__, control) + + +def _json(path: Path, data: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data)) + + +def _model(path: Path, kind: type = SentenceTransformer, module_type: object = "modeling_marker.Marker") -> None: + _json(path / "config_sentence_transformers.json", {"model_type": kind.__name__}) + _json(path / "modules.json", [{"idx": 0, "name": "0", "path": "", "type": module_type}]) + + +def _marker(path: Path, monkeypatch: pytest.MonkeyPatch, *, pause: str = "") -> tuple[str, Path]: + path.mkdir(parents=True, exist_ok=True) + name = "modeling_marker_" + uuid.uuid4().hex + marker = path / "executed.txt" + (path / f"{name}.py").write_text( + "from pathlib import Path\n" + + pause + + f"Path({str(marker)!r}).write_text('controlled import')\n" + + "raise RuntimeError('controlled import marker')\nclass Marker:\n pass\n" + ) + monkeypatch.syspath_prepend(str(path)) + return f"{name}.Marker", marker + + +def _nested(path: Path, class_ref: str, depth: int, *, tokenizer: bool = False) -> None: + router = "sentence_transformers.models.Router" + word = "sentence_transformers.models.WordEmbeddings" + _model(path, module_type=router if depth else word) + current = path + for index in range(depth): + child = router if index < depth - 1 else word if tokenizer else class_ref + _json( + current / "router_config.json", + {"types": {"child": child}, "structure": {"query": ["child"]}, "parameters": {"default_route": "query"}}, + ) + current = current / "child" + if tokenizer: + _json(current / "wordembedding_config.json", {"tokenizer_class": class_ref}) + + +@pytest.mark.parametrize("model_class", [SentenceTransformer, SparseEncoder, CrossEncoder]) +@pytest.mark.parametrize("allowed", [False, None, True]) +def test_root_custom_module_requires_explicit_trust(tmp_path, monkeypatch, model_class, allowed): + reference, marker = _marker(tmp_path, monkeypatch) + _model(tmp_path, model_class, reference) + expected = RuntimeError if allowed else ValueError + with pytest.raises(expected, match="controlled import marker" if allowed else "requires trust_remote_code=True"): + trust.load_sentence_transformer(model_class, str(tmp_path), trust_remote_code=allowed, local_files_only=True) + assert marker.exists() is bool(allowed) + + +@pytest.mark.parametrize(("depth", "tokenizer"), [(1, False), (3, False), (0, True), (1, True), (3, True)]) +@pytest.mark.parametrize("allowed", [False, None, True]) +def test_recursive_router_and_tokenizer_gate(tmp_path, monkeypatch, depth, tokenizer, allowed): + reference, marker = _marker(tmp_path, monkeypatch) + _nested(tmp_path, reference, depth, tokenizer=tokenizer) + expected = RuntimeError if allowed else ValueError + with pytest.raises(expected, match="controlled import marker" if allowed else "requires trust_remote_code=True"): + trust.load_sentence_transformer( + SentenceTransformer, str(tmp_path), trust_remote_code=allowed, local_files_only=True + ) + assert marker.exists() is bool(allowed) + + +@pytest.mark.parametrize( + "reference", + [ + None, + [], + {}, + "os.system", + "sentence_transformersx.models.Router", + "sentence_transformers.models..Router", + "sentence_transformers.models.Router/../../modeling_marker.Marker", + "sentence_transformers.models.Router:Marker", + "sentence_transformers.SentenceTransformer", + "sentence_transformers.util.import_from_string", + "sentence_transformers.models.Router\x00", + ], +) +def test_malformed_or_noncomponent_reference_never_imported(monkeypatch, reference): + boundary = trust._NativeLoadBoundary() + imported = [] + monkeypatch.setattr(importlib, "import_module", lambda *args, **kwargs: imported.append(args)) + with pytest.raises(ValueError, match="requires trust_remote_code=True"): + boundary.resolve_module(reference) + with pytest.raises(ValueError, match="requires trust_remote_code=True"): + boundary.resolve_tokenizer(reference) + assert imported == [] + + +@pytest.mark.parametrize("nested", trust._KWARG_GROUPS) +@pytest.mark.parametrize("allowed", [False, None]) +@pytest.mark.parametrize("mapping", [dict, UserDict]) +def test_nested_kwargs_cannot_elevate_trust(tmp_path, monkeypatch, nested, allowed, mapping): + reference, marker = _marker(tmp_path, monkeypatch) + _nested(tmp_path, reference, 2, tokenizer=True) + with pytest.raises(ValueError, match="requires trust_remote_code=True"): + trust.load_sentence_transformer( + SentenceTransformer, + str(tmp_path), + trust_remote_code=allowed, + local_files_only=True, + **{nested: mapping({"trust_remote_code": True})}, + ) + assert not marker.exists() + + +@pytest.mark.parametrize("nested", trust._KWARG_GROUPS) +@pytest.mark.parametrize("value", [[("trust_remote_code", True)], "trust_remote_code", 1]) +def test_nonmapping_nested_kwargs_fail_before_constructor(tmp_path, monkeypatch, nested, value): + monkeypatch.setattr(SentenceTransformer, "__init__", lambda *args, **kwargs: pytest.fail("constructor reached")) + with pytest.raises(ValueError, match="requires trust_remote_code=True"): + trust.load_sentence_transformer(SentenceTransformer, str(tmp_path), **{nested: value}) + + +@pytest.mark.parametrize("elevated", ["processor_kwargs", "tokenizer_kwargs"]) +def test_canonical_alias_collision_cannot_hide_elevated_trust(tmp_path, monkeypatch, elevated): + monkeypatch.setattr(SentenceTransformer, "__init__", lambda *args, **kwargs: pytest.fail("constructor reached")) + kwargs: dict[str, Mapping[str, bool | None]] = { + "processor_kwargs": {"trust_remote_code": False}, + "tokenizer_kwargs": {"trust_remote_code": None}, + } + kwargs[elevated] = UserDict({"trust_remote_code": True}) + with pytest.raises(ValueError, match="requires trust_remote_code=True"): + trust.load_sentence_transformer( + SentenceTransformer, + str(tmp_path), + processor_kwargs=kwargs["processor_kwargs"], + tokenizer_kwargs=kwargs["tokenizer_kwargs"], + ) + + +@pytest.mark.parametrize("allowed", [False, None]) +@pytest.mark.parametrize("nested", trust._KWARG_GROUPS) +def test_nested_none_is_denying_and_caller_mapping_unchanged(tmp_path, monkeypatch, nested, allowed): + original = UserDict({"trust_remote_code": None, "local_files_only": True}) + seen = [] + + def construct(model, path, *, trust_remote_code, **kwargs): + seen.append((trust_remote_code, kwargs[nested])) + assert kwargs[nested] is not original + + monkeypatch.setattr(SentenceTransformer, "__init__", construct) + loaded = trust.load_sentence_transformer( + SentenceTransformer, str(tmp_path), trust_remote_code=allowed, **{nested: original} + ) + assert seen == [(False, {"trust_remote_code": False, "local_files_only": True})] + assert original == {"trust_remote_code": None, "local_files_only": True} + assert trust._RESOLVER not in vars(loaded) + + +def _bert_checkpoint(path: Path, model_class: type, hf_class: type) -> None: + path.mkdir(parents=True, exist_ok=True) + vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "hello", "world"] + (path / "vocab.txt").write_text("\n".join(vocab) + "\n") + BertTokenizerFast(vocab_file=str(path / "vocab.txt")).save_pretrained(path) + config = BertConfig( + vocab_size=len(vocab), + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_hidden_layers=1, + num_labels=1, + ) + hf_class(config).save_pretrained(path) + + +def _auto_marker(path: Path, kind: str, model_class: type) -> Path: + name = "modeling_auto_marker_" + uuid.uuid4().hex + marker = path / "auto-executed.txt" + (path / f"{name}.py").write_text( + "from pathlib import Path\n" + "from _sie_st_trust_auto_import_control import ControlledAutoImportMarker\n" + f"Path({str(marker)!r}).write_text('controlled Auto import')\n" + "raise ControlledAutoImportMarker('controlled Auto import marker')\nclass Marker:\n pass\n" + ) + filename = "tokenizer_config.json" if kind == "tokenizer" else "config.json" + config_path = path / filename + config = json.loads(config_path.read_text()) + reference = f"{name}.Marker" + if kind == "tokenizer": + config["auto_map"] = {"AutoTokenizer": [reference, reference]} + elif kind == "config": + config["auto_map"] = {"AutoConfig": reference} + else: + auto_classes: dict[type, str] = { + SentenceTransformer: "AutoModel", + SparseEncoder: "AutoModelForMaskedLM", + CrossEncoder: "AutoModelForSequenceClassification", + } + auto = auto_classes[model_class] + config["auto_map"] = {auto: reference} + _json(config_path, config) + return marker + + +@pytest.mark.parametrize( + ("model_class", "hf_class", "alias"), + [ + (SentenceTransformer, BertModel, "tokenizer_kwargs"), + (SparseEncoder, BertForMaskedLM, "tokenizer_kwargs"), + (CrossEncoder, BertForSequenceClassification, "tokenizer_args"), + ], +) +@pytest.mark.parametrize("allowed", [False, None, True]) +@pytest.mark.parametrize("mapping", [dict, UserDict]) +def test_actual_auto_tokenizer_alias_cannot_elevate_trust(tmp_path, model_class, hf_class, alias, allowed, mapping): + _bert_checkpoint(tmp_path, model_class, hf_class) + marker = _auto_marker(tmp_path, "tokenizer", model_class) + with pytest.raises( + _ControlledAutoImportMarker if allowed is True else ValueError, + match="controlled Auto import marker" if allowed is True else "requires trust_remote_code=True", + ): + trust.load_sentence_transformer( + model_class, + str(tmp_path), + trust_remote_code=allowed, + device="cpu", + local_files_only=True, + **{alias: mapping({"trust_remote_code": True})}, + ) + assert marker.exists() is (allowed is True) + + +@pytest.mark.parametrize("kind", ["model", "config"]) +@pytest.mark.parametrize("allowed", [False, None, True]) +def test_actual_cross_encoder_auto_alias_cannot_elevate_trust(tmp_path, kind, allowed): + _bert_checkpoint(tmp_path, CrossEncoder, BertForSequenceClassification) + marker = _auto_marker(tmp_path, kind, CrossEncoder) + alias = "automodel_args" if kind == "model" else "config_args" + with pytest.raises( + _ControlledAutoImportMarker if allowed is True else ValueError, + match="controlled Auto import marker" if allowed is True else "requires trust_remote_code=True", + ): + trust.load_sentence_transformer( + CrossEncoder, + str(tmp_path), + trust_remote_code=allowed, + device="cpu", + local_files_only=True, + **{alias: UserDict({"trust_remote_code": True})}, + ) + assert marker.exists() is (allowed is True) + + +@pytest.mark.parametrize( + ("kind", "key"), + [ + ("model", "model_args"), + ("model", "model_kwargs"), + ("tokenizer", "tokenizer_args"), + ("tokenizer", "processor_kwargs"), + ("config", "config_args"), + ("config", "config_kwargs"), + ], +) +@pytest.mark.parametrize("allowed", [False, None, True]) +def test_checkpoint_saved_kwargs_cannot_elevate_trust(tmp_path, kind, key, allowed): + _bert_checkpoint(tmp_path, SentenceTransformer, BertModel) + _model(tmp_path, module_type="sentence_transformers.models.Transformer") + marker = _auto_marker(tmp_path, kind, SentenceTransformer) + _json(tmp_path / "sentence_bert_config.json", {key: {"trust_remote_code": True}}) + if allowed is True: + with pytest.raises(_ControlledAutoImportMarker, match="controlled Auto import marker"): + trust.load_sentence_transformer( + SentenceTransformer, str(tmp_path), trust_remote_code=True, device="cpu", local_files_only=True + ) + else: + # A native fallback is safe: the relevant invariant is no custom import. + loaded = trust.load_sentence_transformer( + SentenceTransformer, str(tmp_path), trust_remote_code=allowed, device="cpu", local_files_only=True + ) + assert type(loaded) is SentenceTransformer + assert marker.exists() is (allowed is True) + + +def test_private_loader_ignores_nested_import_trust_true(tmp_path, monkeypatch): + reference, marker = _marker(tmp_path, monkeypatch) + _nested(tmp_path, reference, 2, tokenizer=True) + boundary = trust._NativeLoadBoundary() + router = boundary.resolve_module("sentence_transformers.models.Router") + assert issubclass(router, Router) + with pytest.raises(ValueError, match="requires trust_remote_code=True"): + router.load(str(tmp_path), trust_remote_code=True, local_files_only=True) + assert not marker.exists() + + +def test_known_aliases_resolve_to_native_classes_without_imports(monkeypatch): + boundary = trust._NativeLoadBoundary() + imported = [] + monkeypatch.setattr(importlib, "import_module", lambda *args, **kwargs: imported.append(args)) + for name in ( + "sentence_transformers.models.Pooling", + "sentence_transformers.models.Pooling.Pooling", + "sentence_transformers.sentence_transformer.modules.Pooling", + "sentence_transformers.sentence_transformer.modules.pooling.Pooling", + ): + assert boundary.resolve_module(name) is Pooling + for name in ( + "sentence_transformers.models.tokenizer.WhitespaceTokenizer", + "sentence_transformers.models.tokenizer.WhitespaceTokenizer.WhitespaceTokenizer", + "sentence_transformers.sentence_transformer.modules.tokenizer.whitespace.WhitespaceTokenizer", + ): + assert boundary.resolve_tokenizer(name) is WhitespaceTokenizer + assert imported == [] + + +@pytest.mark.parametrize("allowed", [False, None]) +def test_existing_dense_activation_guard_preserved(tmp_path, monkeypatch, allowed): + reference, marker = _marker(tmp_path, monkeypatch) + Dense(4, 4).save(str(tmp_path)) + config = json.loads((tmp_path / "config.json").read_text()) + config["activation_function"] = reference + _json(tmp_path / "config.json", config) + native = trust._NativeLoadBoundary().resolve_module("sentence_transformers.models.Dense") + assert issubclass(native, Dense) + loaded = native.load(str(tmp_path), trust_remote_code=allowed, local_files_only=True) + assert type(loaded) is Dense + assert type(loaded.activation_function) is torch.nn.Tanh + assert not marker.exists() + + +def test_changed_nested_loader_fails_closed(tmp_path, monkeypatch): + _model(tmp_path, module_type="sentence_transformers.models.Router") + monkeypatch.setattr(Router, "load", classmethod(lambda cls, *args, **kwargs: pytest.fail("changed loader reached"))) + with pytest.raises(ValueError, match="Unsupported Sentence Transformers"): + trust.load_sentence_transformer(SentenceTransformer, str(tmp_path), local_files_only=True) + + +@pytest.mark.parametrize("router", [False, True]) +@pytest.mark.parametrize("allowed", [False, None]) +def test_native_loaded_objects_and_forwards_are_unchanged(tmp_path, router, allowed): + tokenizer = WhitespaceTokenizer(vocab=["[PAD]", "hello", "world"], stop_words=[]) + word = WordEmbeddings(tokenizer, torch.arange(12, dtype=torch.float32).reshape(3, 4)) + modules: list[Module] = [word, Pooling(4)] + if router: + modules = [Router({"query": modules}, default_route="query")] + native = SentenceTransformer(modules=modules, device="cpu", local_files_only=True) + native.save(str(tmp_path), create_model_card=False) + expected = native.encode(["hello world", "hello"], show_progress_bar=False) + loaded = trust.load_sentence_transformer( + SentenceTransformer, str(tmp_path), trust_remote_code=allowed, device="cpu", local_files_only=True + ) + assert type(loaded) is SentenceTransformer + assert trust._RESOLVER not in vars(loaded) + assert loaded.forward.__func__ is SentenceTransformer.forward + assert loaded.encode.__func__ is SentenceTransformer.encode + assert [type(item) for item in loaded.modules()] == [type(item) for item in native.modules()] + assert not any("facade" in type(item).__name__.lower() for item in loaded.modules()) + np.testing.assert_array_equal(loaded.encode(["hello world", "hello"], show_progress_bar=False), expected) + + +def test_temporary_binding_is_removed_on_constructor_failure(tmp_path, monkeypatch): + _model(tmp_path) + refs = [] + original = trust._NativeLoadBoundary.resolve_root + + def observe(boundary, model, *args, **kwargs): + refs.append(weakref.ref(model)) + assert trust._RESOLVER in vars(model) + return original(boundary, model, *args, **kwargs) + + monkeypatch.setattr(trust._NativeLoadBoundary, "resolve_root", observe) + with pytest.raises(ValueError, match="requires trust_remote_code=True") as error: + trust.load_sentence_transformer(SentenceTransformer, str(tmp_path), local_files_only=True) + tb = error.value.__traceback__ + while tb: + if tb.tb_frame.f_code.co_name == "load_sentence_transformer": + assert trust._RESOLVER not in vars(tb.tb_frame.f_locals["model"]) + tb = tb.tb_next + del error + gc.collect() + assert len(refs) == 1 + assert refs[0]() is None + + +@pytest.mark.parametrize( + ("model_class", "hf_class"), + [(SentenceTransformer, BertModel), (SparseEncoder, BertForMaskedLM), (CrossEncoder, BertForSequenceClassification)], +) +def test_native_transformer_constructors_and_predictions(tmp_path, model_class, hf_class): + vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "hello", "world", "one", "two"] + vocab_path = tmp_path / "vocab.txt" + vocab_path.write_text("\n".join(vocab) + "\n") + BertTokenizerFast(vocab_file=str(vocab_path)).save_pretrained(tmp_path) + config = BertConfig( + vocab_size=len(vocab), + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_hidden_layers=1, + num_labels=1, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + ) + hf_class(config).save_pretrained(tmp_path) + native = model_class(str(tmp_path), device="cpu", trust_remote_code=False, local_files_only=True) + loaded = trust.load_sentence_transformer( + model_class, str(tmp_path), device="cpu", trust_remote_code=False, local_files_only=True + ) + assert type(loaded) is model_class + assert trust._RESOLVER not in vars(loaded) + assert loaded.forward.__func__ is native.forward.__func__ + assert [type(item) for item in loaded.modules()] == [type(item) for item in native.modules()] + if model_class is CrossEncoder: + actual = loaded.predict([("hello", "world"), ("one", "two")], show_progress_bar=False) + expected = native.predict([("hello", "world"), ("one", "two")], show_progress_bar=False) + assert loaded.predict.__func__ is native.predict.__func__ + else: + actual = loaded.encode(["hello world", "one two"], show_progress_bar=False) + expected = native.encode(["hello world", "one two"], show_progress_bar=False) + assert loaded.encode.__func__ is native.encode.__func__ + if isinstance(actual, torch.Tensor): + assert isinstance(expected, torch.Tensor) + torch.testing.assert_close(actual.to_dense(), expected.to_dense(), rtol=0, atol=0) + else: + np.testing.assert_array_equal(actual, expected) + + +def test_unknown_version_fails_before_constructor(tmp_path, monkeypatch): + monkeypatch.setattr(sentence_transformers, "__version__", "5.8.0") + monkeypatch.setattr(SentenceTransformer, "__init__", lambda *args, **kwargs: pytest.fail("constructor reached")) + with pytest.raises(ValueError, match="Unsupported Sentence Transformers"): + trust.load_sentence_transformer(SentenceTransformer, str(tmp_path), local_files_only=True) + + +class _LoadControl(types.ModuleType): + entered: threading.Event + release: threading.Event + + +def test_concurrent_true_false_loads_do_not_change_library_bindings(tmp_path, monkeypatch): + control_name = "trust_control_" + uuid.uuid4().hex + control = _LoadControl(control_name) + control.entered = threading.Event() + control.release = threading.Event() + monkeypatch.setitem(sys.modules, control_name, control) + pause = f"import {control_name} as control\ncontrol.entered.set()\nassert control.release.wait(10)\n" + reference, marker = _marker(tmp_path, monkeypatch, pause=pause) + _nested(tmp_path, reference, 2, tokenizer=True) + loaders = (Router.load.__func__, WordEmbeddings.load.__func__) + before = [] + for fn in loaders: + assert isinstance(fn, types.FunctionType) + before.append((fn, fn.__globals__, dict(fn.__globals__))) + + def allowed_load(): + with pytest.raises(RuntimeError, match="controlled import marker"): + trust.load_sentence_transformer( + SentenceTransformer, str(tmp_path), trust_remote_code=True, local_files_only=True + ) + + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(allowed_load) + try: + assert control.entered.wait(10) + with pytest.raises(ValueError, match="requires trust_remote_code=True"): + trust.load_sentence_transformer( + SentenceTransformer, str(tmp_path), trust_remote_code=False, local_files_only=True + ) + assert not marker.exists() + finally: + control.release.set() + future.result(timeout=10) + assert marker.exists() + assert loaders == (Router.load.__func__, WordEmbeddings.load.__func__) + for fn, globals_object, entries in before: + assert fn.__globals__ is globals_object + assert entries.keys() == fn.__globals__.keys() + assert all(fn.__globals__[name] is value for name, value in entries.items()) + + +@pytest.mark.parametrize( + ("adapter_class", "model_class"), + [ + (SentenceTransformerDenseAdapter, SentenceTransformer), + (SentenceTransformerSparseAdapter, SparseEncoder), + (CrossEncoderAdapter, CrossEncoder), + (CrossEncoderLoraAdapter, CrossEncoder), + (SparseEncoderVisionAdapter, SparseEncoder), + ], +) +@pytest.mark.parametrize("allowed", [False, None, True]) +def test_actual_adapter_loads_enforce_root_trust(tmp_path, monkeypatch, adapter_class, model_class, allowed): + reference, marker = _marker(tmp_path, monkeypatch) + _model(tmp_path, model_class, reference) + extra = ( + {"lora_paths": {"example": str(tmp_path / "unused-adapter")}} + if adapter_class is CrossEncoderLoraAdapter + else {} + ) + adapter = adapter_class(str(tmp_path), trust_remote_code=allowed, **extra) + registrations = [] + original = SparseEncoderVisionAdapter._register_remote_code_dir + + def register(instance): + registrations.append(instance) + original(instance) + + monkeypatch.setattr(SparseEncoderVisionAdapter, "_register_remote_code_dir", register) + before = list(sys.path) + try: + with pytest.raises( + RuntimeError if allowed is True else ValueError, + match="controlled import marker" if allowed is True else "requires trust_remote_code=True", + ): + adapter.load("cpu") + assert marker.exists() is (allowed is True) + assert adapter._model is None + assert bool(registrations) is (allowed is True and adapter_class is SparseEncoderVisionAdapter) + if allowed is not True: + assert sys.path == before + finally: + adapter.unload() + + +@pytest.mark.parametrize("allowed", [False, None, True]) +def test_actual_dense_adapter_nested_config_cannot_elevate_trust(tmp_path, allowed): + _bert_checkpoint(tmp_path, SentenceTransformer, BertModel) + marker = _auto_marker(tmp_path, "config", SentenceTransformer) + adapter = SentenceTransformerDenseAdapter( + str(tmp_path), trust_remote_code=allowed, config_kwargs={"trust_remote_code": True} + ) + try: + with pytest.raises( + _ControlledAutoImportMarker if allowed is True else ValueError, + match="controlled Auto import marker" if allowed is True else "requires trust_remote_code=True", + ): + adapter.load("cpu") + assert marker.exists() is (allowed is True) + assert adapter._model is None + finally: + adapter.unload() + + +@pytest.mark.parametrize( + ("adapter_class", "model_class", "hf_class"), + [ + (SentenceTransformerDenseAdapter, SentenceTransformer, BertModel), + (SentenceTransformerSparseAdapter, SparseEncoder, BertForMaskedLM), + (CrossEncoderAdapter, CrossEncoder, BertForSequenceClassification), + ], +) +@pytest.mark.parametrize("allowed", [False, None, True]) +def test_actual_adapter_native_models_preserve_types_and_predictions( + tmp_path, adapter_class, model_class, hf_class, allowed +): + _bert_checkpoint(tmp_path, model_class, hf_class) + native = model_class(str(tmp_path), device="cpu", trust_remote_code=False, local_files_only=True) + adapter = adapter_class(str(tmp_path), trust_remote_code=allowed) + try: + adapter.load("cpu") + loaded = adapter._model + assert type(loaded) is model_class + assert trust._RESOLVER not in vars(loaded) + assert loaded.forward.__func__ is native.forward.__func__ + assert [type(item) for item in loaded.modules()] == [type(item) for item in native.modules()] + if model_class is CrossEncoder: + actual = loaded.predict([("hello", "world")], show_progress_bar=False) + expected = native.predict([("hello", "world")], show_progress_bar=False) + else: + actual = loaded.encode(["hello world"], show_progress_bar=False) + expected = native.encode(["hello world"], show_progress_bar=False) + if isinstance(actual, torch.Tensor): + torch.testing.assert_close(actual.to_dense(), expected.to_dense(), rtol=0, atol=0) + else: + np.testing.assert_array_equal(actual, expected) + finally: + adapter.unload() + assert adapter._model is None + + +def test_native_causal_adapter_uses_guard_before_custom_config_import(tmp_path, monkeypatch): + _bert_checkpoint(tmp_path, CrossEncoder, BertForSequenceClassification) + marker = _auto_marker(tmp_path, "config", CrossEncoder) + config = json.loads((tmp_path / "config.json").read_text()) + config["model_type"] = "controlled_custom_model_without_native_fallback" + _json(tmp_path / "config.json", config) + _json( + tmp_path / "modules.json", + [ + {"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.base.modules.transformer.Transformer"}, + { + "idx": 1, + "name": "1", + "path": "1_LogitScore", + "type": "sentence_transformers.cross_encoder.modules.logit_score.LogitScore", + }, + ], + ) + seen = [] + original = trust.load_sentence_transformer + + def observe(model_class, *args, **kwargs): + seen.append((model_class, kwargs["trust_remote_code"])) + return original(model_class, *args, **kwargs) + + monkeypatch.setattr(trust, "load_sentence_transformer", observe) + adapter = NativeCausalCrossEncoderAdapter(str(tmp_path)) + try: + with pytest.raises(ValueError, match="trust_remote_code"): + adapter.load("cpu") + assert seen == [(CrossEncoder, False)] + assert not marker.exists() + assert adapter._model is None + finally: + adapter.unload() diff --git a/packages/sie_server/tests/adapters/test_st_sparse_vision.py b/packages/sie_server/tests/adapters/test_st_sparse_vision.py index da0aab2e8..00ee40cd5 100644 --- a/packages/sie_server/tests/adapters/test_st_sparse_vision.py +++ b/packages/sie_server/tests/adapters/test_st_sparse_vision.py @@ -218,7 +218,7 @@ def test_load_evicts_stale_sibling_custom_module( assert "modeling_st_vsplade" not in sys.modules assert sys.path[0] == str(snapshot.resolve()) - @patch("sie_server.adapters.st_sparse_vision.adapter.SparseEncoder") + @patch("sie_server.adapters.st_sparse_vision.adapter.load_sentence_transformer") def test_load_without_trust_remote_code_skips_registration( self, mock_class: MagicMock, diff --git a/packages/sie_server/tests/config/test_native_zerank_catalog.py b/packages/sie_server/tests/config/test_native_zerank_catalog.py index dd1527952..600c6cfef 100644 --- a/packages/sie_server/tests/config/test_native_zerank_catalog.py +++ b/packages/sie_server/tests/config/test_native_zerank_catalog.py @@ -28,7 +28,7 @@ def test_native_catalog_pin_profile_and_bundle() -> None: assert profile["adapter_options"] == {"loadtime": {}, "runtime": {}} bundle = yaml.safe_load((_SERVER / "bundles/default.yaml").read_text()) assert _ADAPTER in bundle["adapters"] - assert bundle["deps"]["sentence-transformers"] == ">=5.4.1,<6" + assert bundle["deps"]["sentence-transformers"] == "==5.7.0" assert bundle["deps"]["transformers"] == ">=4.57,<5" assert bundle["deps"]["gliner2"] == ">=1.3.2,<=2.0.0" diff --git a/uv.lock b/uv.lock index 70a325af0..6ca7db56a 100644 --- a/uv.lock +++ b/uv.lock @@ -4350,22 +4350,23 @@ wheels = [ [[package]] name = "sentence-transformers" -version = "5.4.1" +version = "5.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, { name = "numpy" }, { name = "scikit-learn" }, { name = "scipy" }, + { name = "tokenizers" }, { name = "torch", version = "2.9.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, { name = "torch", version = "2.9.1+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform == 'linux'" }, { name = "tqdm" }, { name = "transformers" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4d/68/7f98c221940ce783b492ad6140384daf2e2918cd7175009d6a362c22b9ee/sentence_transformers-5.4.1.tar.gz", hash = "sha256:436bcb1182a0ff42a8fb2b1c43498a70d0a75b688d182f2cd0d1dd115af61ddc", size = 428910 } +sdist = { url = "https://files.pythonhosted.org/packages/9d/59/867381b1414a975da6c9953f48a07c05cb0629305e2d37c9bcc9764367b2/sentence_transformers-5.7.0.tar.gz", hash = "sha256:fd8c8fc35e6323631dff9f3760969ebf7980dc3cfda0ab1354bc6a774cc0e5d8", size = 466382 } wheels = [ - { url = "https://files.pythonhosted.org/packages/c5/d9/3a9b6f2ccdedc9dc00fe37b2fc58f58f8efbff44565cf4bf39d8568bb13a/sentence_transformers-5.4.1-py3-none-any.whl", hash = "sha256:a6d640fc363849b63affb8e140e9d328feabab86f83d58ac3e16b1c28140b790", size = 571311 }, + { url = "https://files.pythonhosted.org/packages/e8/c8/f63d99e354532f5b83e735dd1e001bda92495fbfde934f65d924abf2b071/sentence_transformers-5.7.0-py3-none-any.whl", hash = "sha256:b78141da3d8137e70d965866e2ca43190b9266f3d4d8752e250ded75e7136730", size = 611333 }, ] [[package]] @@ -4957,7 +4958,7 @@ requires-dist = [ { name = "requests", specifier = ">=2.32,<3" }, { name = "rich", specifier = ">=13,<14" }, { name = "safetensors", specifier = ">=0.4,<1" }, - { name = "sentence-transformers", specifier = ">=5.0,<6" }, + { name = "sentence-transformers", specifier = "==5.7.0" }, { name = "sie-sdk", extras = ["storage"], editable = "packages/sie_sdk" }, { name = "textual", marker = "extra == 'top'", specifier = ">=3.0,<4" }, { name = "threadpoolctl", specifier = ">=3.6.0" },