From c9259e4c89ab1a54c0f569981516cee873334d2d Mon Sep 17 00:00:00 2001 From: Dushyant Behl Date: Wed, 9 Sep 2026 13:04:47 +0530 Subject: [PATCH 1/3] enable trust remote code models in test framework Signed-off-by: Dushyant Behl --- hf_adapters/auto_spyre_model.py | 11 ++++- tests/conftest.py | 14 +++++- tests/cpu/test_adapter_cpu_accuracy.py | 11 +++-- tests/cpu/test_embed_cpu_accuracy.py | 6 ++- tests/cpu/test_generate_cpu.py | 11 ++++- tests/cpu/test_load_cpu.py | 23 +++++++-- tests/model_lists/exclude.yaml | 1 - tests/model_registry.py | 55 ++++++++++++++++++++- tests/spyre/edge_cases/_shared.py | 9 +++- tests/spyre/test_e2e_multibatch_spyre.py | 16 ++++-- tests/spyre/test_e2e_smoke_spyre.py | 9 +++- tests/spyre/test_e2e_token_compare_spyre.py | 11 ++++- tests/spyre/test_load_spyre.py | 25 +++++++--- 13 files changed, 170 insertions(+), 32 deletions(-) diff --git a/hf_adapters/auto_spyre_model.py b/hf_adapters/auto_spyre_model.py index 436cf5a3..884d5d84 100644 --- a/hf_adapters/auto_spyre_model.py +++ b/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( diff --git a/tests/conftest.py b/tests/conftest.py index 3186eafc..5fd7fc85 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -309,16 +309,22 @@ def load_ref_model( adapter_mod: types.ModuleType | None = None, auto_model_cls: type = AutoModelForCausalLM, ): + 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") + 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 @@ -329,6 +335,10 @@ def resolve_adapter_module_for_test( type[PretrainedConfig], types.ModuleType ] = CONFIG_TO_ADAPTER_MODULE_MAPPING, ) -> types.ModuleType: + from model_registry import 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=str(model_name_or_path) in REMOTE_CODE_PATHS, ) diff --git a/tests/cpu/test_adapter_cpu_accuracy.py b/tests/cpu/test_adapter_cpu_accuracy.py index f015e3db..6fb95fef 100644 --- a/tests/cpu/test_adapter_cpu_accuracy.py +++ b/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") @@ -164,10 +164,15 @@ 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): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - tokenizer = AutoTokenizer.from_pretrained(model_path) + 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( diff --git a/tests/cpu/test_embed_cpu_accuracy.py b/tests/cpu/test_embed_cpu_accuracy.py index f9ecc3bb..33326006 100644 --- a/tests/cpu/test_embed_cpu_accuracy.py +++ b/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") @@ -103,7 +103,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=(model_path in REMOTE_CODE_PATHS) + ) _unwrap_compiled_blocks(model) with torch.no_grad(): adapter_hidden = _run_prefill( diff --git a/tests/cpu/test_generate_cpu.py b/tests/cpu/test_generate_cpu.py index 1cbbfd6e..dcadb546 100644 --- a/tests/cpu/test_generate_cpu.py +++ b/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, ) @@ -59,7 +60,9 @@ def test_multibatch(model_path: str) -> None: hf_common_mod = sys.modules["hf_adapters.hf_common"] adapter_mod = resolve_adapter_module_for_test(model_path) - tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) # HF reference (per-prompt, BEFORE patching for cleanliness) model = load_ref_model(model_path, adapter_mod) @@ -72,7 +75,11 @@ def test_multibatch(model_path: str) -> None: model = load_ref_model(model_path, adapter_mod) 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=(model_path in REMOTE_CODE_PATHS), + ) _set_rope_dtype(model, dtype) sequences = hf_common_mod.generate( adapter_mod._run_forward, diff --git a/tests/cpu/test_load_cpu.py b/tests/cpu/test_load_cpu.py index 62bb5c0e..b49823c6 100644 --- a/tests/cpu/test_load_cpu.py +++ b/tests/cpu/test_load_cpu.py @@ -35,6 +35,7 @@ EMBED_PATHS, MASKED_LM_PATHS, QUESTION_ANSWERING_PATHS, + REMOTE_CODE_PATHS, TOKEN_CLASSIFICATION_PATHS, ) @@ -43,7 +44,9 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) def test_load_causal_lm(model_path): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - model = auto_spyre_model.AutoSpyreModelForCausalLM.from_pretrained(model_path) + model = auto_spyre_model.AutoSpyreModelForCausalLM.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) assert model is not None assert callable( getattr(model, "generate", None) @@ -63,7 +66,9 @@ def test_load_embedding(model_path): def load_embedding(model_path: str) -> Any: auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - model = auto_spyre_model.AutoSpyreModel.from_pretrained(model_path) + model = auto_spyre_model.AutoSpyreModel.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) return model @@ -71,9 +76,13 @@ def load_embedding(model_path: str) -> Any: @pytest.mark.parametrize("model_path", MASKED_LM_PATHS, ids=MASKED_LM_PATHS) def test_load_masked_lm(model_path): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] + trust_remote_code = model_path in REMOTE_CODE_PATHS model = auto_spyre_model.AutoSpyreModelForMaskedLM.from_pretrained( model_path, - dtype=auto_spyre_model.dtype_for_model_path(model_path, target_device="cpu"), + dtype=auto_spyre_model.dtype_for_model_path( + model_path, target_device="cpu", trust_remote_code=trust_remote_code + ), + trust_remote_code=trust_remote_code, ) assert callable(model.forward) del model @@ -86,9 +95,13 @@ def test_load_masked_lm(model_path): ) def test_load_question_answering(model_path): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] + trust_remote_code = model_path in REMOTE_CODE_PATHS model = auto_spyre_model.AutoSpyreModelForQuestionAnswering.from_pretrained( model_path, - dtype=auto_spyre_model.dtype_for_model_path(model_path, target_device="cpu"), + dtype=auto_spyre_model.dtype_for_model_path( + model_path, target_device="cpu", trust_remote_code=trust_remote_code + ), + trust_remote_code=trust_remote_code, ) assert callable(model.forward) assert next(model.qa_outputs.parameters()).device.type == "cpu" @@ -103,7 +116,7 @@ def test_load_question_answering(model_path): def test_load_token_classification(model_path): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] model = auto_spyre_model.AutoSpyreModelForTokenClassification.from_pretrained( - model_path + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) ) assert callable(model.forward) assert next(model.classifier.parameters()).device.type == "cpu" diff --git a/tests/model_lists/exclude.yaml b/tests/model_lists/exclude.yaml index a61dd24f..211bf849 100644 --- a/tests/model_lists/exclude.yaml +++ b/tests/model_lists/exclude.yaml @@ -4,4 +4,3 @@ models: - tiiuae/Falcon3-1B-Base - mistralai/Mistral-Small-3.1-24B-Instruct-2503 - ibm-research/granite-4.1-20b # not yet publicly accessible - - bharatgenai/Param-1-5B # can be enabled post https://github.com/torch-spyre/hf-adapters/issues/482 diff --git a/tests/model_registry.py b/tests/model_registry.py index e7dd0ec7..2e4ccf2b 100644 --- a/tests/model_registry.py +++ b/tests/model_registry.py @@ -42,6 +42,17 @@ def _include_gated() -> bool: return os.getenv("SPYRE_INCLUDE_GATED", "0") == "1" +def _include_trust_remote_code() -> bool: + """Whether ``trust_remote_code`` models should be included in test lists. + + Some checkpoints (e.g. bharatgenai/Param-1-5B) ship custom modelling code on + the Hub and only load with ``trust_remote_code=True``. We do not run them + implicitly in CI so they are excluded by default. + Set ``SPYRE_INCLUDE_TRUST_REMOTE_CODE=1`` to opt them in. + """ + return os.getenv("SPYRE_INCLUDE_TRUST_REMOTE_CODE", "0") == "1" + + # Model registries - shared by all tests CAUSAL_LM_MODELS = { # hf_gpt2.py @@ -651,6 +662,7 @@ def _select_representative_paths( models: dict[str, dict], *, include_gated: bool, + include_trust_remote_code: bool | None = None, predicate=None, ) -> list[str]: """Select representative model paths for each adapter module. @@ -658,14 +670,19 @@ def _select_representative_paths( Groups ``models`` by adapter and picks the smallest (by ``size``) model in each group, breaking ties by key name for determinism. Entries marked ``always_test`` are included in addition to that representative. Gated - models are skipped unless ``include_gated``. An optional + models are skipped unless ``include_gated``; ``trust_remote_code`` models are + skipped unless ``include_trust_remote_code``. An optional ``predicate(info) -> bool`` filters which entries are eligible (e.g. ``kind == "vlm"`` for vision). """ + if include_trust_remote_code is None: + include_trust_remote_code = _include_trust_remote_code() adapter_to_keys: dict[str, list[str]] = {} for key, info in models.items(): if info.get("is_gated", False) and not include_gated: continue + if info.get("trust_remote_code", False) and not include_trust_remote_code: + continue if predicate is not None and not predicate(info): continue adapter = info["adapter"].replace(".py", "") @@ -718,6 +735,7 @@ def _exclude(paths: list[str]) -> list[str]: # ``kind == "vlm"`` excludes bare vision towers. _include_gated_flag = _include_gated() + # ``kind == "dspark_draft"`` entries are speculative-decoding drafters (block # proposers, driven by ``_run_draft_block`` — no ``generate``), so they are # registered for adapter-coverage but excluded from the generate-based CPU/Spyre @@ -782,6 +800,7 @@ def _all_paths( models: dict[str, dict], *, include_gated: bool, + include_trust_remote_code: bool | None = None, predicate=None, ) -> list[str]: """All registered paths (no per-adapter reduction), for explicit selection. @@ -791,14 +810,48 @@ def _all_paths( target a non-representative checkpoint, e.g. a larger model that shares an adapter with a smaller default. """ + if include_trust_remote_code is None: + include_trust_remote_code = _include_trust_remote_code() return [ info["path"] for info in models.values() if (include_gated or not info.get("is_gated", False)) + and (include_trust_remote_code or not info.get("trust_remote_code", False)) and (predicate is None or predicate(info)) ] +def _collect_remote_code_paths() -> frozenset[str]: + """Paths of every registered checkpoint that needs ``trust_remote_code``. + + Spans all category registries so a single lookup covers any harness. Test + call sites check membership in ``REMOTE_CODE_PATHS`` to decide whether to + forward ``trust_remote_code=True`` to ``from_pretrained``. + """ + registries = ( + CAUSAL_LM_MODELS, + EMBEDDING_MODELS, + MASKED_LM_MODELS, + QUESTION_ANSWERING_MODELS, + TOKEN_CLASSIFICATION_MODELS, + VISION_MODELS, + ) + paths: set[str] = set() + for models in registries: + with_trc = _all_paths( + models, include_gated=True, include_trust_remote_code=True + ) + without_trc = _all_paths( + models, include_gated=True, include_trust_remote_code=False + ) + paths.update(set(with_trc) - set(without_trc)) + return frozenset(paths) + + +# Paths that must be loaded with ``trust_remote_code=True`` +REMOTE_CODE_PATHS: frozenset[str] = _collect_remote_code_paths() + + # Every registered path per category, bypassing the smallest-per-adapter # reduction above -- used by generate_test_matrix.py's ``--only`` allowlist so # a caller can target any registered checkpoint, not just the adapter's diff --git a/tests/spyre/edge_cases/_shared.py b/tests/spyre/edge_cases/_shared.py index 8c68c8b5..829e076e 100644 --- a/tests/spyre/edge_cases/_shared.py +++ b/tests/spyre/edge_cases/_shared.py @@ -48,12 +48,15 @@ load_ref_model, resolve_adapter_module_for_test, ) +from tests.model_registry import REMOTE_CODE_PATHS def _load_spyre_model(model_path: str) -> PreTrainedModel: print(f" Loading {model_path} on Spyre ...") t0 = time.time() - model = AutoSpyreModelForCausalLM.from_pretrained(model_path) + model = AutoSpyreModelForCausalLM.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) print(f" Spyre load+prepare: {time.time() - t0:.1f}s") return model @@ -62,7 +65,9 @@ def _setup( model_path: str, need_ref: bool, ): - tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) adapter = resolve_adapter_module_for_test(model_path) ref_model = load_ref_model(model_path, adapter_mod=adapter) if need_ref else None diff --git a/tests/spyre/test_e2e_multibatch_spyre.py b/tests/spyre/test_e2e_multibatch_spyre.py index 7041baa5..bf14acc5 100644 --- a/tests/spyre/test_e2e_multibatch_spyre.py +++ b/tests/spyre/test_e2e_multibatch_spyre.py @@ -69,6 +69,7 @@ from tests.model_registry import ( CAUSAL_PATHS, NON_BLOCKING_CAUSAL_MODELS, + REMOTE_CODE_PATHS, xfail_non_blocking, ) @@ -100,7 +101,9 @@ def test_e2e_multibatch_spyre(model_path: str) -> None: assert len(PROMPTS) > 1, "multi-batch test needs batch_size > 1" - tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) print(f"\n{'=' * 70}") print(f" {model_path} (batch_size={len(PROMPTS)})") @@ -116,8 +119,15 @@ def test_e2e_multibatch_spyre(model_path: str) -> None: # One Spyre model, reused for the batched run and each single run (generate # allocates a fresh KV cache per call, so runs do not contaminate each other). - spyre_dtype = dtype_for_model_path(model_path, target_device="spyre") - model = AutoSpyreModelForCausalLM.from_pretrained(model_path, dtype=spyre_dtype) + trust_remote_code = model_path in REMOTE_CODE_PATHS + spyre_dtype = dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ) + model = AutoSpyreModelForCausalLM.from_pretrained( + model_path, + dtype=spyre_dtype, + trust_remote_code=trust_remote_code, + ) # Exercise the public tokenized-input API attached by AutoSpyreModel. Bind the # model as a default arg because the local name is deleted after the runs. diff --git a/tests/spyre/test_e2e_smoke_spyre.py b/tests/spyre/test_e2e_smoke_spyre.py index 306e92d2..491dcf73 100644 --- a/tests/spyre/test_e2e_smoke_spyre.py +++ b/tests/spyre/test_e2e_smoke_spyre.py @@ -35,6 +35,7 @@ from tests.model_registry import ( CAUSAL_PATHS, NON_BLOCKING_CAUSAL_MODELS, + REMOTE_CODE_PATHS, xfail_non_blocking, ) @@ -52,11 +53,15 @@ def run_smoke_test(model_path: str) -> dict[str, Any]: print(f"{'=' * 70}") t0 = time.time() - model = AutoSpyreModelForCausalLM.from_pretrained(model_path) + model = AutoSpyreModelForCausalLM.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) load_time = time.time() - t0 print(f" Load time: {load_time:.1f}s") - tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) prompt = "The capital of France is" print(f" Prompt: {prompt!r}") diff --git a/tests/spyre/test_e2e_token_compare_spyre.py b/tests/spyre/test_e2e_token_compare_spyre.py index 3b03f7f0..16b6a818 100644 --- a/tests/spyre/test_e2e_token_compare_spyre.py +++ b/tests/spyre/test_e2e_token_compare_spyre.py @@ -42,6 +42,7 @@ from tests.model_registry import ( CAUSAL_PATHS, NON_BLOCKING_CAUSAL_MODELS, + REMOTE_CODE_PATHS, xfail_non_blocking, ) @@ -281,7 +282,9 @@ def _run_model_test(model_path: str, num_decode: int = 4) -> list[dict[str, Any] print(f" {model_path}") print(f"{'=' * 70}") - tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) model = load_ref_model(model_path=model_path, adapter_mod=adapter) prompt = "The capital of France is" @@ -297,7 +300,11 @@ def _run_model_test(model_path: str, num_decode: int = 4) -> list[dict[str, Any] # Use bf16/fp16 dtype, requested by the registry or based on the model config. # (Spyre does not support float32, so float32 entries will use fp16.) - spyre_dtype = dtype_for_model_path(model_path, target_device="spyre") + spyre_dtype = dtype_for_model_path( + model_path, + target_device="spyre", + trust_remote_code=(model_path in REMOTE_CODE_PATHS), + ) move_model_to_spyre(model=model, module=adapter, dtype=spyre_dtype) print(" Running adapter on Spyre ...") adapter_results = adapter_greedy_steps( diff --git a/tests/spyre/test_load_spyre.py b/tests/spyre/test_load_spyre.py index 9cb567f6..24d509bd 100644 --- a/tests/spyre/test_load_spyre.py +++ b/tests/spyre/test_load_spyre.py @@ -44,6 +44,7 @@ MASKED_LM_PATHS, NON_BLOCKING_CAUSAL_MODELS, QUESTION_ANSWERING_PATHS, + REMOTE_CODE_PATHS, SEQ_CLASSIFICATION_PATHS, TOKEN_CLASSIFICATION_PATHS, xfail_non_blocking, @@ -73,7 +74,9 @@ def load_causal_lm(model_path: str) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForCausalLM t0 = time.time() - model = AutoSpyreModelForCausalLM.from_pretrained(model_path) + model = AutoSpyreModelForCausalLM.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) load_s = time.time() - t0 model_is_not_none = model is not None @@ -86,7 +89,9 @@ def load_embedding(model_path: str) -> tuple[Any, float]: from hf_adapters import AutoSpyreModel t0 = time.time() - model = AutoSpyreModel.from_pretrained(model_path) + model = AutoSpyreModel.from_pretrained( + model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + ) load_s = time.time() - t0 return model is not None, load_s @@ -108,10 +113,15 @@ def load_masked_lm(model_path: str) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForMaskedLM from hf_adapters.auto_spyre_model import dtype_for_model_path - dtype = dtype_for_model_path(model_path, target_device="spyre") + trust_remote_code = model_path in REMOTE_CODE_PATHS + dtype = dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ) t0 = time.time() - model = AutoSpyreModelForMaskedLM.from_pretrained(model_path, dtype=dtype) + model = AutoSpyreModelForMaskedLM.from_pretrained( + model_path, dtype=dtype, trust_remote_code=trust_remote_code + ) load_s = time.time() - t0 model_is_not_none = model is not None @@ -138,10 +148,13 @@ def load_question_answering(model_path: str) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForQuestionAnswering from hf_adapters.auto_spyre_model import dtype_for_model_path - dtype = dtype_for_model_path(model_path, target_device="spyre") + trust_remote_code = model_path in REMOTE_CODE_PATHS + dtype = dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ) t0 = time.time() model: Any = AutoSpyreModelForQuestionAnswering.from_pretrained( - model_path, dtype=dtype + model_path, dtype=dtype, trust_remote_code=trust_remote_code ) load_s = time.time() - t0 head_on_cpu = next(model.qa_outputs.parameters()).device.type == "cpu" From bad4be779c165a41cbb15c208ab24d2c78595bfe Mon Sep 17 00:00:00 2001 From: Dushyant Behl Date: Thu, 10 Sep 2026 13:40:28 +0530 Subject: [PATCH 2/3] update code and add trust_remote_code to remaining files Signed-off-by: Dushyant Behl --- hf_adapters/st_backend.py | 14 +++++- tests/_vision_helpers.py | 14 +++++- tests/conftest.py | 9 +++- tests/cpu/_seq_classification_helpers.py | 11 +++- tests/cpu/test_adapter_cpu_accuracy.py | 8 ++- tests/cpu/test_embed_cpu_accuracy.py | 16 ++++-- tests/cpu/test_generate_cpu.py | 13 +++-- tests/cpu/test_load_cpu.py | 6 ++- tests/cpu/test_masked_lm_cpu_accuracy.py | 18 +++++-- .../test_question_answering_cpu_accuracy.py | 18 +++++-- .../test_token_classification_cpu_accuracy.py | 11 ++-- tests/cpu/test_vlm_e2e_cpu.py | 20 ++++++-- tests/model_registry.py | 32 ++++-------- tests/spyre/_seq_classification_helpers.py | 13 ++++- tests/spyre/edge_cases/_shared.py | 43 ++++++++++++---- tests/spyre/test_dspark_draft_spyre.py | 18 +++++-- .../spyre/test_e2e_masked_lm_compare_spyre.py | 18 +++++-- tests/spyre/test_e2e_multibatch_spyre.py | 10 ++-- ...st_e2e_question_answering_compare_spyre.py | 18 +++++-- tests/spyre/test_e2e_smoke_vision_spyre.py | 23 +++++++-- ..._e2e_token_classification_compare_spyre.py | 18 +++++-- tests/spyre/test_e2e_token_compare_spyre.py | 18 +++++-- tests/spyre/test_load_spyre.py | 50 ++++++++++++++----- tests/spyre/test_multicard_spyre.py | 12 ++++- tests/spyre/test_vlm_e2e_spyre.py | 21 ++++++-- 25 files changed, 332 insertions(+), 120 deletions(-) diff --git a/hf_adapters/st_backend.py b/hf_adapters/st_backend.py index 929baf3d..66412dd2 100644 --- a/hf_adapters/st_backend.py +++ b/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 diff --git a/tests/_vision_helpers.py b/tests/_vision_helpers.py index 3fc87514..4f26020a 100644 --- a/tests/_vision_helpers.py +++ b/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, ) diff --git a/tests/conftest.py b/tests/conftest.py index 5fd7fc85..180ec76c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -308,13 +308,15 @@ 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 - trust_remote_code = model_path in REMOTE_CODE_PATHS + 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 ) @@ -334,11 +336,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=str(model_name_or_path) in REMOTE_CODE_PATHS, + trust_remote_code=trust_remote_code, ) diff --git a/tests/cpu/_seq_classification_helpers.py b/tests/cpu/_seq_classification_helpers.py index fc09093c..fa51a5e7 100644 --- a/tests/cpu/_seq_classification_helpers.py +++ b/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) diff --git a/tests/cpu/test_adapter_cpu_accuracy.py b/tests/cpu/test_adapter_cpu_accuracy.py index 6fb95fef..41747050 100644 --- a/tests/cpu/test_adapter_cpu_accuracy.py +++ b/tests/cpu/test_adapter_cpu_accuracy.py @@ -188,8 +188,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( diff --git a/tests/cpu/test_embed_cpu_accuracy.py b/tests/cpu/test_embed_cpu_accuracy.py index 33326006..35ec6400 100644 --- a/tests/cpu/test_embed_cpu_accuracy.py +++ b/tests/cpu/test_embed_cpu_accuracy.py @@ -85,14 +85,22 @@ def _run_prefill( def test_auto_loader(model_path: str) -> 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) + 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(): @@ -104,7 +112,7 @@ def test_auto_loader(model_path: str) -> None: # Auto-loader path model = auto_spyre_model.AutoSpyreModel.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) _unwrap_compiled_blocks(model) with torch.no_grad(): diff --git a/tests/cpu/test_generate_cpu.py b/tests/cpu/test_generate_cpu.py index dcadb546..7b4741d5 100644 --- a/tests/cpu/test_generate_cpu.py +++ b/tests/cpu/test_generate_cpu.py @@ -58,27 +58,30 @@ def test_multibatch(model_path: str) -> 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) + 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, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + 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", - trust_remote_code=(model_path in REMOTE_CODE_PATHS), + trust_remote_code=trust_remote_code, ) _set_rope_dtype(model, dtype) sequences = hf_common_mod.generate( diff --git a/tests/cpu/test_load_cpu.py b/tests/cpu/test_load_cpu.py index b49823c6..a806c2c6 100644 --- a/tests/cpu/test_load_cpu.py +++ b/tests/cpu/test_load_cpu.py @@ -64,10 +64,12 @@ def test_load_embedding(model_path): gc.collect() -def load_embedding(model_path: str) -> Any: +def load_embedding(model_path: str, trust_remote_code: bool | None = None) -> Any: auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS model = auto_spyre_model.AutoSpyreModel.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) return model diff --git a/tests/cpu/test_masked_lm_cpu_accuracy.py b/tests/cpu/test_masked_lm_cpu_accuracy.py index 82d9b794..783004fa 100644 --- a/tests/cpu/test_masked_lm_cpu_accuracy.py +++ b/tests/cpu/test_masked_lm_cpu_accuracy.py @@ -22,7 +22,7 @@ resolve_adapter_module_for_test, ) from tests.cpu.conftest import _unwrap_compiled_blocks -from tests.model_registry import MASKED_LM_PATHS +from tests.model_registry import MASKED_LM_PATHS, REMOTE_CODE_PATHS pytestmark = pytest.mark.model_harness("masked_lm") @@ -36,9 +36,16 @@ @pytest.mark.parametrize("model_path", MASKED_LM_PATHS, ids=MASKED_LM_PATHS) def test_auto_loader(model_path: str) -> None: auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - dtype = dtype_for_model_path(model_path, target_device="cpu") - adapter_module = resolve_adapter_module_for_test(model_path) - tokenizer = AutoTokenizer.from_pretrained(model_path) + 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 + ) + adapter_module = resolve_adapter_module_for_test( + model_path, trust_remote_code=trust_remote_code + ) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) assert tokenizer.mask_token is not None encoded = tokenizer( @@ -52,6 +59,7 @@ def test_auto_loader(model_path: str) -> None: model_path=model_path, adapter_mod=adapter_module, auto_model_cls=AutoModelForMaskedLM, + trust_remote_code=trust_remote_code, ) with torch.no_grad(): ref_logits = ref_model(**encoded, return_dict=True).logits.float() @@ -59,7 +67,7 @@ def test_auto_loader(model_path: str) -> None: gc.collect() model = auto_spyre_model.AutoSpyreModelForMaskedLM.from_pretrained( - model_path, dtype=dtype + model_path, dtype=dtype, trust_remote_code=trust_remote_code ) _unwrap_compiled_blocks(model) with torch.no_grad(): diff --git a/tests/cpu/test_question_answering_cpu_accuracy.py b/tests/cpu/test_question_answering_cpu_accuracy.py index b35a4e66..acab68ac 100644 --- a/tests/cpu/test_question_answering_cpu_accuracy.py +++ b/tests/cpu/test_question_answering_cpu_accuracy.py @@ -22,7 +22,7 @@ resolve_adapter_module_for_test, ) from tests.cpu.conftest import _unwrap_compiled_blocks -from tests.model_registry import QUESTION_ANSWERING_PATHS +from tests.model_registry import QUESTION_ANSWERING_PATHS, REMOTE_CODE_PATHS pytestmark = pytest.mark.model_harness("question_answering") @@ -39,9 +39,16 @@ ) def test_native_forward(model_path: str) -> None: auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - dtype = dtype_for_model_path(model_path, target_device="cpu") - adapter_module = resolve_adapter_module_for_test(model_path) - tokenizer = AutoTokenizer.from_pretrained(model_path) + 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 + ) + adapter_module = resolve_adapter_module_for_test( + model_path, trust_remote_code=trust_remote_code + ) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) encoded = tokenizer( QUESTIONS, CONTEXTS, @@ -55,6 +62,7 @@ def test_native_forward(model_path: str) -> None: model_path=model_path, adapter_mod=adapter_module, auto_model_cls=AutoModelForQuestionAnswering, + trust_remote_code=trust_remote_code, ) with torch.no_grad(): ref_outputs = ref_model(**encoded, return_dict=True) @@ -62,7 +70,7 @@ def test_native_forward(model_path: str) -> None: gc.collect() model = auto_spyre_model.AutoSpyreModelForQuestionAnswering.from_pretrained( - model_path, dtype=dtype + model_path, dtype=dtype, trust_remote_code=trust_remote_code ) _unwrap_compiled_blocks(model) with torch.no_grad(): diff --git a/tests/cpu/test_token_classification_cpu_accuracy.py b/tests/cpu/test_token_classification_cpu_accuracy.py index a23ea044..c6f0919f 100644 --- a/tests/cpu/test_token_classification_cpu_accuracy.py +++ b/tests/cpu/test_token_classification_cpu_accuracy.py @@ -22,7 +22,7 @@ resolve_adapter_module_for_test, ) from tests.cpu.conftest import _unwrap_compiled_blocks -from tests.model_registry import TOKEN_CLASSIFICATION_PATHS +from tests.model_registry import REMOTE_CODE_PATHS, TOKEN_CLASSIFICATION_PATHS pytestmark = pytest.mark.model_harness("token_classification") @@ -38,12 +38,16 @@ ) def test_native_forward(model_path: str) -> None: auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] + trust_remote_code = model_path in REMOTE_CODE_PATHS dtype = get_dtype_for_cpu(model_path) adapter_module = resolve_adapter_module_for_test( model_path, mapping=auto_spyre_model.TOKEN_CLASSIFICATION_CONFIG_TO_ADAPTER_MODULE_MAPPING, + trust_remote_code=trust_remote_code, + ) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code ) - tokenizer = AutoTokenizer.from_pretrained(model_path) encoded = tokenizer( SENTENCES, return_tensors="pt", @@ -55,6 +59,7 @@ def test_native_forward(model_path: str) -> None: model_path=model_path, adapter_mod=adapter_module, auto_model_cls=AutoModelForTokenClassification, + trust_remote_code=trust_remote_code, ) with torch.no_grad(): ref_logits = ref_model(**encoded, return_dict=True).logits.float() @@ -62,7 +67,7 @@ def test_native_forward(model_path: str) -> None: gc.collect() model = auto_spyre_model.AutoSpyreModelForTokenClassification.from_pretrained( - model_path, dtype=dtype + model_path, dtype=dtype, trust_remote_code=trust_remote_code ) _unwrap_compiled_blocks(model) with torch.no_grad(): diff --git a/tests/cpu/test_vlm_e2e_cpu.py b/tests/cpu/test_vlm_e2e_cpu.py index 49f47c14..3c76b5c9 100644 --- a/tests/cpu/test_vlm_e2e_cpu.py +++ b/tests/cpu/test_vlm_e2e_cpu.py @@ -48,7 +48,7 @@ resolve_adapter_module_for_test, ) from tests.cpu.conftest import _set_rope_dtype, _unwrap_compiled_blocks -from tests.model_registry import VISION_PATHS +from tests.model_registry import REMOTE_CODE_PATHS, VISION_PATHS pytestmark = pytest.mark.model_harness("vision") @@ -60,12 +60,19 @@ def test_vlm_generate(model_path: str) -> None: from hf_adapters.auto_spyre_model import dtype_for_model_path + trust_remote_code = model_path in REMOTE_CODE_PATHS adapter = resolve_adapter_module_for_test( - model_path, mapping=IMAGE_TEXT_TO_TEXT_CONFIG_TO_ADAPTER_MODULE_MAPPING + model_path, + mapping=IMAGE_TEXT_TO_TEXT_CONFIG_TO_ADAPTER_MODULE_MAPPING, + trust_remote_code=trust_remote_code, + ) + dtype = dtype_for_model_path( + model_path, target_device="cpu", trust_remote_code=trust_remote_code ) - dtype = dtype_for_model_path(model_path, target_device="cpu") - processor, batch = build_vlm_batch(model_path, PROMPT) + processor, batch = build_vlm_batch( + model_path, PROMPT, trust_remote_code=trust_remote_code + ) batch["pixel_values"] = batch["pixel_values"].to(dtype) # --- Stock reference: the FULL model.generate() (real deepstack) --- @@ -76,11 +83,14 @@ def test_vlm_generate(model_path: str) -> None: batch=batch, max_new_tokens=MAX_NEW_TOKENS, adapter_mod=adapter, + trust_remote_code=trust_remote_code, ) gc.collect() # --- Adapter generate (greedy) --- - model = AutoSpyreModelForImageTextToText.from_pretrained(model_path, dtype=dtype) + model = AutoSpyreModelForImageTextToText.from_pretrained( + model_path, dtype=dtype, trust_remote_code=trust_remote_code + ) _set_rope_dtype(model, dtype) _unwrap_compiled_blocks(model) prompt_len = batch["input_ids"].shape[1] diff --git a/tests/model_registry.py b/tests/model_registry.py index 2e4ccf2b..b89b1d7f 100644 --- a/tests/model_registry.py +++ b/tests/model_registry.py @@ -821,14 +821,13 @@ def _all_paths( ] -def _collect_remote_code_paths() -> frozenset[str]: - """Paths of every registered checkpoint that needs ``trust_remote_code``. - - Spans all category registries so a single lookup covers any harness. Test - call sites check membership in ``REMOTE_CODE_PATHS`` to decide whether to - forward ``trust_remote_code=True`` to ``from_pretrained``. - """ - registries = ( +# Paths that must be loaded with ``trust_remote_code=True`` +# Spans all category registries so a single lookup covers any harness. Test +# call sites check membership in ``REMOTE_CODE_PATHS`` to decide whether to +# forward ``trust_remote_code=True`` to ``from_pretrained``. +REMOTE_CODE_PATHS: frozenset[str] = frozenset( + info["path"] + for models in ( CAUSAL_LM_MODELS, EMBEDDING_MODELS, MASKED_LM_MODELS, @@ -836,20 +835,9 @@ def _collect_remote_code_paths() -> frozenset[str]: TOKEN_CLASSIFICATION_MODELS, VISION_MODELS, ) - paths: set[str] = set() - for models in registries: - with_trc = _all_paths( - models, include_gated=True, include_trust_remote_code=True - ) - without_trc = _all_paths( - models, include_gated=True, include_trust_remote_code=False - ) - paths.update(set(with_trc) - set(without_trc)) - return frozenset(paths) - - -# Paths that must be loaded with ``trust_remote_code=True`` -REMOTE_CODE_PATHS: frozenset[str] = _collect_remote_code_paths() + for info in models.values() + if info.get("trust_remote_code", False) +) # Every registered path per category, bypassing the smallest-per-adapter diff --git a/tests/spyre/_seq_classification_helpers.py b/tests/spyre/_seq_classification_helpers.py index d72cc951..37f9b676 100644 --- a/tests/spyre/_seq_classification_helpers.py +++ b/tests/spyre/_seq_classification_helpers.py @@ -32,12 +32,14 @@ from hf_adapters.auto_spyre_model import dtype_for_model_path from hf_adapters.hf_common import move_model_to_spyre, prefill_sequence_classification from tests.conftest import load_ref_model +from tests.model_registry import REMOTE_CODE_PATHS def run_seq_classification_cpu_vs_spyre( model_path: str, adapter: types.ModuleType, inputs: list[str] | list[tuple[str, str]], + trust_remote_code: bool | None = None, ) -> dict: """Load a seq-classification model, run a CPU reference forward, then run the adapter on Spyre via ``prefill_sequence_classification``. @@ -55,18 +57,25 @@ def run_seq_classification_cpu_vs_spyre( ``spyre_logits`` – ``[B, num_labels]`` float CPU tensor (adapter on Spyre). ``dtype`` – dtype used for the Spyre model. """ - dtype = dtype_for_model_path(model_path, target_device="spyre") + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS + dtype = dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ) print(f"\n{'=' * 70}") print(f" {model_path}") print(f" dtype: {dtype}") print(f"{'=' * 70}") - tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) model = load_ref_model( model_path=model_path, adapter_mod=adapter, auto_model_cls=AutoModelForSequenceClassification, + trust_remote_code=trust_remote_code, ) encoded = tokenizer( diff --git a/tests/spyre/edge_cases/_shared.py b/tests/spyre/edge_cases/_shared.py index 829e076e..e3371f15 100644 --- a/tests/spyre/edge_cases/_shared.py +++ b/tests/spyre/edge_cases/_shared.py @@ -51,11 +51,15 @@ from tests.model_registry import REMOTE_CODE_PATHS -def _load_spyre_model(model_path: str) -> PreTrainedModel: +def _load_spyre_model( + model_path: str, trust_remote_code: bool | None = None +) -> PreTrainedModel: print(f" Loading {model_path} on Spyre ...") + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS t0 = time.time() model = AutoSpyreModelForCausalLM.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) print(f" Spyre load+prepare: {time.time() - t0:.1f}s") return model @@ -64,14 +68,25 @@ def _load_spyre_model(model_path: str) -> PreTrainedModel: def _setup( model_path: str, need_ref: bool, + trust_remote_code: bool | None = None, ): + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS tokenizer = AutoTokenizer.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code + ) + adapter = resolve_adapter_module_for_test( + model_path, trust_remote_code=trust_remote_code ) - adapter = resolve_adapter_module_for_test(model_path) - ref_model = load_ref_model(model_path, adapter_mod=adapter) if need_ref else None - spyre_model = _load_spyre_model(model_path) + ref_model = ( + load_ref_model( + model_path, adapter_mod=adapter, trust_remote_code=trust_remote_code + ) + if need_ref + else None + ) + spyre_model = _load_spyre_model(model_path, trust_remote_code=trust_remote_code) return model_path, tokenizer, ref_model, spyre_model @@ -85,9 +100,13 @@ def _teardown( gc.collect() -def run_greedy_case(model_path: str, case_id: str) -> tuple[bool, str]: +def run_greedy_case( + model_path: str, case_id: str, trust_remote_code: bool | None = None +) -> tuple[bool, str]: """Greedy-generate case: HF reference == Spyre output, per row.""" - info, tokenizer, ref_model, model = _setup(model_path, need_ref=True) + info, tokenizer, ref_model, model = _setup( + model_path, need_ref=True, trust_remote_code=trust_remote_code + ) try: targets, max_new = CASES[case_id] prompts = make_prompts(tokenizer, targets) @@ -111,9 +130,13 @@ def run_greedy_case(model_path: str, case_id: str) -> tuple[bool, str]: _teardown(model, ref_model) -def run_eos_case(model_path: str, case_id: str) -> tuple[bool, str]: +def run_eos_case( + model_path: str, case_id: str, trust_remote_code: bool | None = None +) -> tuple[bool, str]: """Forced-EOS case: shared eos_token_id stops each row at its requested offset.""" - info, tokenizer, ref_model, model = _setup(model_path, need_ref=True) + info, tokenizer, ref_model, model = _setup( + model_path, need_ref=True, trust_remote_code=trust_remote_code + ) try: eos_offsets, max_new = EOS_CASES[case_id] batch_size = len(eos_offsets) diff --git a/tests/spyre/test_dspark_draft_spyre.py b/tests/spyre/test_dspark_draft_spyre.py index 94bc07c6..0e1db2eb 100644 --- a/tests/spyre/test_dspark_draft_spyre.py +++ b/tests/spyre/test_dspark_draft_spyre.py @@ -34,7 +34,7 @@ import pytest import torch -from tests.model_registry import DSPARK_PATHS +from tests.model_registry import DSPARK_PATHS, REMOTE_CODE_PATHS pytest.importorskip("deepspec", reason="DSpark drafter modeling requires DeepSpec") @@ -52,16 +52,26 @@ def test_dspark_draft_block(ckpt): dev = torch.device("spyre:0") hf_common.DEVICE = dev + trust_remote_code = ckpt in REMOTE_CODE_PATHS + # The library resolves the checkpoint to its adapter by architecture # (``*DSparkModel``); confirm it lands on a DSpark draft adapter module. - resolved = resolve_adapter_module(ckpt) + resolved = resolve_adapter_module(ckpt, trust_remote_code=trust_remote_code) assert resolved.__name__.rsplit(".", 1)[-1].startswith( "hf_dspark_" ), f"{ckpt} resolved to {resolved.__name__}, expected a DSpark draft adapter" - arch = (AutoConfig.from_pretrained(ckpt).architectures or [""])[0] + arch = ( + AutoConfig.from_pretrained( + ckpt, trust_remote_code=trust_remote_code + ).architectures + or [""] + )[0] model = AutoModelForCausalLM.from_pretrained( - ckpt, dtype=torch.float16, attn_implementation="sdpa" + ckpt, + dtype=torch.float16, + attn_implementation="sdpa", + trust_remote_code=trust_remote_code, ).eval() model.requires_grad_(False) diff --git a/tests/spyre/test_e2e_masked_lm_compare_spyre.py b/tests/spyre/test_e2e_masked_lm_compare_spyre.py index 27d2d685..eba3ab34 100644 --- a/tests/spyre/test_e2e_masked_lm_compare_spyre.py +++ b/tests/spyre/test_e2e_masked_lm_compare_spyre.py @@ -17,7 +17,7 @@ from hf_adapters import AutoSpyreModelForMaskedLM from hf_adapters.auto_spyre_model import dtype_for_model_path -from tests.model_registry import MASKED_LM_PATHS +from tests.model_registry import MASKED_LM_PATHS, REMOTE_CODE_PATHS pytestmark = pytest.mark.model_harness("masked_lm") @@ -30,7 +30,10 @@ @pytest.mark.parametrize("model_path", MASKED_LM_PATHS, ids=MASKED_LM_PATHS) def test_e2e_masked_lm_compare_spyre(model_path: str) -> None: - tokenizer = AutoTokenizer.from_pretrained(model_path) + trust_remote_code = model_path in REMOTE_CODE_PATHS + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) assert tokenizer.mask_token is not None encoded = tokenizer( [prompt.format(mask=tokenizer.mask_token) for prompt in PROMPTS], @@ -41,8 +44,11 @@ def test_e2e_masked_lm_compare_spyre(model_path: str) -> None: ref_model = AutoModelForMaskedLM.from_pretrained( model_path, - 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 + ), device_map="cpu", + trust_remote_code=trust_remote_code, ).eval() with torch.no_grad(): ref_logits = ref_model(**encoded, return_dict=True).logits.float() @@ -50,7 +56,11 @@ def test_e2e_masked_lm_compare_spyre(model_path: str) -> None: gc.collect() model = AutoSpyreModelForMaskedLM.from_pretrained( - model_path, dtype=dtype_for_model_path(model_path, target_device="spyre") + model_path, + dtype=dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ), + trust_remote_code=trust_remote_code, ) with torch.no_grad(): logits = model(**encoded, return_dict=True).logits.float() diff --git a/tests/spyre/test_e2e_multibatch_spyre.py b/tests/spyre/test_e2e_multibatch_spyre.py index bf14acc5..4fbbcdc2 100644 --- a/tests/spyre/test_e2e_multibatch_spyre.py +++ b/tests/spyre/test_e2e_multibatch_spyre.py @@ -97,12 +97,15 @@ def _print_table(model_path: str, rows: list[dict[str, Any]]) -> None: "model_path", xfail_non_blocking(CAUSAL_PATHS, table=NON_BLOCKING_CAUSAL_MODELS) ) def test_e2e_multibatch_spyre(model_path: str) -> None: - adapter_mod = resolve_adapter_module_for_test(model_path) + trust_remote_code = model_path in REMOTE_CODE_PATHS + adapter_mod = resolve_adapter_module_for_test( + model_path, trust_remote_code=trust_remote_code + ) assert len(PROMPTS) > 1, "multi-batch test needs batch_size > 1" tokenizer = AutoTokenizer.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) print(f"\n{'=' * 70}") @@ -111,7 +114,7 @@ def test_e2e_multibatch_spyre(model_path: str) -> None: # HF reference (per-prompt) — informational only. Run BEFORE prepare_for_spyre # patches RMSNorm globally. Loaded on CPU via load_ref_model, then discarded. - model = load_ref_model(model_path, adapter_mod) + model = load_ref_model(model_path, adapter_mod, trust_remote_code=trust_remote_code) print(" Running HF reference on CPU (per-prompt, informational) ...") hf_outputs = hf_reference_outputs(model, tokenizer, PROMPTS, MAX_NEW_TOKENS) del model @@ -119,7 +122,6 @@ def test_e2e_multibatch_spyre(model_path: str) -> None: # One Spyre model, reused for the batched run and each single run (generate # allocates a fresh KV cache per call, so runs do not contaminate each other). - trust_remote_code = model_path in REMOTE_CODE_PATHS spyre_dtype = dtype_for_model_path( model_path, target_device="spyre", trust_remote_code=trust_remote_code ) diff --git a/tests/spyre/test_e2e_question_answering_compare_spyre.py b/tests/spyre/test_e2e_question_answering_compare_spyre.py index e52fc3b1..101a9a20 100644 --- a/tests/spyre/test_e2e_question_answering_compare_spyre.py +++ b/tests/spyre/test_e2e_question_answering_compare_spyre.py @@ -17,7 +17,7 @@ from hf_adapters import AutoSpyreModelForQuestionAnswering from hf_adapters.auto_spyre_model import dtype_for_model_path -from tests.model_registry import QUESTION_ANSWERING_PATHS +from tests.model_registry import QUESTION_ANSWERING_PATHS, REMOTE_CODE_PATHS pytestmark = pytest.mark.model_harness("question_answering") @@ -33,7 +33,10 @@ "model_path", QUESTION_ANSWERING_PATHS, ids=QUESTION_ANSWERING_PATHS ) def test_e2e_question_answering_compare_spyre(model_path: str) -> None: - tokenizer = AutoTokenizer.from_pretrained(model_path) + trust_remote_code = model_path in REMOTE_CODE_PATHS + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) encoded = tokenizer( QUESTIONS, CONTEXTS, @@ -45,8 +48,11 @@ def test_e2e_question_answering_compare_spyre(model_path: str) -> None: ref_model = AutoModelForQuestionAnswering.from_pretrained( model_path, - 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 + ), device_map="cpu", + trust_remote_code=trust_remote_code, ).eval() with torch.no_grad(): ref_outputs = ref_model(**encoded, return_dict=True) @@ -54,7 +60,11 @@ def test_e2e_question_answering_compare_spyre(model_path: str) -> None: gc.collect() model = AutoSpyreModelForQuestionAnswering.from_pretrained( - model_path, dtype=dtype_for_model_path(model_path, target_device="spyre") + model_path, + dtype=dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ), + trust_remote_code=trust_remote_code, ) with torch.no_grad(): outputs = model(**encoded, return_dict=True) diff --git a/tests/spyre/test_e2e_smoke_vision_spyre.py b/tests/spyre/test_e2e_smoke_vision_spyre.py index 69df3487..b9dfc7ac 100644 --- a/tests/spyre/test_e2e_smoke_vision_spyre.py +++ b/tests/spyre/test_e2e_smoke_vision_spyre.py @@ -53,6 +53,7 @@ from tests._vision_helpers import build_vlm_batch, load_smoke_test_images from tests.model_registry import ( NON_BLOCKING_VISION_MODELS, + REMOTE_CODE_PATHS, VISION_PATHS, xfail_non_blocking, ) @@ -71,6 +72,7 @@ def _run_single_image( label: str, prompt: str, image: Any, + trust_remote_code: bool | None = None, ) -> dict[str, Any]: """Run generate for one image and return a per-image result dict. @@ -79,7 +81,9 @@ def _run_single_image( TTFT and ITL are printed by the shared generate loop via ``timing=True``. """ - processor_i, batch = build_vlm_batch(model_path, prompt, image) + processor_i, batch = build_vlm_batch( + model_path, prompt, image, trust_remote_code=trust_remote_code + ) batch["pixel_values"] = batch["pixel_values"].to(dtype) # --- Full generation (timing printed by the shared generate loop) --------- @@ -136,11 +140,14 @@ def run_vision_smoke_test(model_path: str) -> dict[str, Any]: print(f" loading from {model_path}") print(f"{'=' * 70}") - dtype = dtype_for_model_path(model_path, target_device="spyre") + trust_remote_code = model_path in REMOTE_CODE_PATHS + dtype = dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ) t0 = time.time() model = AutoSpyreModelForImageTextToText.from_pretrained( - model_name_or_path=model_path, dtype=dtype + model_name_or_path=model_path, dtype=dtype, trust_remote_code=trust_remote_code ) load_time = time.time() - t0 print(f" Load time: {load_time:.1f}s") @@ -152,7 +159,15 @@ def run_vision_smoke_test(model_path: str) -> dict[str, Any]: image_results = [] for label, prompt, image in smoke_images: print(f"\n [{label}] prompt: {prompt!r}") - result = _run_single_image(model, dtype, model_path, label, prompt, image) + result = _run_single_image( + model, + dtype, + model_path, + label, + prompt, + image, + trust_remote_code=trust_remote_code, + ) print(f" [{label}] output: {result['text']!r}") print(f" [{label}] gen: {result['gen_s']:.1f}s status: {result['status']}") image_results.append(result) diff --git a/tests/spyre/test_e2e_token_classification_compare_spyre.py b/tests/spyre/test_e2e_token_classification_compare_spyre.py index 0481350c..b6097c9c 100644 --- a/tests/spyre/test_e2e_token_classification_compare_spyre.py +++ b/tests/spyre/test_e2e_token_classification_compare_spyre.py @@ -34,7 +34,7 @@ import pytest import torch import torch.nn.functional as F -from model_registry import TOKEN_CLASSIFICATION_PATHS +from model_registry import REMOTE_CODE_PATHS, TOKEN_CLASSIFICATION_PATHS from transformers import AutoModelForTokenClassification, AutoTokenizer from hf_adapters import AutoSpyreModelForTokenClassification @@ -55,7 +55,10 @@ "model_path", TOKEN_CLASSIFICATION_PATHS, ids=TOKEN_CLASSIFICATION_PATHS ) def test_e2e_token_classification_compare_spyre(model_path: str) -> None: - tokenizer = AutoTokenizer.from_pretrained(model_path) + trust_remote_code = model_path in REMOTE_CODE_PATHS + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) encoded = tokenizer( SENTENCES, return_tensors="pt", @@ -65,8 +68,11 @@ def test_e2e_token_classification_compare_spyre(model_path: str) -> None: ref_model = AutoModelForTokenClassification.from_pretrained( model_path, - 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 + ), device_map="cpu", + trust_remote_code=trust_remote_code, ).eval() with torch.no_grad(): ref_logits = ref_model(**encoded, return_dict=True).logits.float() @@ -74,7 +80,11 @@ def test_e2e_token_classification_compare_spyre(model_path: str) -> None: gc.collect() model = AutoSpyreModelForTokenClassification.from_pretrained( - model_path, dtype=dtype_for_model_path(model_path, target_device="spyre") + model_path, + dtype=dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code + ), + trust_remote_code=trust_remote_code, ) with torch.no_grad(): logits = model(**encoded, return_dict=True).logits.float() diff --git a/tests/spyre/test_e2e_token_compare_spyre.py b/tests/spyre/test_e2e_token_compare_spyre.py index 16b6a818..68127c6a 100644 --- a/tests/spyre/test_e2e_token_compare_spyre.py +++ b/tests/spyre/test_e2e_token_compare_spyre.py @@ -272,20 +272,28 @@ def _print_table(rows: list[dict[str, Any]]) -> None: ) -def _run_model_test(model_path: str, num_decode: int = 4) -> list[dict[str, Any]]: +def _run_model_test( + model_path: str, num_decode: int = 4, trust_remote_code: bool | None = None +) -> list[dict[str, Any]]: """Full comparison for one model. Returns the list of comparison rows.""" from transformers import AutoTokenizer - adapter = resolve_adapter_module_for_test(model_path) + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS + adapter = resolve_adapter_module_for_test( + model_path, trust_remote_code=trust_remote_code + ) print(f"\n{'=' * 70}") print(f" {model_path}") print(f"{'=' * 70}") tokenizer = AutoTokenizer.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code + ) + model = load_ref_model( + model_path=model_path, adapter_mod=adapter, trust_remote_code=trust_remote_code ) - model = load_ref_model(model_path=model_path, adapter_mod=adapter) prompt = "The capital of France is" # Tokenize following the model's canonical scheme (chat template for @@ -303,7 +311,7 @@ def _run_model_test(model_path: str, num_decode: int = 4) -> list[dict[str, Any] spyre_dtype = dtype_for_model_path( model_path, target_device="spyre", - trust_remote_code=(model_path in REMOTE_CODE_PATHS), + trust_remote_code=trust_remote_code, ) move_model_to_spyre(model=model, module=adapter, dtype=spyre_dtype) print(" Running adapter on Spyre ...") diff --git a/tests/spyre/test_load_spyre.py b/tests/spyre/test_load_spyre.py index 24d509bd..19e4ad20 100644 --- a/tests/spyre/test_load_spyre.py +++ b/tests/spyre/test_load_spyre.py @@ -70,12 +70,16 @@ def test_load_causal_lm(model_path: str) -> None: ), f"{model_path}: AutoSpyreModelForCausalLM did not attach generate()" -def load_causal_lm(model_path: str) -> tuple[Any, Any, float]: +def load_causal_lm( + model_path: str, trust_remote_code: bool | None = None +) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForCausalLM + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS t0 = time.time() model = AutoSpyreModelForCausalLM.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) load_s = time.time() - t0 @@ -85,12 +89,16 @@ def load_causal_lm(model_path: str) -> tuple[Any, Any, float]: return model_is_not_none, callables, load_s -def load_embedding(model_path: str) -> tuple[Any, float]: +def load_embedding( + model_path: str, trust_remote_code: bool | None = None +) -> tuple[Any, float]: from hf_adapters import AutoSpyreModel + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS t0 = time.time() model = AutoSpyreModel.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) load_s = time.time() - t0 return model is not None, load_s @@ -109,11 +117,14 @@ def test_load_embedding(model_path: str) -> None: print(f"| {model_path} | embedding | PASS | {load_s:.1f} |") -def load_masked_lm(model_path: str) -> tuple[Any, Any, float]: +def load_masked_lm( + model_path: str, trust_remote_code: bool | None = None +) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForMaskedLM from hf_adapters.auto_spyre_model import dtype_for_model_path - trust_remote_code = model_path in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS dtype = dtype_for_model_path( model_path, target_device="spyre", trust_remote_code=trust_remote_code ) @@ -144,11 +155,14 @@ def test_load_masked_lm(model_path: str) -> None: assert callables, f"{model_path}: AutoSpyreModelForMaskedLM forward is not callable" -def load_question_answering(model_path: str) -> tuple[Any, Any, float]: +def load_question_answering( + model_path: str, trust_remote_code: bool | None = None +) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForQuestionAnswering from hf_adapters.auto_spyre_model import dtype_for_model_path - trust_remote_code = model_path in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS dtype = dtype_for_model_path( model_path, target_device="spyre", trust_remote_code=trust_remote_code ) @@ -172,11 +186,17 @@ def test_load_question_answering(model_path: str) -> None: assert ready, f"{model_path}: native forward or CPU QA head is not ready" -def load_seq_classification(model_path: str) -> tuple[Any, Any, float]: +def load_seq_classification( + model_path: str, trust_remote_code: bool | None = None +) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForSequenceClassification + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS t0 = time.time() - model: Any = AutoSpyreModelForSequenceClassification.from_pretrained(model_path) + model: Any = AutoSpyreModelForSequenceClassification.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) load_s = time.time() - t0 head_on_cpu = next(model.classifier.parameters()).device.type == "cpu" return model is not None, callable(model.forward) and head_on_cpu, load_s @@ -197,11 +217,17 @@ def test_load_seq_classification(model_path: str) -> None: assert ready, f"{model_path}: native forward or CPU classifier head is not ready" -def load_token_classification(model_path: str) -> tuple[Any, Any, float]: +def load_token_classification( + model_path: str, trust_remote_code: bool | None = None +) -> tuple[Any, Any, float]: from hf_adapters import AutoSpyreModelForTokenClassification + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS t0 = time.time() - model: Any = AutoSpyreModelForTokenClassification.from_pretrained(model_path) + model: Any = AutoSpyreModelForTokenClassification.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) load_s = time.time() - t0 head_on_cpu = next(model.classifier.parameters()).device.type == "cpu" return model is not None, callable(model.forward) and head_on_cpu, load_s diff --git a/tests/spyre/test_multicard_spyre.py b/tests/spyre/test_multicard_spyre.py index 2ed780b6..7782149c 100644 --- a/tests/spyre/test_multicard_spyre.py +++ b/tests/spyre/test_multicard_spyre.py @@ -146,6 +146,7 @@ def run_multicard_smoke_test( max_new_tokens: int = _DEFAULT_MAX_NEW_TOKENS, dtype: "torch.dtype | None" = None, batch_size: int = _DEFAULT_BATCH_SIZE, + trust_remote_code: bool | None = None, ) -> dict[str, Any]: """Load model and generate tokens; return a diagnostics dict. @@ -251,17 +252,24 @@ def run_multicard_smoke_test( # ── Phase 1: model load ──────────────────────────────────────────────── print(f"\n{'=' * 20} Loading Model...") + if trust_remote_code is None: + from tests.model_registry import REMOTE_CODE_PATHS + + trust_remote_code = model_path in REMOTE_CODE_PATHS + model = None tokenizer = None load_t0 = time.time() try: tp = "auto" if world_size > 1 else None - kwargs: dict[str, Any] = {"tp_plan": tp} + kwargs: dict[str, Any] = {"tp_plan": tp, "trust_remote_code": trust_remote_code} if dtype is not None: kwargs["dtype"] = dtype model = AutoSpyreModelForCausalLM.from_pretrained(model_path, **kwargs) result["load_s"] = time.time() - load_t0 - tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=trust_remote_code + ) print(f" Load time : {result['load_s']:.1f}s [OK]") except Exception: result["load_s"] = time.time() - load_t0 diff --git a/tests/spyre/test_vlm_e2e_spyre.py b/tests/spyre/test_vlm_e2e_spyre.py index 69c5f484..b3d6ad93 100644 --- a/tests/spyre/test_vlm_e2e_spyre.py +++ b/tests/spyre/test_vlm_e2e_spyre.py @@ -84,6 +84,7 @@ from tests.conftest import load_ref_model from tests.model_registry import ( NON_BLOCKING_VISION_MODELS, + REMOTE_CODE_PATHS, VISION_PATHS, xfail_non_blocking, ) @@ -218,6 +219,7 @@ def _stock_vlm_greedy_steps( adapter_mod, num_steps: int, ref_model=None, + trust_remote_code: bool | None = None, ) -> tuple[list[torch.Tensor], list[int]]: """Stock HF per-step greedy logits + token ids over prefill + decode. @@ -242,6 +244,7 @@ def _stock_vlm_greedy_steps( model_path=model_path, adapter_mod=adapter_mod, auto_model_cls=AutoModelForImageTextToText, + trust_remote_code=trust_remote_code, ) with torch.no_grad(): gen = ref_model.generate( @@ -262,12 +265,19 @@ def _stock_vlm_greedy_steps( "model_path", xfail_non_blocking(VISION_PATHS, table=NON_BLOCKING_VISION_MODELS) ) def test_vlm_generate_spyre(model_path: str) -> None: + trust_remote_code = model_path in REMOTE_CODE_PATHS adapter = resolve_adapter_module( - model_path, mapping=IMAGE_TEXT_TO_TEXT_CONFIG_TO_ADAPTER_MODULE_MAPPING + model_path, + mapping=IMAGE_TEXT_TO_TEXT_CONFIG_TO_ADAPTER_MODULE_MAPPING, + trust_remote_code=trust_remote_code, + ) + dtype = dtype_for_model_path( + model_path, target_device="spyre", trust_remote_code=trust_remote_code ) - dtype = dtype_for_model_path(model_path, target_device="spyre") - processor, batch = build_vlm_batch(model_path, PROMPT) + processor, batch = build_vlm_batch( + model_path, PROMPT, trust_remote_code=trust_remote_code + ) batch["pixel_values"] = batch["pixel_values"].to(dtype) tokenizer = processor.tokenizer @@ -287,6 +297,7 @@ def test_vlm_generate_spyre(model_path: str) -> None: model_path=model_path, adapter_mod=adapter, auto_model_cls=AutoModelForImageTextToText, + trust_remote_code=trust_remote_code, ) ref_logits, ref_tokens = _stock_vlm_greedy_steps( model_path=model_path, @@ -294,6 +305,7 @@ def test_vlm_generate_spyre(model_path: str) -> None: num_steps=NUM_COMPARE_STEPS, adapter_mod=adapter, ref_model=ref_model, + trust_remote_code=trust_remote_code, ) ref_text = stock_vlm_generate( model_path=model_path, @@ -302,6 +314,7 @@ def test_vlm_generate_spyre(model_path: str) -> None: max_new_tokens=MAX_NEW_TOKENS, adapter_mod=adapter, ref_model=ref_model, + trust_remote_code=trust_remote_code, ) del ref_model gc.collect() @@ -309,7 +322,7 @@ def test_vlm_generate_spyre(model_path: str) -> None: # --- Adapter on Spyre --- print(" Loading model for Spyre ...") model = AutoSpyreModelForImageTextToText.from_pretrained( - model_name_or_path=model_path, dtype=dtype + model_name_or_path=model_path, dtype=dtype, trust_remote_code=trust_remote_code ) # Per-step adapter logits on Spyre, teacher-forced on stock's tokens (so the From bb11adbe57e3b84573d28e13023c12604514e6aa Mon Sep 17 00:00:00 2001 From: Dushyant Behl Date: Fri, 11 Sep 2026 13:43:38 +0530 Subject: [PATCH 3/3] Trust remote code should be a method parameter Signed-off-by: Dushyant Behl --- tests/conftest.py | 22 ++++++++++ tests/cpu/test_adapter_cpu_accuracy.py | 5 ++- tests/cpu/test_embed_cpu_accuracy.py | 5 ++- tests/cpu/test_generate_cpu.py | 5 ++- tests/cpu/test_load_cpu.py | 26 ++++++----- tests/cpu/test_masked_lm_cpu_accuracy.py | 5 ++- .../test_question_answering_cpu_accuracy.py | 5 ++- tests/cpu/test_reranker_cpu_accuracy.py | 4 +- .../test_seq_classification_cpu_accuracy.py | 4 +- .../test_token_classification_cpu_accuracy.py | 5 ++- tests/cpu/test_vlm_e2e_cpu.py | 5 ++- tests/spyre/edge_cases/test_eos_spyre.py | 44 ++++++++++++++----- .../edge_cases/test_long_multi_block_spyre.py | 8 +++- tests/spyre/edge_cases/test_mixed_spyre.py | 14 ++++-- .../test_no_pad_token_fallback_spyre.py | 8 +++- .../test_prompt_exactly_block_spyre.py | 8 +++- tests/spyre/edge_cases/test_sampling_spyre.py | 16 +++++-- tests/spyre/edge_cases/test_short_spyre.py | 30 +++++++++---- .../test_single_token_prompt_spyre.py | 8 +++- tests/spyre/test_dspark_draft_spyre.py | 5 ++- .../spyre/test_e2e_masked_lm_compare_spyre.py | 7 ++- tests/spyre/test_e2e_multibatch_spyre.py | 5 ++- ...st_e2e_question_answering_compare_spyre.py | 7 ++- .../spyre/test_e2e_reranker_compare_spyre.py | 13 ++++-- ...st_e2e_seq_classification_compare_spyre.py | 13 ++++-- tests/spyre/test_e2e_smoke_spyre.py | 15 ++++--- tests/spyre/test_e2e_smoke_vision_spyre.py | 13 ++++-- ..._e2e_token_classification_compare_spyre.py | 7 ++- tests/spyre/test_e2e_token_compare_spyre.py | 11 +++-- tests/spyre/test_load_spyre.py | 42 +++++++++++++----- tests/spyre/test_multicard_spyre.py | 6 ++- tests/spyre/test_vlm_e2e_spyre.py | 5 ++- 32 files changed, 268 insertions(+), 108 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 180ec76c..d3b23730 100644 --- a/tests/conftest.py +++ b/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: diff --git a/tests/cpu/test_adapter_cpu_accuracy.py b/tests/cpu/test_adapter_cpu_accuracy.py index 41747050..a5c6eeaa 100644 --- a/tests/cpu/test_adapter_cpu_accuracy.py +++ b/tests/cpu/test_adapter_cpu_accuracy.py @@ -162,9 +162,10 @@ 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"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + 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 ) diff --git a/tests/cpu/test_embed_cpu_accuracy.py b/tests/cpu/test_embed_cpu_accuracy.py index 35ec6400..70cda977 100644 --- a/tests/cpu/test_embed_cpu_accuracy.py +++ b/tests/cpu/test_embed_cpu_accuracy.py @@ -82,10 +82,11 @@ 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"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + 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 ) diff --git a/tests/cpu/test_generate_cpu.py b/tests/cpu/test_generate_cpu.py index 7b4741d5..f6e44383 100644 --- a/tests/cpu/test_generate_cpu.py +++ b/tests/cpu/test_generate_cpu.py @@ -54,11 +54,12 @@ @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"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + 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 ) diff --git a/tests/cpu/test_load_cpu.py b/tests/cpu/test_load_cpu.py index a806c2c6..c1e2082a 100644 --- a/tests/cpu/test_load_cpu.py +++ b/tests/cpu/test_load_cpu.py @@ -42,10 +42,12 @@ @pytest.mark.model_harness("causal") @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) -def test_load_causal_lm(model_path): +def test_load_causal_lm(model_path, trust_remote_code): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS model = auto_spyre_model.AutoSpyreModelForCausalLM.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) assert model is not None assert callable( @@ -57,8 +59,8 @@ def test_load_causal_lm(model_path): @pytest.mark.model_harness("embedding") @pytest.mark.parametrize("model_path", EMBED_PATHS, ids=EMBED_PATHS) -def test_load_embedding(model_path): - model = load_embedding(model_path=model_path) +def test_load_embedding(model_path, trust_remote_code): + model = load_embedding(model_path=model_path, trust_remote_code=trust_remote_code) assert model is not None del model gc.collect() @@ -76,9 +78,10 @@ def load_embedding(model_path: str, trust_remote_code: bool | None = None) -> An @pytest.mark.model_harness("masked_lm") @pytest.mark.parametrize("model_path", MASKED_LM_PATHS, ids=MASKED_LM_PATHS) -def test_load_masked_lm(model_path): +def test_load_masked_lm(model_path, trust_remote_code): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS model = auto_spyre_model.AutoSpyreModelForMaskedLM.from_pretrained( model_path, dtype=auto_spyre_model.dtype_for_model_path( @@ -95,9 +98,10 @@ def test_load_masked_lm(model_path): @pytest.mark.parametrize( "model_path", QUESTION_ANSWERING_PATHS, ids=QUESTION_ANSWERING_PATHS ) -def test_load_question_answering(model_path): +def test_load_question_answering(model_path, trust_remote_code): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS model = auto_spyre_model.AutoSpyreModelForQuestionAnswering.from_pretrained( model_path, dtype=auto_spyre_model.dtype_for_model_path( @@ -115,10 +119,12 @@ def test_load_question_answering(model_path): @pytest.mark.parametrize( "model_path", TOKEN_CLASSIFICATION_PATHS, ids=TOKEN_CLASSIFICATION_PATHS ) -def test_load_token_classification(model_path): +def test_load_token_classification(model_path, trust_remote_code): auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS model = auto_spyre_model.AutoSpyreModelForTokenClassification.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) assert callable(model.forward) assert next(model.classifier.parameters()).device.type == "cpu" diff --git a/tests/cpu/test_masked_lm_cpu_accuracy.py b/tests/cpu/test_masked_lm_cpu_accuracy.py index 783004fa..fe419373 100644 --- a/tests/cpu/test_masked_lm_cpu_accuracy.py +++ b/tests/cpu/test_masked_lm_cpu_accuracy.py @@ -34,9 +34,10 @@ @pytest.mark.parametrize("model_path", MASKED_LM_PATHS, ids=MASKED_LM_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"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + 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 ) diff --git a/tests/cpu/test_question_answering_cpu_accuracy.py b/tests/cpu/test_question_answering_cpu_accuracy.py index acab68ac..8816dfea 100644 --- a/tests/cpu/test_question_answering_cpu_accuracy.py +++ b/tests/cpu/test_question_answering_cpu_accuracy.py @@ -37,9 +37,10 @@ @pytest.mark.parametrize( "model_path", QUESTION_ANSWERING_PATHS, ids=QUESTION_ANSWERING_PATHS ) -def test_native_forward(model_path: str) -> None: +def test_native_forward(model_path: str, trust_remote_code: bool | None) -> None: auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + 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 ) diff --git a/tests/cpu/test_reranker_cpu_accuracy.py b/tests/cpu/test_reranker_cpu_accuracy.py index cf62bff9..70d27d70 100644 --- a/tests/cpu/test_reranker_cpu_accuracy.py +++ b/tests/cpu/test_reranker_cpu_accuracy.py @@ -79,10 +79,10 @@ def _assert_reranker_logits( @pytest.mark.parametrize("model_path", RERANKER_PATHS, ids=RERANKER_PATHS) -def test_auto_loader(model_path: str) -> None: +def test_auto_loader(model_path: str, trust_remote_code: bool | None) -> None: """AutoSpyreModelForSequenceClassification logits match HF CPU reference.""" ref_logits, adapter_logits = run_seq_classification_auto_loader_vs_ref( - model_path, PAIRS + model_path, PAIRS, trust_remote_code=trust_remote_code ) gc.collect() _assert_reranker_logits(ref_logits, adapter_logits) diff --git a/tests/cpu/test_seq_classification_cpu_accuracy.py b/tests/cpu/test_seq_classification_cpu_accuracy.py index f555c341..b2c8fd7b 100644 --- a/tests/cpu/test_seq_classification_cpu_accuracy.py +++ b/tests/cpu/test_seq_classification_cpu_accuracy.py @@ -74,10 +74,10 @@ def _assert_seq_classification_logits( @pytest.mark.parametrize( "model_path", SEQ_CLASSIFICATION_PATHS, ids=SEQ_CLASSIFICATION_PATHS ) -def test_auto_loader(model_path: str) -> None: +def test_auto_loader(model_path: str, trust_remote_code: bool | None) -> None: """AutoSpyreModelForSequenceClassification logits match HF CPU reference.""" ref_logits, adapter_logits = run_seq_classification_auto_loader_vs_ref( - model_path, TEXTS + model_path, TEXTS, trust_remote_code=trust_remote_code ) gc.collect() _assert_seq_classification_logits(ref_logits, adapter_logits) diff --git a/tests/cpu/test_token_classification_cpu_accuracy.py b/tests/cpu/test_token_classification_cpu_accuracy.py index c6f0919f..227aa6ce 100644 --- a/tests/cpu/test_token_classification_cpu_accuracy.py +++ b/tests/cpu/test_token_classification_cpu_accuracy.py @@ -36,9 +36,10 @@ @pytest.mark.parametrize( "model_path", TOKEN_CLASSIFICATION_PATHS, ids=TOKEN_CLASSIFICATION_PATHS ) -def test_native_forward(model_path: str) -> None: +def test_native_forward(model_path: str, trust_remote_code: bool | None) -> None: auto_spyre_model = sys.modules["hf_adapters.auto_spyre_model"] - trust_remote_code = model_path in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS dtype = get_dtype_for_cpu(model_path) adapter_module = resolve_adapter_module_for_test( model_path, diff --git a/tests/cpu/test_vlm_e2e_cpu.py b/tests/cpu/test_vlm_e2e_cpu.py index 3c76b5c9..4b0d2359 100644 --- a/tests/cpu/test_vlm_e2e_cpu.py +++ b/tests/cpu/test_vlm_e2e_cpu.py @@ -57,10 +57,11 @@ @pytest.mark.parametrize("model_path", VISION_PATHS, ids=VISION_PATHS) -def test_vlm_generate(model_path: str) -> None: +def test_vlm_generate(model_path: str, trust_remote_code: bool | None) -> None: from hf_adapters.auto_spyre_model import dtype_for_model_path - trust_remote_code = model_path in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS adapter = resolve_adapter_module_for_test( model_path, mapping=IMAGE_TEXT_TO_TEXT_CONFIG_TO_ADAPTER_MODULE_MAPPING, diff --git a/tests/spyre/edge_cases/test_eos_spyre.py b/tests/spyre/edge_cases/test_eos_spyre.py index e3a9e65c..1343d786 100644 --- a/tests/spyre/edge_cases/test_eos_spyre.py +++ b/tests/spyre/edge_cases/test_eos_spyre.py @@ -34,22 +34,32 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_eos_first_of_second_block_spyre(model_path: str) -> None: - ok, detail = run_eos_case(model_path, "eos_first_of_second_block") +def test_eos_first_of_second_block_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_eos_case( + model_path, "eos_first_of_second_block", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_eos_first_token_spyre(model_path: str) -> None: - ok, detail = run_eos_case(model_path, "eos_first_token") +def test_eos_first_token_spyre(model_path: str, trust_remote_code: bool | None) -> None: + ok, detail = run_eos_case( + model_path, "eos_first_token", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_eos_inside_prompt_spyre(model_path: str) -> None: - info, tokenizer, ref_model, model = _setup(model_path, need_ref=True) +def test_eos_inside_prompt_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + info, tokenizer, ref_model, model = _setup( + model_path, need_ref=True, trust_remote_code=trust_remote_code + ) try: if tokenizer.eos_token_id is None: pytest.skip("tokenizer has no eos_token_id") @@ -81,22 +91,32 @@ def test_eos_inside_prompt_spyre(model_path: str) -> None: @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_eos_mid_block_spyre(model_path: str) -> None: - ok, detail = run_eos_case(model_path, "eos_mid_block") +def test_eos_mid_block_spyre(model_path: str, trust_remote_code: bool | None) -> None: + ok, detail = run_eos_case( + model_path, "eos_mid_block", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_eos_on_last_step_spyre(model_path: str) -> None: - ok, detail = run_eos_case(model_path, "eos_on_last_step") +def test_eos_on_last_step_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_eos_case( + model_path, "eos_on_last_step", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_no_eos_runs_full_budget_spyre(model_path: str) -> None: - info, tokenizer, ref_model, model = _setup(model_path, need_ref=True) +def test_no_eos_runs_full_budget_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + info, tokenizer, ref_model, model = _setup( + model_path, need_ref=True, trust_remote_code=trust_remote_code + ) try: no_eos_prompts = make_prompts(tokenizer, [5, 12]) no_eos_max_new = 64 + 7 diff --git a/tests/spyre/edge_cases/test_long_multi_block_spyre.py b/tests/spyre/edge_cases/test_long_multi_block_spyre.py index d30a9ae3..12315d95 100644 --- a/tests/spyre/edge_cases/test_long_multi_block_spyre.py +++ b/tests/spyre/edge_cases/test_long_multi_block_spyre.py @@ -24,6 +24,10 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_long_multi_block_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "long_multi_block") +def test_long_multi_block_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_greedy_case( + model_path, "long_multi_block", trust_remote_code=trust_remote_code + ) assert ok, detail diff --git a/tests/spyre/edge_cases/test_mixed_spyre.py b/tests/spyre/edge_cases/test_mixed_spyre.py index 3769fe1e..ef94848a 100644 --- a/tests/spyre/edge_cases/test_mixed_spyre.py +++ b/tests/spyre/edge_cases/test_mixed_spyre.py @@ -24,13 +24,19 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_mixed_short_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "mixed_short") +def test_mixed_short_spyre(model_path: str, trust_remote_code: bool | None) -> None: + ok, detail = run_greedy_case( + model_path, "mixed_short", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_mixed_with_single_token_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "mixed_with_single_token") +def test_mixed_with_single_token_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_greedy_case( + model_path, "mixed_with_single_token", trust_remote_code=trust_remote_code + ) assert ok, detail diff --git a/tests/spyre/edge_cases/test_no_pad_token_fallback_spyre.py b/tests/spyre/edge_cases/test_no_pad_token_fallback_spyre.py index 02d60813..934924c5 100644 --- a/tests/spyre/edge_cases/test_no_pad_token_fallback_spyre.py +++ b/tests/spyre/edge_cases/test_no_pad_token_fallback_spyre.py @@ -34,8 +34,12 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_no_pad_token_fallback_spyre(model_path: str) -> None: - info, tokenizer, ref_model, model = _setup(model_path, need_ref=True) +def test_no_pad_token_fallback_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + info, tokenizer, ref_model, model = _setup( + model_path, need_ref=True, trust_remote_code=trust_remote_code + ) try: no_pad_prompts = make_prompts(tokenizer, [5, 12]) no_pad_max_new = 16 diff --git a/tests/spyre/edge_cases/test_prompt_exactly_block_spyre.py b/tests/spyre/edge_cases/test_prompt_exactly_block_spyre.py index a8cdcf4c..8f392ca3 100644 --- a/tests/spyre/edge_cases/test_prompt_exactly_block_spyre.py +++ b/tests/spyre/edge_cases/test_prompt_exactly_block_spyre.py @@ -24,6 +24,10 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_prompt_exactly_block_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "prompt_exactly_block") +def test_prompt_exactly_block_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_greedy_case( + model_path, "prompt_exactly_block", trust_remote_code=trust_remote_code + ) assert ok, detail diff --git a/tests/spyre/edge_cases/test_sampling_spyre.py b/tests/spyre/edge_cases/test_sampling_spyre.py index 6af04222..ef7865dd 100644 --- a/tests/spyre/edge_cases/test_sampling_spyre.py +++ b/tests/spyre/edge_cases/test_sampling_spyre.py @@ -36,8 +36,12 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_sampling_determinism_spyre(model_path: str) -> None: - info, tokenizer, _, model = _setup(model_path, need_ref=False) +def test_sampling_determinism_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + info, tokenizer, _, model = _setup( + model_path, need_ref=False, trust_remote_code=trust_remote_code + ) try: sampling_prompts = make_prompts(tokenizer, SAMPLING_TARGETS) encoded = encode_generation_inputs(tokenizer, sampling_prompts) @@ -71,8 +75,12 @@ def test_sampling_determinism_spyre(model_path: str) -> None: @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_sampling_top_k_zero_spyre(model_path: str) -> None: - info, tokenizer, _, model = _setup(model_path, need_ref=False) +def test_sampling_top_k_zero_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + info, tokenizer, _, model = _setup( + model_path, need_ref=False, trust_remote_code=trust_remote_code + ) try: sampling_prompts = make_prompts(tokenizer, SAMPLING_TARGETS) encoded = encode_generation_inputs(tokenizer, sampling_prompts) diff --git a/tests/spyre/edge_cases/test_short_spyre.py b/tests/spyre/edge_cases/test_short_spyre.py index 1e2770a6..16e69e1a 100644 --- a/tests/spyre/edge_cases/test_short_spyre.py +++ b/tests/spyre/edge_cases/test_short_spyre.py @@ -24,27 +24,41 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_short_block_minus_one_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "short_block_minus_one") +def test_short_block_minus_one_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_greedy_case( + model_path, "short_block_minus_one", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_short_cross_block_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "short_cross_block") +def test_short_cross_block_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_greedy_case( + model_path, "short_cross_block", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_short_one_token_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "short_one_token") +def test_short_one_token_spyre(model_path: str, trust_remote_code: bool | None) -> None: + ok, detail = run_greedy_case( + model_path, "short_one_token", trust_remote_code=trust_remote_code + ) assert ok, detail @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_short_two_blocks_plus_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "short_two_blocks_plus") +def test_short_two_blocks_plus_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_greedy_case( + model_path, "short_two_blocks_plus", trust_remote_code=trust_remote_code + ) assert ok, detail diff --git a/tests/spyre/edge_cases/test_single_token_prompt_spyre.py b/tests/spyre/edge_cases/test_single_token_prompt_spyre.py index c6052a15..f91aa86e 100644 --- a/tests/spyre/edge_cases/test_single_token_prompt_spyre.py +++ b/tests/spyre/edge_cases/test_single_token_prompt_spyre.py @@ -24,6 +24,10 @@ @pytest.mark.parametrize("model_path", CAUSAL_PATHS, ids=CAUSAL_PATHS) @pytest.mark.slow -def test_single_token_prompt_spyre(model_path: str) -> None: - ok, detail = run_greedy_case(model_path, "single_token_prompt") +def test_single_token_prompt_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + ok, detail = run_greedy_case( + model_path, "single_token_prompt", trust_remote_code=trust_remote_code + ) assert ok, detail diff --git a/tests/spyre/test_dspark_draft_spyre.py b/tests/spyre/test_dspark_draft_spyre.py index 0e1db2eb..1813ec30 100644 --- a/tests/spyre/test_dspark_draft_spyre.py +++ b/tests/spyre/test_dspark_draft_spyre.py @@ -40,7 +40,7 @@ @pytest.mark.parametrize("ckpt", DSPARK_PATHS) -def test_dspark_draft_block(ckpt): +def test_dspark_draft_block(ckpt, trust_remote_code): """prepare_for_spyre + _run_draft_block produce finite block hidden states.""" import torch_spyre # noqa: F401 from transformers import AutoConfig, AutoModelForCausalLM @@ -52,7 +52,8 @@ def test_dspark_draft_block(ckpt): dev = torch.device("spyre:0") hf_common.DEVICE = dev - trust_remote_code = ckpt in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = ckpt in REMOTE_CODE_PATHS # The library resolves the checkpoint to its adapter by architecture # (``*DSparkModel``); confirm it lands on a DSpark draft adapter module. diff --git a/tests/spyre/test_e2e_masked_lm_compare_spyre.py b/tests/spyre/test_e2e_masked_lm_compare_spyre.py index eba3ab34..9727e63d 100644 --- a/tests/spyre/test_e2e_masked_lm_compare_spyre.py +++ b/tests/spyre/test_e2e_masked_lm_compare_spyre.py @@ -29,8 +29,11 @@ @pytest.mark.parametrize("model_path", MASKED_LM_PATHS, ids=MASKED_LM_PATHS) -def test_e2e_masked_lm_compare_spyre(model_path: str) -> None: - trust_remote_code = model_path in REMOTE_CODE_PATHS +def test_e2e_masked_lm_compare_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + 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 ) diff --git a/tests/spyre/test_e2e_multibatch_spyre.py b/tests/spyre/test_e2e_multibatch_spyre.py index 4fbbcdc2..d1c5724c 100644 --- a/tests/spyre/test_e2e_multibatch_spyre.py +++ b/tests/spyre/test_e2e_multibatch_spyre.py @@ -96,8 +96,9 @@ def _print_table(model_path: str, rows: list[dict[str, Any]]) -> None: @pytest.mark.parametrize( "model_path", xfail_non_blocking(CAUSAL_PATHS, table=NON_BLOCKING_CAUSAL_MODELS) ) -def test_e2e_multibatch_spyre(model_path: str) -> None: - trust_remote_code = model_path in REMOTE_CODE_PATHS +def test_e2e_multibatch_spyre(model_path: str, trust_remote_code: bool | None) -> None: + 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 ) diff --git a/tests/spyre/test_e2e_question_answering_compare_spyre.py b/tests/spyre/test_e2e_question_answering_compare_spyre.py index 101a9a20..01271e7a 100644 --- a/tests/spyre/test_e2e_question_answering_compare_spyre.py +++ b/tests/spyre/test_e2e_question_answering_compare_spyre.py @@ -32,8 +32,11 @@ @pytest.mark.parametrize( "model_path", QUESTION_ANSWERING_PATHS, ids=QUESTION_ANSWERING_PATHS ) -def test_e2e_question_answering_compare_spyre(model_path: str) -> None: - trust_remote_code = model_path in REMOTE_CODE_PATHS +def test_e2e_question_answering_compare_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + 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 ) diff --git a/tests/spyre/test_e2e_reranker_compare_spyre.py b/tests/spyre/test_e2e_reranker_compare_spyre.py index fd9c6334..79617454 100644 --- a/tests/spyre/test_e2e_reranker_compare_spyre.py +++ b/tests/spyre/test_e2e_reranker_compare_spyre.py @@ -40,7 +40,7 @@ SEQUENCE_CLASSIFICATION_CONFIG_TO_ADAPTER_MODULE_MAPPING, resolve_adapter_module, ) -from tests.model_registry import RERANKER_PATHS +from tests.model_registry import REMOTE_CODE_PATHS, RERANKER_PATHS from tests.spyre._seq_classification_helpers import run_seq_classification_cpu_vs_spyre pytestmark = pytest.mark.model_harness("reranker") @@ -65,12 +65,19 @@ @pytest.mark.parametrize("model_path", RERANKER_PATHS, ids=RERANKER_PATHS) -def test_e2e_reranker_compare_spyre(model_path: str) -> None: +def test_e2e_reranker_compare_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS adapter = resolve_adapter_module( model_path, mapping=SEQUENCE_CLASSIFICATION_CONFIG_TO_ADAPTER_MODULE_MAPPING, + trust_remote_code=trust_remote_code, + ) + result = run_seq_classification_cpu_vs_spyre( + model_path, adapter, PAIRS, trust_remote_code=trust_remote_code ) - result = run_seq_classification_cpu_vs_spyre(model_path, adapter, PAIRS) ref_scores = result["ref_logits"][:, 0] spyre_scores = result["spyre_logits"][:, 0] diff --git a/tests/spyre/test_e2e_seq_classification_compare_spyre.py b/tests/spyre/test_e2e_seq_classification_compare_spyre.py index 6a252136..361c82e4 100644 --- a/tests/spyre/test_e2e_seq_classification_compare_spyre.py +++ b/tests/spyre/test_e2e_seq_classification_compare_spyre.py @@ -40,7 +40,7 @@ import torch import torch.nn.functional as F from _seq_classification_helpers import run_seq_classification_cpu_vs_spyre -from model_registry import SEQ_CLASSIFICATION_PATHS +from model_registry import REMOTE_CODE_PATHS, SEQ_CLASSIFICATION_PATHS from hf_adapters.auto_spyre_model import ( SEQUENCE_CLASSIFICATION_CONFIG_TO_ADAPTER_MODULE_MAPPING, @@ -62,12 +62,19 @@ @pytest.mark.parametrize( "model_path", SEQ_CLASSIFICATION_PATHS, ids=SEQ_CLASSIFICATION_PATHS ) -def test_e2e_seq_classification_compare_spyre(model_path: str) -> None: +def test_e2e_seq_classification_compare_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS adapter = resolve_adapter_module( model_path, mapping=SEQUENCE_CLASSIFICATION_CONFIG_TO_ADAPTER_MODULE_MAPPING, + trust_remote_code=trust_remote_code, + ) + result = run_seq_classification_cpu_vs_spyre( + model_path, adapter, TEXTS, trust_remote_code=trust_remote_code ) - result = run_seq_classification_cpu_vs_spyre(model_path, adapter, TEXTS) ref_logits = result["ref_logits"] spyre_logits = result["spyre_logits"] diff --git a/tests/spyre/test_e2e_smoke_spyre.py b/tests/spyre/test_e2e_smoke_spyre.py index 491dcf73..d4072da1 100644 --- a/tests/spyre/test_e2e_smoke_spyre.py +++ b/tests/spyre/test_e2e_smoke_spyre.py @@ -42,25 +42,30 @@ pytestmark = pytest.mark.model_harness("causal") -def run_smoke_test(model_path: str) -> dict[str, Any]: +def run_smoke_test( + model_path: str, trust_remote_code: bool | None = None +) -> dict[str, Any]: """Load model, generate 5 tokens, validate output. Returns a result dict.""" from transformers import AutoTokenizer from hf_adapters import AutoSpyreModelForCausalLM + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS + print(f"\n{'=' * 70}") print(f" loading from {model_path}") print(f"{'=' * 70}") t0 = time.time() model = AutoSpyreModelForCausalLM.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) load_time = time.time() - t0 print(f" Load time: {load_time:.1f}s") tokenizer = AutoTokenizer.from_pretrained( - model_path, trust_remote_code=(model_path in REMOTE_CODE_PATHS) + model_path, trust_remote_code=trust_remote_code ) prompt = "The capital of France is" print(f" Prompt: {prompt!r}") @@ -112,8 +117,8 @@ def run_smoke_test(model_path: str) -> dict[str, Any]: @pytest.mark.parametrize( "model_path", xfail_non_blocking(CAUSAL_PATHS, table=NON_BLOCKING_CAUSAL_MODELS) ) -def test_e2e_smoke_spyre(model_path: str) -> None: - result = run_smoke_test(model_path) +def test_e2e_smoke_spyre(model_path: str, trust_remote_code: bool | None) -> None: + result = run_smoke_test(model_path, trust_remote_code=trust_remote_code) print("\n## E2E Smoke Test Results\n") print("| Model | Status | Tokens | Generated Text | Load (s) | Gen (s) |") print("|-------|--------|--------|----------------|----------|---------|") diff --git a/tests/spyre/test_e2e_smoke_vision_spyre.py b/tests/spyre/test_e2e_smoke_vision_spyre.py index b9dfc7ac..911204c6 100644 --- a/tests/spyre/test_e2e_smoke_vision_spyre.py +++ b/tests/spyre/test_e2e_smoke_vision_spyre.py @@ -130,7 +130,9 @@ def _run_single_image( } -def run_vision_smoke_test(model_path: str) -> dict[str, Any]: +def run_vision_smoke_test( + model_path: str, trust_remote_code: bool | None = None +) -> dict[str, Any]: """Load model once, then run all SMOKE_TEST_IMAGES through it in sequence. Returns a result dict with per-image results and overall load time. @@ -140,7 +142,8 @@ def run_vision_smoke_test(model_path: str) -> dict[str, Any]: print(f" loading from {model_path}") print(f"{'=' * 70}") - trust_remote_code = model_path in REMOTE_CODE_PATHS + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS dtype = dtype_for_model_path( model_path, target_device="spyre", trust_remote_code=trust_remote_code ) @@ -184,8 +187,10 @@ def run_vision_smoke_test(model_path: str) -> dict[str, Any]: @pytest.mark.parametrize( "model_path", xfail_non_blocking(VISION_PATHS, table=NON_BLOCKING_VISION_MODELS) ) -def test_e2e_smoke_vision_spyre(model_path: str) -> None: - result = run_vision_smoke_test(model_path) +def test_e2e_smoke_vision_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + result = run_vision_smoke_test(model_path, trust_remote_code=trust_remote_code) print("\n## E2E Vision Smoke Test Results\n") print("| Image | Status | Tokens | Generated Text | Gen (s) |") diff --git a/tests/spyre/test_e2e_token_classification_compare_spyre.py b/tests/spyre/test_e2e_token_classification_compare_spyre.py index b6097c9c..f2b40f10 100644 --- a/tests/spyre/test_e2e_token_classification_compare_spyre.py +++ b/tests/spyre/test_e2e_token_classification_compare_spyre.py @@ -54,8 +54,11 @@ @pytest.mark.parametrize( "model_path", TOKEN_CLASSIFICATION_PATHS, ids=TOKEN_CLASSIFICATION_PATHS ) -def test_e2e_token_classification_compare_spyre(model_path: str) -> None: - trust_remote_code = model_path in REMOTE_CODE_PATHS +def test_e2e_token_classification_compare_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + 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 ) diff --git a/tests/spyre/test_e2e_token_compare_spyre.py b/tests/spyre/test_e2e_token_compare_spyre.py index 68127c6a..a234c6b5 100644 --- a/tests/spyre/test_e2e_token_compare_spyre.py +++ b/tests/spyre/test_e2e_token_compare_spyre.py @@ -327,8 +327,9 @@ def _run_model_test( def token_compare_spyre( model_path: str, + trust_remote_code: bool | None = None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - rows = _run_model_test(model_path) + rows = _run_model_test(model_path, trust_remote_code=trust_remote_code) mismatches = [r for r in rows if not r["top1_match"]] return mismatches, rows @@ -336,8 +337,12 @@ def token_compare_spyre( @pytest.mark.parametrize( "model_path", xfail_non_blocking(CAUSAL_PATHS, table=NON_BLOCKING_CAUSAL_MODELS) ) -def test_e2e_token_compare_spyre(model_path: str) -> None: - mismatches, rows = token_compare_spyre(model_path) +def test_e2e_token_compare_spyre( + model_path: str, trust_remote_code: bool | None +) -> None: + mismatches, rows = token_compare_spyre( + model_path, trust_remote_code=trust_remote_code + ) _print_table(rows) n_match = sum(1 for r in rows if r["top1_match"]) print(f"\nTop-1 agreement: {n_match}/{len(rows)} steps") diff --git a/tests/spyre/test_load_spyre.py b/tests/spyre/test_load_spyre.py index 19e4ad20..70a8ca2a 100644 --- a/tests/spyre/test_load_spyre.py +++ b/tests/spyre/test_load_spyre.py @@ -55,9 +55,11 @@ @pytest.mark.parametrize( "model_path", xfail_non_blocking(CAUSAL_PATHS, table=NON_BLOCKING_CAUSAL_MODELS) ) -def test_load_causal_lm(model_path: str) -> None: +def test_load_causal_lm(model_path: str, trust_remote_code: bool | None) -> None: - model_is_not_none, callables, load_s = load_causal_lm(model_path) + model_is_not_none, callables, load_s = load_causal_lm( + model_path, trust_remote_code=trust_remote_code + ) print(f" [{model_path}] causal-LM load time: {load_s:.1f}s") print("\n## Spyre Load Test Results\n") @@ -106,9 +108,11 @@ def load_embedding( @pytest.mark.model_harness("embedding") @pytest.mark.parametrize("model_path", EMBED_PATHS, ids=EMBED_PATHS) -def test_load_embedding(model_path: str) -> None: +def test_load_embedding(model_path: str, trust_remote_code: bool | None) -> None: - model_loaded, load_s = load_embedding(model_path) + model_loaded, load_s = load_embedding( + model_path, trust_remote_code=trust_remote_code + ) assert model_loaded, f"{model_path}: from_pretrained returned None" print(f" [{model_path}] embedding load time: {load_s:.1f}s") print("\n## Spyre Load Test Results\n") @@ -142,9 +146,11 @@ def load_masked_lm( @pytest.mark.model_harness("masked_lm") @pytest.mark.parametrize("model_path", MASKED_LM_PATHS, ids=MASKED_LM_PATHS) -def test_load_masked_lm(model_path: str) -> None: +def test_load_masked_lm(model_path: str, trust_remote_code: bool | None) -> None: - model_is_not_none, callables, load_s = load_masked_lm(model_path) + model_is_not_none, callables, load_s = load_masked_lm( + model_path, trust_remote_code=trust_remote_code + ) print(f" [{model_path}] masked-LM load time: {load_s:.1f}s") print("\n## Spyre Load Test Results\n") @@ -179,8 +185,12 @@ def load_question_answering( @pytest.mark.parametrize( "model_path", QUESTION_ANSWERING_PATHS, ids=QUESTION_ANSWERING_PATHS ) -def test_load_question_answering(model_path: str) -> None: - model_is_not_none, ready, load_s = load_question_answering(model_path) +def test_load_question_answering( + model_path: str, trust_remote_code: bool | None +) -> None: + model_is_not_none, ready, load_s = load_question_answering( + model_path, trust_remote_code=trust_remote_code + ) print(f" [{model_path}] question-answering load time: {load_s:.1f}s") assert model_is_not_none, f"{model_path}: from_pretrained returned None" assert ready, f"{model_path}: native forward or CPU QA head is not ready" @@ -206,8 +216,12 @@ def load_seq_classification( @pytest.mark.parametrize( "model_path", SEQ_CLASSIFICATION_PATHS, ids=SEQ_CLASSIFICATION_PATHS ) -def test_load_seq_classification(model_path: str) -> None: - model_is_not_none, ready, load_s = load_seq_classification(model_path) +def test_load_seq_classification( + model_path: str, trust_remote_code: bool | None +) -> None: + model_is_not_none, ready, load_s = load_seq_classification( + model_path, trust_remote_code=trust_remote_code + ) print(f" [{model_path}] seq-classification load time: {load_s:.1f}s") print("\n## Spyre Load Test Results\n") print("| Path | Kind | Status | Load (s) |") @@ -237,8 +251,12 @@ def load_token_classification( @pytest.mark.parametrize( "model_path", TOKEN_CLASSIFICATION_PATHS, ids=TOKEN_CLASSIFICATION_PATHS ) -def test_load_token_classification(model_path: str) -> None: - model_is_not_none, ready, load_s = load_token_classification(model_path) +def test_load_token_classification( + model_path: str, trust_remote_code: bool | None +) -> None: + model_is_not_none, ready, load_s = load_token_classification( + model_path, trust_remote_code=trust_remote_code + ) print(f" [{model_path}] token-classification load time: {load_s:.1f}s") print("\n## Spyre Load Test Results\n") print("| Path | Kind | Status | Load (s) |") diff --git a/tests/spyre/test_multicard_spyre.py b/tests/spyre/test_multicard_spyre.py index 7782149c..793b471e 100644 --- a/tests/spyre/test_multicard_spyre.py +++ b/tests/spyre/test_multicard_spyre.py @@ -386,9 +386,11 @@ def _run_generate() -> tuple[list[str], str]: @pytest.mark.parametrize("model_path", [_DEFAULT_MODEL]) -def test_multicard_smoke_single_card(model_path: str) -> None: +def test_multicard_smoke_single_card( + model_path: str, trust_remote_code: bool | None +) -> None: """Single-card smoke test: load, generate, verify output passes all checks.""" - result = run_multicard_smoke_test(model_path) + result = run_multicard_smoke_test(model_path, trust_remote_code=trust_remote_code) assert result["status"] == "PASS", ( f"Smoke test failed with status {result['status']}.\n" f"Checks: {result.get('seq_checks')}\n" diff --git a/tests/spyre/test_vlm_e2e_spyre.py b/tests/spyre/test_vlm_e2e_spyre.py index b3d6ad93..aed835cf 100644 --- a/tests/spyre/test_vlm_e2e_spyre.py +++ b/tests/spyre/test_vlm_e2e_spyre.py @@ -264,8 +264,9 @@ def _stock_vlm_greedy_steps( @pytest.mark.parametrize( "model_path", xfail_non_blocking(VISION_PATHS, table=NON_BLOCKING_VISION_MODELS) ) -def test_vlm_generate_spyre(model_path: str) -> None: - trust_remote_code = model_path in REMOTE_CODE_PATHS +def test_vlm_generate_spyre(model_path: str, trust_remote_code: bool | None) -> None: + if trust_remote_code is None: + trust_remote_code = model_path in REMOTE_CODE_PATHS adapter = resolve_adapter_module( model_path, mapping=IMAGE_TEXT_TO_TEXT_CONFIG_TO_ADAPTER_MODULE_MAPPING,