fix readme, subfolder, add test

Author: stephantul

model2vec/persistence/persistence.py

@@ -101,12 +101,13 @@ def load_pretrained(
     folder_or_repo_path = Path(folder_or_repo_path)
 
     # We resolve a folder or repo path to an actual local folder.
-    folder = _resolve_folder(
-        folder_or_repo_path=folder_or_repo_path, token=token, force_download=force_download, subfolder=subfolder
-    )
+    folder = _resolve_folder(folder_or_repo_path=folder_or_repo_path, token=token, force_download=force_download)
+
+    if subfolder:
+        folder = folder / subfolder
 
     selected_layout = _get_paths(folder)
-    readme_path = folder_or_repo_path / "README.md"
+    readme_path = folder / "README.md"
 
     opened_tensor_file = cast(SafeOpenProtocol, safetensors.safe_open(selected_layout.embeddings, framework="numpy"))
     embedding_name = "embedding.weight" if selected_layout.is_sentence_transformers else "embeddings"
@@ -133,7 +134,7 @@ def load_pretrained(
     return embeddings, tokenizer, config, metadata, weights, mapping
 
 
-def _resolve_folder(folder_or_repo_path: Path, subfolder: str | None, token: str | None, force_download: bool) -> Path:
+def _resolve_folder(folder_or_repo_path: Path, token: str | None, force_download: bool) -> Path:
     """Resolve a folder locally or from hugging face hub."""
     if folder_or_repo_path.exists():
         return folder_or_repo_path

tests/test_persistence.py

@@ -17,12 +17,12 @@ def test_local_loading(mock_static_model: StaticModel) -> None:
         with patch(
             "model2vec.persistence.persistence.huggingface_hub.snapshot_download",
         ) as mock_snapshot:
-            mock_snapshot.return_value = dir_name
+            mock_snapshot.return_value = Path(dir_name)
             # Patch the cache to return the local path, simulate a cache hit
             # we pass a fake model name to actually hit the cache and snapshot download paths
             with patch("model2vec.persistence.persistence.maybe_get_cached_model_path") as cache:
                 # Simulate cache hit
-                cache.return_value = dir_name
+                cache.return_value = Path(dir_name)
                 s = StaticModel.from_pretrained("haha", force_download=True)
                 assert s.tokens == mock_static_model.tokens
                 s = StaticModel.from_pretrained("haha", force_download=False)
@@ -52,6 +52,17 @@ def test_garbage(mock_static_model: StaticModel) -> None:
             StaticModel.from_pretrained(dir_name)
 
 
+def test_subfolder_loading(mock_static_model: StaticModel) -> None:
+    """Test that garbage loading crashes."""
+    with TemporaryDirectory() as dir_name:
+        dir_name_path = Path(dir_name)
+        mock_static_model.save_pretrained(dir_name_path / "subfolder")
+        with pytest.raises(ValueError):
+            StaticModel.from_pretrained(dir_name)
+        model = StaticModel.from_pretrained(dir_name, subfolder="subfolder")
+        assert isinstance(model, StaticModel)
+
+
 def test_maybe_get_cached_model_path() -> None:
     """Test cached model path."""
     model_id = "t/t"