fix: fix loading of local models

Author: lsorber

src/wtpsplit_lite/_config.py

@@ -38,10 +38,11 @@ def num_labels(self) -> int:
 
     @classmethod
     @cache
-    def from_pretrained(cls, pretrained_model_name_or_path: str) -> "SubwordXLMConfig":
-        is_local = Path.is_file(Path(pretrained_model_name_or_path))
+    def from_pretrained(cls, pretrained_model_name_or_path: str | Path) -> "SubwordXLMConfig":
+        model_path = Path(pretrained_model_name_or_path)
+        is_local = model_path.is_dir() and (model_path / "config.json").is_file()
         model_config_filepath = Path(
-            pretrained_model_name_or_path
+            (model_path / "config.json")
             if is_local
             else hf_hub_download(pretrained_model_name_or_path, "config.json")
         )

src/wtpsplit_lite/_tokenizer.py

@@ -64,10 +64,13 @@ def __call__(
 
     @classmethod
     @cache
-    def from_pretrained(cls, pretrained_model_name_or_path: str) -> "XLMRobertaTokenizerFast":
-        is_local = Path.is_file(Path(pretrained_model_name_or_path))
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: str | Path
+    ) -> "XLMRobertaTokenizerFast":
+        model_path = Path(pretrained_model_name_or_path)
+        is_local = model_path.is_dir() and (model_path / "tokenizer.json").is_file()
         tokenizer_filepath = Path(
-            pretrained_model_name_or_path
+            (model_path / "tokenizer.json")
             if is_local
             else hf_hub_download(pretrained_model_name_or_path, "tokenizer.json")
         )

tests/test_sat.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import pytest
+from huggingface_hub import hf_hub_download
 from wtpsplit import SaT as SaTOriginal
 
 from wtpsplit_lite import SaT as SaTLite
@@ -103,3 +104,13 @@ def test_weighting(sat: SaTLite, text: str) -> None:
     output_lite = sat.split(text, stride=128, block_size=256, weighting="hat")
     reconstructed_text_lite = "".join(output_lite)
     assert text == reconstructed_text_lite
+
+
+def test_local() -> None:
+    """Test loading a local model."""
+    model_filepath = hf_hub_download("segment-any-text/sat-3l-sm", filename="model_optimized.onnx")
+    model_dir = Path(model_filepath).parent
+    sat_lite = SaTLite(model_dir)
+    text = "This is a test This is another test."
+    output_lite = sat_lite.split(text)
+    assert output_lite == ["This is a test ", "This is another test."]