diff --git a/README.md b/README.md index d4c88261..18d43077 100644 --- a/README.md +++ b/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 diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index 35e68236..d4574846 100644 --- a/fastembed/common/model_management.py +++ b/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 diff --git a/tests/test_common.py b/tests/test_common.py index f7cae5ba..3bb6c1f0 100644 --- a/tests/test_common.py +++ b/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, + ) diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index 2050a42c..a3cdbde0 100644 --- a/tests/test_custom_models.py +++ b/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()