diff --git a/python/python/lancedb/embeddings/instructor.py b/python/python/lancedb/embeddings/instructor.py index 675a0139c..4ea605bd8 100644 --- a/python/python/lancedb/embeddings/instructor.py +++ b/python/python/lancedb/embeddings/instructor.py @@ -3,6 +3,7 @@ from typing import List +from urllib.parse import unquote, urlparse import numpy as np @@ -125,6 +126,10 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction): @weak_lru(maxsize=1) def get_model(self): + huggingface_hub = attempt_import_or_raise("huggingface_hub", "huggingface-hub") + if not hasattr(huggingface_hub, "cached_download"): + huggingface_hub.cached_download = _cached_download(huggingface_hub) + instructor_embedding = attempt_import_or_raise( "InstructorEmbedding", "InstructorEmbedding" ) @@ -140,3 +145,44 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction): model, {torch.nn.Linear}, dtype=torch.qint8 ) return model + + +def _cached_download(huggingface_hub): + """Provide the legacy download API used by sentence-transformers 2.2.x.""" + + def cached_download( + *, + url, + cache_dir=None, + force_filename=None, + library_name=None, + library_version=None, + user_agent=None, + use_auth_token=None, + **_, + ): + path = urlparse(url).path.lstrip("/") + try: + repo_id, resolved_path = path.split("/resolve/", maxsplit=1) + revision, filename = resolved_path.split("/", maxsplit=1) + except ValueError as err: + raise ValueError(f"Unsupported Hugging Face Hub URL: {url}") from err + + repo_id = unquote(repo_id) + revision = unquote(revision) + filename = unquote(filename) + if force_filename is not None: + filename = force_filename + + return huggingface_hub.hf_hub_download( + repo_id=repo_id, + filename=filename, + revision=revision, + local_dir=cache_dir, + library_name=library_name, + library_version=library_version, + user_agent=user_agent, + token=use_auth_token, + ) + + return cached_download diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 5efb7d98a..393d55625 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -3,6 +3,7 @@ import os import pickle +from types import SimpleNamespace from typing import List, Optional, Union from unittest.mock import MagicMock, patch @@ -522,6 +523,58 @@ def test_embedding_function_safe_model_dump(embedding_type): ) +def test_instructor_embedding_supports_huggingface_hub_without_cached_download(): + from lancedb.embeddings.instructor import InstructorEmbeddingFunction + + hub_download = MagicMock(return_value="/cache/1_Pooling/config.json") + huggingface_hub = SimpleNamespace(hf_hub_download=hub_download) + instructor_model = MagicMock() + instructor_embedding = SimpleNamespace( + INSTRUCTOR=MagicMock(return_value=instructor_model) + ) + + def import_dependency(module, _mitigation): + if module == "huggingface_hub": + return huggingface_hub + if module == "InstructorEmbedding": + assert hasattr(huggingface_hub, "cached_download") + return instructor_embedding + if module == "torch": + return SimpleNamespace() + raise AssertionError(f"Unexpected import: {module}") + + with patch( + "lancedb.embeddings.instructor.attempt_import_or_raise", + side_effect=import_dependency, + ): + embedding = InstructorEmbeddingFunction.create(show_progress_bar=False) + assert embedding.get_model() is instructor_model + + path = huggingface_hub.cached_download( + url=( + "https://huggingface.co/hkunlp/instructor-base/resolve/abc123/" + "1_Pooling/config.json" + ), + cache_dir="/cache", + force_filename="1_Pooling/config.json", + library_name="sentence-transformers", + library_version="2.2.2", + use_auth_token="token", + ) + + assert path == "/cache/1_Pooling/config.json" + hub_download.assert_called_once_with( + repo_id="hkunlp/instructor-base", + filename="1_Pooling/config.json", + revision="abc123", + local_dir="/cache", + library_name="sentence-transformers", + library_version="2.2.2", + user_agent=None, + token="token", + ) + + @patch("time.sleep") def test_retry(mock_sleep): test_function = MagicMock(side_effect=[Exception] * 9 + ["result"])