slimmify code, add tests

Author: stephantul

model2vec/persistence/hf.py

@@ -35,6 +35,9 @@ def maybe_get_cached_model_path(model_id: str) -> Path | None:
 
     Returns None if there is no cached model. In this case, the model will be downloaded.
     """
+    # Early exit: a model_id must have a single '/'
+    if model_id.count("/") != 1:
+        return None
     # Make path object
     cache_dir = Path(HF_HUB_CACHE)
     # This is specific to how HF stores the files.

model2vec/persistence/persistence.py

@@ -98,20 +98,14 @@ def load_pretrained(
     :return: The embeddings, tokenizer, config, metadata, weights and mapping.
 
     """
-    # Logic: by setting cached_folder, we no longer have a valid hub identifier.
-    # So, in _get_paths, we always hit a local folder, and never pull from the hub.
-    # If force_download is False, we never set the cached_folder.
-    if not force_download:
-        cached_folder = maybe_get_cached_model_path(str(folder_or_repo_path))
-        if cached_folder:
-            logger.info(f"Found cached model at {cached_folder}, loading from cache.")
-            folder_or_repo_path = cached_folder
-        else:
-            logger.info(f"No cached model found for {folder_or_repo_path}, loading from hub.")
-
     folder_or_repo_path = Path(folder_or_repo_path)
 
-    selected_layout = _get_paths(folder_or_repo_path, subfolder, token)
+    # 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
+    )
+
+    selected_layout = _get_paths(folder)
     readme_path = folder_or_repo_path / "README.md"
 
     opened_tensor_file = cast(SafeOpenProtocol, safetensors.safe_open(selected_layout.embeddings, framework="numpy"))
@@ -139,19 +133,23 @@ def load_pretrained(
     return embeddings, tokenizer, config, metadata, weights, mapping
 
 
-def _resolve_folder(folder_or_repo_path: Path, subfolder: str | None, token: str | None) -> Path:
+def _resolve_folder(folder_or_repo_path: Path, subfolder: str | None, token: str | None, force_download: bool) -> Path:
     """Resolve a folder locally or from hugging face hub."""
-    local_candidate = folder_or_repo_path / subfolder if subfolder else folder_or_repo_path
-    if local_candidate.exists():
-        folder = local_candidate
-    else:
-        folder = Path(huggingface_hub.snapshot_download(str(folder_or_repo_path), repo_type="model", token=token))
-    return folder / subfolder if subfolder else folder
+    if folder_or_repo_path.exists():
+        return folder_or_repo_path
+    # We now know we're dealing with either an invalid path, or
+    # a HF model ID.
+    if not force_download:
+        if folder := maybe_get_cached_model_path(str(folder_or_repo_path)):
+            return folder
+
+    folder = Path(huggingface_hub.snapshot_download(str(folder_or_repo_path), repo_type="model", token=token))
 
+    return folder
 
-def _get_paths(folder_or_repo_path: Path, subfolder: str | None, token: str | None) -> Layout:
-    folder = _resolve_folder(folder_or_repo_path, subfolder, token)
 
+def _get_paths(folder: Path) -> Layout:
+    """Get all paths by trying out multiple layouts."""
     for layout in FOLDER_LAYOUTS:
         layout = layout.with_parent(folder)
         if layout.is_valid():

tests/conftest.py

@@ -14,6 +14,7 @@
 from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 
 from model2vec.inference import StaticModelPipeline
+from model2vec.model import StaticModel
 from model2vec.train import StaticModelForClassification
 
 _TOKENIZER_TYPES = ["wordpiece", "bpe", "unigram"]
@@ -112,6 +113,17 @@ def mock_inference_pipeline(mock_trained_pipeline: StaticModelForClassification)
     return mock_trained_pipeline.to_pipeline()
 
 
+@pytest.fixture(scope="session")
+def mock_static_model() -> StaticModel:
+    """A small static model to test saving/loading."""
+    tokenizer = AutoTokenizer.from_pretrained("tests/data/test_tokenizer").backend_tokenizer
+    generator = np.random.RandomState(42)
+    vectors = generator.randn(len(tokenizer.get_vocab()), 12)
+    model = StaticModel(vectors=vectors, tokenizer=tokenizer)
+
+    return model
+
+
 @pytest.fixture(
     params=[
         (False, "single_label", "str"),

tests/test_persistence.py

@@ -0,0 +1,51 @@
+import shutil
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest.mock import patch
+
+import pytest
+
+from model2vec.model import StaticModel
+
+
+def test_local_loading(mock_static_model: StaticModel) -> None:
+    """Test saving and loading."""
+    with TemporaryDirectory() as dir_name:
+        mock_static_model.save_pretrained(dir_name)
+
+        with patch(
+            "model2vec.persistence.persistence.huggingface_hub.snapshot_download",
+        ) as mock_snapshot:
+            mock_snapshot.return_value = 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
+                s = StaticModel.from_pretrained("haha", force_download=True)
+                assert s.tokens == mock_static_model.tokens
+                s = StaticModel.from_pretrained("haha", force_download=False)
+                assert s.tokens == mock_static_model.tokens
+
+                # Simulate cache miss
+                cache.return_value = None
+                s = StaticModel.from_pretrained("haha", force_download=True)
+                assert s.tokens == mock_static_model.tokens
+                s = StaticModel.from_pretrained("haha", force_download=False)
+                assert s.tokens == mock_static_model.tokens
+
+                # Called twice, only when `force_download` is False
+                assert cache.call_count == 2
+            # Called three times, two times when force download is True,
+            # and once when it was false but the cache returned None
+            assert mock_snapshot.call_count == 3
+
+
+def test_garbage(mock_static_model: StaticModel) -> None:
+    """Test that garbage loading crashes."""
+    with TemporaryDirectory() as dir_name:
+        mock_static_model.save_pretrained(dir_name)
+        dir_name_path = Path(dir_name)
+        shutil.move(dir_name_path / "model.safetensors", dir_name_path / "model.safetenso")
+        with pytest.raises(ValueError):
+            StaticModel.from_pretrained(dir_name)