feat: Enable trust remote code in CI.

hf_adapters/auto_spyre_model.py

@@ -283,8 +283,14 @@ def _autoconfig_with_subfolder_fallback(
 def dtype_for_model_path(
     model_name_or_path: Union[str, os.PathLike[str]],
     target_device: str | torch.device,
+    trust_remote_code: bool | None = None,
 ) -> torch.dtype:
     """Resolve one concrete dtype before loading a model."""
+    # trust_remote_code is forwarded to AutoConfig.from_pretrained only when the
+    # dtype has to be read from the config (no explicit policy). Checkpoints that
+    # ship custom config code otherwise block on an interactive opt-in prompt
+    # here, before the load ever reaches from_pretrained.
+
     device_str = (
         target_device.type
         if isinstance(target_device, torch.device)
@@ -296,7 +302,9 @@ def dtype_for_model_path(
     elif policy.dtype is not None:
         dtype = policy.dtype
     else:
-        config = _autoconfig_with_subfolder_fallback(model_name_or_path)
+        config = _autoconfig_with_subfolder_fallback(
+            model_name_or_path, trust_remote_code=trust_remote_code
+        )
         dtype = (
             getattr(config, "dtype", None) or torch.float16 if config else torch.float16
         )
@@ -375,6 +383,7 @@ def from_pretrained(
             dtype = dtype_for_model_path(
                 model_name_or_path,
                 target_device=hf_common.DEVICE,
+                trust_remote_code=trust_remote_code,
             )
 
         model: PreTrainedModel = load_model_common(

hf_adapters/st_backend.py

@@ -86,8 +86,18 @@ def _spyre_load_model(
     if dtype is None:
         dtype = model_kwargs.pop("torch_dtype", None)
 
-    model = AutoSpyreModel.from_pretrained(model_name_or_path, dtype=dtype)
-    adapter_module = resolve_adapter_module(model_name_or_path)
+    # Remote-code checkpoints ship custom modeling code that both the backbone
+    # load and the adapter-module resolution must opt into; ST callers pass it
+    # via ``model_kwargs`` (e.g. ``SentenceTransformer(path, backend="spyre",
+    # model_kwargs={"trust_remote_code": True})``).
+    trust_remote_code = model_kwargs.pop("trust_remote_code", None)
+
+    model = AutoSpyreModel.from_pretrained(
+        model_name_or_path, dtype=dtype, trust_remote_code=trust_remote_code
+    )
+    adapter_module = resolve_adapter_module(
+        model_name_or_path, trust_remote_code=trust_remote_code
+    )
 
     run_backbone_forward = adapter_module._run_backbone_forward
 

tests/_vision_helpers.py

@@ -32,6 +32,7 @@
 from transformers import AutoProcessor
 
 from tests.conftest import load_ref_model
+from tests.model_registry import REMOTE_CODE_PATHS
 
 # ── VLM (image→text) end-to-end helpers ──────────────────────────────────────
 #
@@ -146,6 +147,7 @@ def build_vlm_batch(
     model_path: str,
     prompt: str,
     image: Image.Image | None = None,
+    trust_remote_code: bool | None = None,
 ) -> tuple[AutoProcessor, dict[str, torch.Tensor]]:
     """Processor + tokenized (image + prompt) batch, the official VLM way.
 
@@ -158,10 +160,16 @@ def build_vlm_batch(
     convention. Returns ``(processor, batch)``; ``batch`` carries whatever image
     inputs the model needs (``pixel_values``, ``image_sizes``, …).
     """
+    if trust_remote_code is None:
+        trust_remote_code = model_path in REMOTE_CODE_PATHS
     if "mistral" in model_path.lower():
-        processor = AutoProcessor.from_pretrained(model_path, fix_mistral_regex=True)
+        processor = AutoProcessor.from_pretrained(
+            model_path, fix_mistral_regex=True, trust_remote_code=trust_remote_code
+        )
     else:
-        processor = AutoProcessor.from_pretrained(model_path)
+        processor = AutoProcessor.from_pretrained(
+            model_path, trust_remote_code=trust_remote_code
+        )
     processor.tokenizer.padding_side = "left"
 
     if image is None:
@@ -192,6 +200,7 @@ def stock_vlm_generate(
     adapter_mod,
     max_new_tokens: int,
     ref_model=None,
+    trust_remote_code: bool | None = None,
 ) -> str:
     """Reference: stock ``AutoModelForImageTextToText.generate`` on ``batch``.
 
@@ -206,6 +215,7 @@ def stock_vlm_generate(
         ref_model = load_ref_model(
             model_path=model_path,
             adapter_mod=adapter_mod,
+            trust_remote_code=trust_remote_code,
             auto_model_cls=AutoModelForImageTextToText,
         )
 

tests/conftest.py

@@ -182,6 +182,28 @@ def pytest_addoption(parser: Parser) -> None:
             "in the test decorators are ignored."
         ),
     )
+    parser.addoption(
+        "--trust-remote-code",
+        action="store_true",
+        default=False,
+        help=(
+            "Force trust_remote_code=True for every test that loads a model. "
+            "Use with --model-path to run a remote-code checkpoint that is not "
+            "listed in tests/model_registry.py's REMOTE_CODE_PATHS."
+        ),
+    )
+
+
+@pytest.fixture
+def trust_remote_code(request) -> bool | None:
+    """CLI override for trust_remote_code.
+
+    ``None`` when the flag is absent — every test then falls back to its
+    ``model_path in REMOTE_CODE_PATHS`` check, so registry-driven runs are
+    unchanged. ``True`` when ``--trust-remote-code`` is passed, which wins over
+    the registry (the escape hatch for off-registry ``--model-path`` runs).
+    """
+    return True if request.config.getoption("--trust-remote-code") else None
 
 
 def pytest_generate_tests(metafunc: Metafunc) -> None:
@@ -308,17 +330,25 @@ def load_ref_model(
     model_path: str,
     adapter_mod: types.ModuleType | None = None,
     auto_model_cls: type = AutoModelForCausalLM,
+    trust_remote_code: bool | None = None,
 ):
+    from model_registry import REMOTE_CODE_PATHS
+
     from hf_adapters.auto_spyre_model import dtype_for_model_path
     from hf_adapters.hf_common import load_model_common
 
-    dtype = dtype_for_model_path(model_path, target_device="cpu")
+    if trust_remote_code is None:
+        trust_remote_code = model_path in REMOTE_CODE_PATHS
+    dtype = dtype_for_model_path(
+        model_path, target_device="cpu", trust_remote_code=trust_remote_code
+    )
 
     ref_model = load_model_common(
         model_path=model_path,
         module=adapter_mod,
         dtype=dtype,
         auto_model_cls=auto_model_cls,
+        trust_remote_code=trust_remote_code,
     )
     return ref_model
 
@@ -328,7 +358,14 @@ def resolve_adapter_module_for_test(
     mapping: dict[
         type[PretrainedConfig], types.ModuleType
     ] = CONFIG_TO_ADAPTER_MODULE_MAPPING,
+    trust_remote_code: bool | None = None,
 ) -> types.ModuleType:
+    from model_registry import REMOTE_CODE_PATHS
+
+    if trust_remote_code is None:
+        trust_remote_code = str(model_name_or_path) in REMOTE_CODE_PATHS
     return resolve_adapter_module(
-        model_name_or_path=model_name_or_path, mapping=mapping, trust_remote_code=False
+        model_name_or_path=model_name_or_path,
+        mapping=mapping,
+        trust_remote_code=trust_remote_code,
     )

tests/cpu/_seq_classification_helpers.py

@@ -30,11 +30,13 @@
 
 from tests.conftest import load_ref_model
 from tests.cpu.conftest import _unwrap_compiled_blocks
+from tests.model_registry import REMOTE_CODE_PATHS
 
 
 def run_seq_classification_auto_loader_vs_ref(
     model_path: str,
     inputs: list[str] | list[tuple[str, str]],
+    trust_remote_code: bool | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Run the auto-loader seq-classification path against a stock HF reference.
 
@@ -51,7 +53,11 @@ def run_seq_classification_auto_loader_vs_ref(
         ``(ref_logits, adapter_logits)`` — both ``[B, num_labels]`` float CPU tensors.
     """
     auto_spyre_model_mod = sys.modules["hf_adapters.auto_spyre_model"]
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    if trust_remote_code is None:
+        trust_remote_code = model_path in REMOTE_CODE_PATHS
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code
+    )
 
     encoded = tokenizer(
         inputs,
@@ -66,6 +72,7 @@ def run_seq_classification_auto_loader_vs_ref(
     ref_model = load_ref_model(
         model_path=model_path,
         auto_model_cls=AutoModelForSequenceClassification,
+        trust_remote_code=trust_remote_code,
     )
     ref_model.eval()
     with torch.no_grad():
@@ -75,7 +82,7 @@ def run_seq_classification_auto_loader_vs_ref(
     # --- Auto-loader path ---
     model = (
         auto_spyre_model_mod.AutoSpyreModelForSequenceClassification.from_pretrained(
-            model_path
+            model_path, trust_remote_code=trust_remote_code
         )
     )
     _unwrap_compiled_blocks(model)

tests/cpu/test_adapter_cpu_accuracy.py

@@ -45,7 +45,7 @@
     resolve_adapter_module_for_test,
 )
 from tests.cpu.conftest import _unwrap_compiled_blocks
-from tests.model_registry import CAUSAL_PATHS
+from tests.model_registry import CAUSAL_PATHS, REMOTE_CODE_PATHS
 
 pytestmark = pytest.mark.model_harness("causal")
 
@@ -162,12 +162,18 @@ def adapter_greedy_steps(run_forward_fn, model, input_ids, num_decode=NUM_DECODE
 
 
 @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS)
-def test_auto_loader(model_path):
+def test_auto_loader(model_path, trust_remote_code):
     auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"]
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    if trust_remote_code is None:
+        trust_remote_code = model_path in REMOTE_CODE_PATHS
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code
+    )
 
     # Phase 1: auto-loader generate
-    model = auto_spyre_model.AutoSpyreModelForCausalLM.from_pretrained(model_path)
+    model = auto_spyre_model.AutoSpyreModelForCausalLM.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code
+    )
     _unwrap_compiled_blocks(model)
     encoded = encode_generation_inputs(tokenizer, [PROMPT])
     auto_sequences = model.generate(
@@ -183,8 +189,12 @@ def test_auto_loader(model_path):
     gc.collect()
 
     # Phase 2: HF reference (fresh)
-    adapter_mod = resolve_adapter_module_for_test(model_path)
-    hf_model = load_ref_model(model_path, adapter_mod)
+    adapter_mod = resolve_adapter_module_for_test(
+        model_path, trust_remote_code=trust_remote_code
+    )
+    hf_model = load_ref_model(
+        model_path, adapter_mod, trust_remote_code=trust_remote_code
+    )
     encoded = encode_prompts(tokenizer, PROMPT)
     with torch.no_grad():
         hf_out = hf_model.generate(

tests/cpu/test_embed_cpu_accuracy.py

@@ -52,7 +52,7 @@
     resolve_adapter_module_for_test,
 )
 from tests.cpu.conftest import _unwrap_compiled_blocks, encode_padded, min_cosine
-from tests.model_registry import EMBED_PATHS
+from tests.model_registry import EMBED_PATHS, REMOTE_CODE_PATHS
 
 pytestmark = pytest.mark.model_harness("embedding")
 
@@ -82,17 +82,26 @@ def _run_prefill(
 
 
 @pytest.mark.parametrize("model_path", EMBED_PATHS, ids=EMBED_PATHS)
-def test_auto_loader(model_path: str) -> None:
+def test_auto_loader(model_path: str, trust_remote_code: bool | None) -> None:
     auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"]
     hf_common_mod = sys.modules["hf_adapters.hf_common"]
-    adapter_module = resolve_adapter_module_for_test(model_path)
+    if trust_remote_code is None:
+        trust_remote_code = model_path in REMOTE_CODE_PATHS
+    adapter_module = resolve_adapter_module_for_test(
+        model_path, trust_remote_code=trust_remote_code
+    )
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code
+    )
     input_ids, attention_mask = encode_padded(tokenizer, PROMPTS)
 
     # HF reference (loaded fresh, before the auto-loader path).
     ref_model = load_ref_model(
-        model_path=model_path, adapter_mod=adapter_module, auto_model_cls=AutoModel
+        model_path=model_path,
+        adapter_mod=adapter_module,
+        auto_model_cls=AutoModel,
+        trust_remote_code=trust_remote_code,
     )
 
     with torch.no_grad():
@@ -103,7 +112,9 @@ def test_auto_loader(model_path: str) -> None:
     gc.collect()
 
     # Auto-loader path
-    model = auto_spyre_model.AutoSpyreModel.from_pretrained(model_path)
+    model = auto_spyre_model.AutoSpyreModel.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code
+    )
     _unwrap_compiled_blocks(model)
     with torch.no_grad():
         adapter_hidden = _run_prefill(

tests/cpu/test_generate_cpu.py

@@ -44,6 +44,7 @@
 from tests.model_registry import (
     CAUSAL_PATHS,
     NON_BLOCKING_CAUSAL_MODELS,
+    REMOTE_CODE_PATHS,
     xfail_non_blocking,
 )
 
@@ -53,26 +54,36 @@
 @pytest.mark.parametrize(
     "model_path", xfail_non_blocking(CAUSAL_PATHS, table=NON_BLOCKING_CAUSAL_MODELS)
 )
-def test_multibatch(model_path: str) -> None:
+def test_multibatch(model_path: str, trust_remote_code: bool | None) -> None:
     from hf_adapters.auto_spyre_model import dtype_for_model_path
 
     hf_common_mod = sys.modules["hf_adapters.hf_common"]
-    adapter_mod = resolve_adapter_module_for_test(model_path)
+    if trust_remote_code is None:
+        trust_remote_code = model_path in REMOTE_CODE_PATHS
+    adapter_mod = resolve_adapter_module_for_test(
+        model_path, trust_remote_code=trust_remote_code
+    )
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code
+    )
 
     # HF reference (per-prompt, BEFORE patching for cleanliness)
-    model = load_ref_model(model_path, adapter_mod)
+    model = load_ref_model(model_path, adapter_mod, trust_remote_code=trust_remote_code)
     hf_outputs = hf_reference_outputs(model, tokenizer, PROMPTS, MAX_NEW_TOKENS)
     del model
     gc.collect()
 
     # Adapter batched generate
     encoded = encode_generation_inputs(tokenizer, PROMPTS)
-    model = load_ref_model(model_path, adapter_mod)
+    model = load_ref_model(model_path, adapter_mod, trust_remote_code=trust_remote_code)
     adapter_mod.prepare_for_spyre(model)
     _unwrap_compiled_blocks(model)
-    dtype = dtype_for_model_path(model_path, target_device="cpu")
+    dtype = dtype_for_model_path(
+        model_path,
+        target_device="cpu",
+        trust_remote_code=trust_remote_code,
+    )
     _set_rope_dtype(model, dtype)
     sequences = hf_common_mod.generate(
         adapter_mod._run_forward,