fix(inference): pass the token through when loading the encoder (#362)

* fix(inference): pass the token to the encoder download

* test: cover token passthrough on both pipeline load paths

Author: serhiizghama

model2vec/inference/model.py

@@ -221,7 +221,7 @@ def _load_pipeline(folder_or_repo_path: PathLike, token: str | None = None) -> t
         except EntryNotFoundError:
             return _load_legacy_pipeline(folder_or_repo_path, token)
 
-    model = StaticModel.from_pretrained(folder_or_repo_path)
+    model = StaticModel.from_pretrained(folder_or_repo_path, token=token)
 
     head_config = model.config.get("head_config", {})
     activation = Activation(head_config.get("activation", Activation.IDENTITY.value))
@@ -291,7 +291,7 @@ def convert_legacy_pipeline(
             folder_or_repo_path.as_posix(), _LEGACY_HEAD_FILENAME, token=token
         )
 
-    model = StaticModel.from_pretrained(folder_or_repo_path)
+    model = StaticModel.from_pretrained(folder_or_repo_path, token=token)
     model.embedding = np.nan_to_num(model.embedding)
 
     untrusted_types = skops.io.get_untrusted_types(file=legacy_head_path)

tests/test_inference.py

@@ -135,6 +135,24 @@ def test_load_pipeline_from_hub(mock_inference_pipeline: StaticModelPipeline) ->
         assert loaded.predict(["dog"]).tolist() == mock_inference_pipeline.predict(["dog"]).tolist()
 
 
+def test_load_pipeline_from_hub_with_token(mock_inference_pipeline: StaticModelPipeline) -> None:
+    """Test that the token also reaches the encoder download, not just the head download."""
+    with TemporaryDirectory() as temp_dir:
+        mock_inference_pipeline.save_pretrained(temp_dir)
+        downloaded_model = StaticModel.from_pretrained(temp_dir)
+        head_path = os.path.join(temp_dir, "head.safetensors")
+
+        with (
+            patch("model2vec.inference.model.huggingface_hub.hf_hub_download", return_value=head_path),
+            patch(
+                "model2vec.inference.model.StaticModel.from_pretrained", return_value=downloaded_model
+            ) as mock_from_pretrained,
+        ):
+            StaticModelPipeline.from_pretrained("fake/repo-id", token="secret")
+
+        assert mock_from_pretrained.call_args.kwargs.get("token") == "secret"
+
+
 def test_push_to_hub(mock_inference_pipeline: StaticModelPipeline) -> None:
     """Test that push_to_hub saves the pipeline to a temp folder before pushing it to the hub."""
     captured: dict[str, object] = {}
@@ -278,6 +296,31 @@ def _fake_download(repo_id: str, filename: str, token: str | None = None) -> str
     assert np.allclose(loaded.predict(["dog", "cat"]), legacy_pipeline.predict(encoded))
 
 
+def test_convert_legacy_pipeline_with_token(mock_static_model: StaticModel) -> None:
+    """Test that the token also reaches the encoder download on the legacy path."""
+    rng = np.random.RandomState(0)
+    X = rng.randn(30, mock_static_model.dim)
+    y = rng.randn(30, 4)
+    legacy_pipeline = make_pipeline(MLPRegressor(hidden_layer_sizes=(8,), max_iter=1, random_state=0).fit(X, y))
+
+    with TemporaryDirectory() as temp_dir:
+        _dump_legacy_pipeline(temp_dir, mock_static_model, legacy_pipeline)
+        downloaded_model = StaticModel.from_pretrained(temp_dir)
+
+        with (
+            patch(
+                "model2vec.inference.model.huggingface_hub.hf_hub_download",
+                return_value=os.path.join(temp_dir, "pipeline.skops"),
+            ),
+            patch(
+                "model2vec.inference.model.StaticModel.from_pretrained", return_value=downloaded_model
+            ) as mock_from_pretrained,
+        ):
+            convert_legacy_pipeline("fake/repo-id", token="secret")
+
+    assert mock_from_pretrained.call_args.kwargs.get("token") == "secret"
+
+
 def test_convert_legacy_pipeline_untrusted_type(mock_static_model: StaticModel) -> None:
     """Test that an untrusted type in the legacy pipeline is rejected by default."""
     rng = np.random.RandomState(0)