fix: allow custom models to source from a local directory

README.md

@@ -81,6 +81,22 @@ model = TextEmbedding(model_name="intfloat/multilingual-e5-small")
 embeddings = list(model.embed(documents))
 ```
 
+`hf` also accepts a path to a local directory holding the model files, which is handy for models which are not published on the Hub yet.
+
+```python
+from fastembed import TextEmbedding
+from fastembed.common.model_description import ModelSource, PoolingType
+
+TextEmbedding.add_custom_model(
+    model="my-org/my-model",
+    pooling=PoolingType.MEAN,
+    normalization=True,
+    sources=ModelSource(hf="/path/to/my-model"),  # a local directory is used as is, nothing is downloaded
+    dim=384,
+    model_file="onnx/model.onnx",  # resolved relative to the local directory
+)
+```
+
 
 ### 🔱 Sparse text embeddings
 

fastembed/common/model_management.py

@@ -281,6 +281,49 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]])
 
         return result
 
+    @classmethod
+    def resolve_local_source(cls, model: T, hf_source: str) -> Path | None:
+        """
+        Resolves an `hf` source which points to a local directory instead of a HuggingFace repo id.
+
+        A repo id never refers to an existing directory, so an `hf` source is treated as a local
+        model whenever it resolves to one. This mirrors the way HuggingFace libraries accept either
+        a repo id or a path, and makes it possible to load custom models straight from disk.
+
+        Args:
+            model (T): The model description.
+            hf_source (str): The `hf` field of the model source.
+
+        raises:
+            ValueError: If the directory does not contain all the files the model requires.
+
+        Returns:
+            Optional[Path]: The path to the local model directory, None if the source is a repo id.
+        """
+        model_dir = Path(hf_source).expanduser()
+        if not model_dir.is_dir():
+            return None
+
+        local_root = model_dir.absolute()
+
+        def _is_required_local_file(file: str) -> bool:
+            candidate = Path(os.path.abspath(local_root / file))
+            try:
+                candidate.relative_to(local_root)
+            except ValueError:
+                return False
+            return candidate.is_file()
+
+        required_files = [model.model_file, *model.additional_files]
+        missing_files = [file for file in required_files if not _is_required_local_file(file)]
+        if missing_files:
+            raise ValueError(
+                f"Local directory {model_dir} for model {model.model} is missing the following "
+                f"files: {', '.join(missing_files)}."
+            )
+
+        return model_dir
+
     @classmethod
     def decompress_to_cache(cls, targz_path: str, cache_dir: str) -> str:
         """
@@ -410,6 +453,10 @@ def download_model(cls, model: T, cache_dir: str, retries: int = 3, **kwargs: An
         extra_patterns.extend(model.additional_files)
 
         if hf_source:
+            local_source = cls.resolve_local_source(model, hf_source)
+            if local_source is not None:
+                return local_source
+
             try:
                 cache_kwargs = deepcopy(kwargs)
                 cache_kwargs["local_files_only"] = True

tests/test_common.py

@@ -1,4 +1,8 @@
+import re
+from unittest.mock import patch
+
 import numpy as np
+import pytest
 
 from fastembed import (
     TextEmbedding,
@@ -7,6 +11,8 @@
     LateInteractionMultimodalEmbedding,
     LateInteractionTextEmbedding,
 )
+from fastembed.common.model_description import BaseModelDescription, ModelSource
+from fastembed.common.model_management import ModelManagement
 from fastembed.common.utils import last_token_pooling
 
 
@@ -59,3 +65,113 @@ def test_last_token_pooling_with_left_padding():
     pooled = last_token_pooling(token_embeddings, attention_mask)
 
     assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])
+
+
+def _make_local_model_dir(root, model_file="onnx/model.onnx", additional_files=()):
+    """Create a directory that looks like a downloaded model snapshot."""
+    for rel_path in ("config.json", "tokenizer.json", model_file, *additional_files):
+        file_path = root / rel_path
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_path.write_text("{}")
+    return root
+
+
+def _local_model_description(hf_source, model_file="onnx/model.onnx", additional_files=None):
+    return BaseModelDescription(
+        model="test-org/test-model",
+        sources=ModelSource(hf=hf_source),
+        model_file=model_file,
+        description="",
+        license="",
+        size_in_GB=0.1,
+        additional_files=additional_files or [],
+    )
+
+
+def test_local_directory_hf_source_is_used_as_is(tmp_path):
+    """An `hf` source pointing to a local directory must be used without touching the hub."""
+    model_dir = _make_local_model_dir(tmp_path / "my-model")
+    model = _local_model_description(str(model_dir))
+
+    with patch.object(ModelManagement, "download_files_from_huggingface") as mock_download:
+        resolved_path = ModelManagement.download_model(model, cache_dir=str(tmp_path / "cache"))
+
+    assert resolved_path == model_dir
+    mock_download.assert_not_called()
+
+
+def test_local_directory_hf_source_expands_user(tmp_path, monkeypatch):
+    """`~` in a local directory source must be expanded."""
+    monkeypatch.setenv("HOME", str(tmp_path))
+    monkeypatch.setenv("USERPROFILE", str(tmp_path))  # windows
+    model_dir = _make_local_model_dir(tmp_path / "my-model")
+    model = _local_model_description("~/my-model")
+
+    with patch.object(ModelManagement, "download_files_from_huggingface") as mock_download:
+        resolved_path = ModelManagement.download_model(model, cache_dir=str(tmp_path / "cache"))
+
+    assert resolved_path == model_dir
+    mock_download.assert_not_called()
+
+
+def test_local_directory_hf_source_with_additional_files(tmp_path):
+    """`additional_files` must be resolved relative to the local directory as well."""
+    model_dir = _make_local_model_dir(
+        tmp_path / "my-model", additional_files=("stopwords/en.txt",)
+    )
+    model = _local_model_description(str(model_dir), additional_files=["stopwords/en.txt"])
+
+    resolved_path = ModelManagement.download_model(model, cache_dir=str(tmp_path / "cache"))
+
+    assert resolved_path == model_dir
+
+
+def test_local_directory_hf_source_missing_files(tmp_path):
+    """An incomplete local directory must fail loudly instead of falling back to the hub."""
+    model_dir = _make_local_model_dir(tmp_path / "my-model")
+    model = _local_model_description(
+        str(model_dir), additional_files=["stopwords/en.txt", "vocab.txt"]
+    )
+
+    with patch.object(ModelManagement, "download_files_from_huggingface") as mock_download:
+        with pytest.raises(ValueError, match="stopwords/en.txt, vocab.txt"):
+            ModelManagement.download_model(model, cache_dir=str(tmp_path / "cache"))
+
+    mock_download.assert_not_called()
+
+
+def test_local_directory_hf_source_rejects_path_traversal(tmp_path):
+    """A required file escaping the local directory (via `..` or an absolute path) must not resolve."""
+    model_dir = _make_local_model_dir(tmp_path / "my-model")
+    outside_file = tmp_path / "outside.onnx"
+    outside_file.write_text("{}")
+
+    traversal_model = _local_model_description(str(model_dir), model_file="../outside.onnx")
+    absolute_model = _local_model_description(str(model_dir), model_file=str(outside_file))
+
+    with patch.object(ModelManagement, "download_files_from_huggingface") as mock_download:
+        with pytest.raises(ValueError, match=re.escape("../outside.onnx")):
+            ModelManagement.download_model(traversal_model, cache_dir=str(tmp_path / "cache"))
+        with pytest.raises(ValueError, match=re.escape(str(outside_file))):
+            ModelManagement.download_model(absolute_model, cache_dir=str(tmp_path / "cache"))
+
+    mock_download.assert_not_called()
+
+
+def test_repo_id_hf_source_is_not_treated_as_local_directory(tmp_path):
+    """A repo id, which does not resolve to a directory, must still go through the hub."""
+    model = _local_model_description("test-org/test-model")
+    snapshot_dir = _make_local_model_dir(tmp_path / "snapshot")
+
+    with patch.object(
+        ModelManagement, "download_files_from_huggingface", return_value=str(snapshot_dir)
+    ) as mock_download:
+        resolved_path = ModelManagement.download_model(model, cache_dir=str(tmp_path / "cache"))
+
+    assert resolved_path == snapshot_dir
+    mock_download.assert_called_with(
+        "test-org/test-model",
+        cache_dir=str(tmp_path / "cache"),
+        extra_patterns=["onnx/model.onnx"],
+        local_files_only=True,
+    )

tests/test_custom_models.py

@@ -250,3 +250,39 @@ def test_do_not_add_existing_cross_encoder():
         )
 
     CustomTextCrossEncoder.SUPPORTED_MODELS.clear()
+
+
+def test_text_custom_model_from_local_directory():
+    """A custom model can be sourced from a local directory instead of a HuggingFace repo id."""
+    is_ci = os.getenv("CI")
+    base_model_name = "sentence-transformers/all-MiniLM-L6-v2"
+    custom_model_name = "custom/all-MiniLM-L6-v2-local"
+    dim = 384
+    docs = ["hello world", "flag embedding"]
+
+    base_model = TextEmbedding(base_model_name)
+    local_dir = base_model.model._model_dir
+    expected = np.stack(list(base_model.embed(docs)), axis=0)
+
+    TextEmbedding.add_custom_model(
+        custom_model_name,
+        pooling=PoolingType.MEAN,
+        normalization=True,
+        sources=ModelSource(hf=str(local_dir)),
+        dim=dim,
+        model_file="model.onnx",
+        size_in_gb=0.09,
+    )
+
+    model = TextEmbedding(custom_model_name)
+    assert model.model._model_dir == local_dir
+
+    embeddings = np.stack(list(model.embed(docs)), axis=0)
+    assert embeddings.shape == (2, dim)
+    assert np.allclose(embeddings, expected, atol=1e-3)
+
+    if is_ci:
+        delete_model_cache(local_dir)
+
+    CustomTextEmbedding.SUPPORTED_MODELS.clear()
+    CustomTextEmbedding.POSTPROCESSING_MAPPING.clear()