From 873589bc8e87990d98d0f45d1b23bd9d32ce57bc Mon Sep 17 00:00:00 2001 From: asimurka Date: Tue, 27 Oct 2026 10:46:30 +0100 Subject: [PATCH] Implemented sentence transformers provider --- src/providers/inference/registry.py | 4 + .../inference/sentence_transformers/README.md | 7 + .../sentence_transformers/__init__.py | 6 + .../sentence_transformers/model_cache.py | 72 ++++ .../sentence_transformers/provider.py | 185 ++++++++++ .../inference/sentence_transformers/README.md | 13 + .../sentence_transformers/__init__.py | 1 + .../sentence_transformers/test_model_cache.py | 39 ++ .../sentence_transformers/test_provider.py | 332 ++++++++++++++++++ 9 files changed, 659 insertions(+) create mode 100644 src/providers/inference/sentence_transformers/model_cache.py create mode 100644 src/providers/inference/sentence_transformers/provider.py create mode 100644 tests/unit/providers/inference/sentence_transformers/README.md create mode 100644 tests/unit/providers/inference/sentence_transformers/__init__.py create mode 100644 tests/unit/providers/inference/sentence_transformers/test_model_cache.py create mode 100644 tests/unit/providers/inference/sentence_transformers/test_provider.py diff --git a/src/providers/inference/registry.py b/src/providers/inference/registry.py index 7daee4e27..5458bd85d 100644 --- a/src/providers/inference/registry.py +++ b/src/providers/inference/registry.py @@ -8,6 +8,9 @@ from providers.inference.bedrock.provider import BedrockInferenceProvider from providers.inference.errors import ProviderNotFoundError from providers.inference.openai.provider import OpenAIInferenceProvider +from providers.inference.sentence_transformers.provider import ( + SentenceTransformersInferenceProvider, +) from providers.inference.vertexai.provider import VertexAIInferenceProvider from providers.inference.watsonx.provider import WatsonxInferenceProvider from utils.types import Singleton @@ -24,6 +27,7 @@ "vertexai": VertexAIInferenceProvider, "bedrock": BedrockInferenceProvider, "watsonx": WatsonxInferenceProvider, + "sentence_transformers": SentenceTransformersInferenceProvider, } diff --git a/src/providers/inference/sentence_transformers/README.md b/src/providers/inference/sentence_transformers/README.md index 14e2e9656..ba9ebfc8c 100644 --- a/src/providers/inference/sentence_transformers/README.md +++ b/src/providers/inference/sentence_transformers/README.md @@ -4,3 +4,10 @@ Sentence-transformers inference provider package. +## [model_cache.py](model_cache.py) + +In-process cache for sentence-transformers embedding and rerank models. + +## [provider.py](provider.py) + +Local sentence-transformers inference provider (embedding and rerank). diff --git a/src/providers/inference/sentence_transformers/__init__.py b/src/providers/inference/sentence_transformers/__init__.py index f0803640e..2db732089 100644 --- a/src/providers/inference/sentence_transformers/__init__.py +++ b/src/providers/inference/sentence_transformers/__init__.py @@ -1 +1,7 @@ """Sentence-transformers inference provider package.""" + +from providers.inference.sentence_transformers.provider import ( + SentenceTransformersInferenceProvider, +) + +__all__ = ["SentenceTransformersInferenceProvider"] diff --git a/src/providers/inference/sentence_transformers/model_cache.py b/src/providers/inference/sentence_transformers/model_cache.py new file mode 100644 index 000000000..32e407904 --- /dev/null +++ b/src/providers/inference/sentence_transformers/model_cache.py @@ -0,0 +1,72 @@ +"""In-process cache for sentence-transformers embedding and rerank models.""" + +from __future__ import annotations + +from typing import Literal, Optional, overload + +from pydantic_ai.embeddings import EmbeddingModel +from sentence_transformers import CrossEncoder + + +class SentenceTransformersModelCache: + """Process-local cache of loaded embedding and rerank model instances.""" + + def __init__(self) -> None: + """Initialize empty embedding and rerank caches.""" + self._embeddings: dict[str, EmbeddingModel] = {} + self._rerankers: dict[str, CrossEncoder] = {} + + @overload + def get( + self, + provider_model_id: str, + model_type: Literal["embedding"], + ) -> Optional[EmbeddingModel]: ... + + @overload + def get( + self, + provider_model_id: str, + model_type: Literal["rerank"], + ) -> Optional[CrossEncoder]: ... + + def get( + self, + provider_model_id: str, + model_type: Literal["embedding", "rerank"], + ) -> Optional[EmbeddingModel | CrossEncoder]: + """Return a cached model when present. + + Args: + provider_model_id: Cache key; the Hugging Face model id. + model_type: Capability whose cache to read (embedding or rerank). + + Returns: + The cached model, or None if it has not been stored. + + """ + match model_type: + case "embedding": + return self._embeddings.get(provider_model_id) + case "rerank": + return self._rerankers.get(provider_model_id) + + def put( + self, + provider_model_id: str, + model: EmbeddingModel | CrossEncoder, + ) -> None: + """Store a model in the cache. + + The target store is chosen from the instance type: CrossEncoder goes + to the rerank cache, anything else to the embedding cache. + + Args: + provider_model_id: Cache key; the Hugging Face model id. + model: Model instance to cache. + + """ + if isinstance(model, CrossEncoder): + self._rerankers[provider_model_id] = model + else: + self._embeddings[provider_model_id] = model diff --git a/src/providers/inference/sentence_transformers/provider.py b/src/providers/inference/sentence_transformers/provider.py new file mode 100644 index 000000000..68bd997ff --- /dev/null +++ b/src/providers/inference/sentence_transformers/provider.py @@ -0,0 +1,185 @@ +"""Local sentence-transformers inference provider (embedding and rerank).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional + +from pydantic_ai.embeddings.sentence_transformers import ( + SentenceTransformerEmbeddingModel, +) +from pydantic_ai.models import Model +from pydantic_ai.providers import Provider +from sentence_transformers import CrossEncoder, SentenceTransformer + +from log import get_logger +from models.config import SentenceTransformersInferenceProviderEntry +from providers.inference.base import InferenceProvider +from providers.inference.catalog import CatalogModel, DiscoveredModel +from providers.inference.errors import ModelTypeNotSupportedError +from providers.inference.sentence_transformers.model_cache import ( + SentenceTransformersModelCache, +) + +if TYPE_CHECKING: + from pydantic_ai.embeddings import EmbeddingModel + + from providers.inference.types import ModelType + +logger = get_logger(__name__) + + +class SentenceTransformersInferenceProvider( + InferenceProvider[SentenceTransformersInferenceProviderEntry, Provider[Any]] +): + """Inference provider for local sentence-transformers embedding and rerank. + + This provider has no HTTP client and no remote discovery. The configured + allowed_models list is the catalog. Chat models are never exposed. + """ + + def __init__(self, entry: SentenceTransformersInferenceProviderEntry) -> None: + """Initialize with an empty in-process model cache. + + Args: + entry: Provider configuration entry with a non-empty allow-list. + + """ + super().__init__(entry) + self._cache = SentenceTransformersModelCache() + + def _build_provider(self) -> Optional[Provider[Any]]: + """Return no Pydantic AI provider; models run in-process. + + Returns: + Always None. + + """ + return None + + async def _discover_models(self) -> list[DiscoveredModel]: + """Remote discovery is unsupported for sentence-transformers. + + Returns: + An empty list. Catalog construction uses allowed_models only. + + """ + return [] + + async def _load_catalog(self) -> list[CatalogModel]: + """Build the catalog exclusively from configured allowed models. + + Equivalent to the base-class skip_model_discovery path. Entries typed + as llm are omitted so this provider never exposes chat models. + + Returns: + Catalog rows for embedding and rerank models only. + + """ + catalog: list[CatalogModel] = [] + for allowed in self._entry.allowed_models: + model_type = allowed.model_type or self._infer_model_type( + allowed.provider_model_id + ) + if model_type == "llm": + logger.warning( + "Omitting configured model %r with provider model ID %r " + "from sentence-transformers provider %r; llm is unsupported", + allowed.name, + allowed.provider_model_id, + self.provider_id, + ) + continue + catalog.append( + CatalogModel( + name=allowed.name, + provider_model_id=allowed.provider_model_id, + model_type=model_type, + metadata=dict(allowed.metadata), + ) + ) + return catalog + + def _create_llm(self, model: CatalogModel) -> Model: + """Reject LLM construction; this provider never serves chat models. + + Args: + model: Catalog entry that requested LLM support. + + Raises: + ModelTypeNotSupportedError: Always, because llm is unsupported. + + """ + _ = model + raise ModelTypeNotSupportedError("llm", self.provider_id) + + def _create_embedding(self, model: CatalogModel) -> EmbeddingModel: + """Create or reuse a cached Pydantic AI sentence-transformers embedding model. + + Args: + model: Catalog entry describing the embedding model to instantiate. + + Returns: + A cached Pydantic AI embedding model for the provider model id. + + """ + provider_model_id = model.provider_model_id + cached = self._cache.get(provider_model_id, "embedding") + if cached is not None: + return cached + + config = self._entry.config + embedding = SentenceTransformerEmbeddingModel( + SentenceTransformer( + provider_model_id, + device=config.device, + cache_folder=config.cache_dir, + ) + ) + self._cache.put(provider_model_id, embedding) + return embedding + + def _create_rerank(self, model: CatalogModel) -> CrossEncoder: + """Create or reuse a cached sentence-transformers cross-encoder. + + Args: + model: Catalog entry describing the reranker to instantiate. + + Returns: + A cached CrossEncoder for the provider model id. + + """ + provider_model_id = model.provider_model_id + cached = self._cache.get(provider_model_id, "rerank") + if cached is not None: + return cached + + config = self._entry.config + reranker = CrossEncoder( + provider_model_id, + device=config.device, + cache_folder=config.cache_dir, + ) + self._cache.put(provider_model_id, reranker) + return reranker + + def _infer_model_type( + self, + provider_model_id: str, + **kwargs: Any, + ) -> ModelType: + """Infer embedding vs rerank when allowed_models model_type is omitted. + + Args: + provider_model_id: Hugging Face model id or local path. + **kwargs: Unused; accepted for base-class compatibility. + + Returns: + rerank when the id looks like a cross-encoder or reranker, + otherwise embedding. Never llm. + + """ + _ = kwargs + normalized = provider_model_id.lower().replace("-", "").replace("_", "") + if "crossencoder" in normalized or "rerank" in normalized: + return "rerank" + return "embedding" diff --git a/tests/unit/providers/inference/sentence_transformers/README.md b/tests/unit/providers/inference/sentence_transformers/README.md new file mode 100644 index 000000000..e75a277aa --- /dev/null +++ b/tests/unit/providers/inference/sentence_transformers/README.md @@ -0,0 +1,13 @@ +# List of source files stored in `tests/unit/providers/inference/sentence_transformers` directory + +## [__init__.py](__init__.py) + +Unit tests for the sentence-transformers inference provider. + +## [test_model_cache.py](test_model_cache.py) + +Unit tests for the in-process sentence-transformers model cache. + +## [test_provider.py](test_provider.py) + +Unit tests for the sentence-transformers inference provider. diff --git a/tests/unit/providers/inference/sentence_transformers/__init__.py b/tests/unit/providers/inference/sentence_transformers/__init__.py new file mode 100644 index 000000000..e8fac1717 --- /dev/null +++ b/tests/unit/providers/inference/sentence_transformers/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the sentence-transformers inference provider.""" diff --git a/tests/unit/providers/inference/sentence_transformers/test_model_cache.py b/tests/unit/providers/inference/sentence_transformers/test_model_cache.py new file mode 100644 index 000000000..caa60e119 --- /dev/null +++ b/tests/unit/providers/inference/sentence_transformers/test_model_cache.py @@ -0,0 +1,39 @@ +"""Unit tests for the in-process sentence-transformers model cache.""" + +from typing import Any + +from sentence_transformers import CrossEncoder + +from providers.inference.sentence_transformers.model_cache import ( + SentenceTransformersModelCache, +) + + +def test_get_and_put_embedding_caches_by_provider_model_id(mocker: Any) -> None: + """Stores embeddings by instance type and provider model id.""" + cache = SentenceTransformersModelCache() + embed_a = mocker.Mock(name="embed-a") + embed_b = mocker.Mock(name="embed-b") + + assert cache.get("model-a", "embedding") is None + cache.put("model-a", embed_a) + cache.put("model-b", embed_b) + + assert cache.get("model-a", "embedding") is embed_a + assert cache.get("model-b", "embedding") is embed_b + assert cache.get("model-a", "rerank") is None + + +def test_get_and_put_rerank_caches_by_provider_model_id(mocker: Any) -> None: + """Stores cross-encoders by instance type and provider model id.""" + cache = SentenceTransformersModelCache() + rerank_a = mocker.Mock(name="rerank-a", spec=CrossEncoder) + rerank_b = mocker.Mock(name="rerank-b", spec=CrossEncoder) + + assert cache.get("model-a", "rerank") is None + cache.put("model-a", rerank_a) + cache.put("model-b", rerank_b) + + assert cache.get("model-a", "rerank") is rerank_a + assert cache.get("model-b", "rerank") is rerank_b + assert cache.get("model-a", "embedding") is None diff --git a/tests/unit/providers/inference/sentence_transformers/test_provider.py b/tests/unit/providers/inference/sentence_transformers/test_provider.py new file mode 100644 index 000000000..e639f4f25 --- /dev/null +++ b/tests/unit/providers/inference/sentence_transformers/test_provider.py @@ -0,0 +1,332 @@ +"""Unit tests for the sentence-transformers inference provider.""" + +# pylint: disable=redefined-outer-name,protected-access + +from typing import Any, Optional + +import pytest +from sentence_transformers import CrossEncoder + +from models.config import SentenceTransformersInferenceProviderEntry +from providers.inference.catalog import CatalogModel +from providers.inference.errors import ( + ModelNotAvailableError, + ModelTypeNotSupportedError, +) +from providers.inference.registry import PROVIDER_CLASSES +from providers.inference.sentence_transformers.provider import ( + SentenceTransformersInferenceProvider, +) + + +def _entry( + *, + allowed_models: Optional[list[Any]] = None, + device: Optional[str] = None, + cache_dir: Optional[str] = None, +) -> SentenceTransformersInferenceProviderEntry: + """Build a sentence-transformers provider configuration.""" + config: dict[str, Any] = {} + if device is not None: + config["device"] = device + if cache_dir is not None: + config["cache_dir"] = cache_dir + return SentenceTransformersInferenceProviderEntry.model_validate( + { + "type": "sentence_transformers", + "id": "sentence_transformers", + "allowed_models": allowed_models + or [ + { + "name": "embed", + "provider_model_id": "sentence-transformers/all-MiniLM-L6-v2", + "model_type": "embedding", + }, + { + "name": "rerank", + "provider_model_id": "cross-encoder/ms-marco-MiniLM-L-6-v2", + "model_type": "rerank", + }, + ], + "config": config, + } + ) + + +def _catalog_model( + *, + name: str, + provider_model_id: str, + model_type: str, +) -> CatalogModel: + """Build a catalog model for create_* tests.""" + return CatalogModel( + name=name, + provider_model_id=provider_model_id, + model_type=model_type, # type: ignore[arg-type] + ) + + +@pytest.fixture +def provider() -> SentenceTransformersInferenceProvider: + """Create a sentence-transformers provider with embed and rerank models.""" + return SentenceTransformersInferenceProvider( + _entry(device="cpu", cache_dir="/tmp/hf-cache") + ) + + +def test_build_provider_returns_none( + provider: SentenceTransformersInferenceProvider, +) -> None: + """Inline provider has no PAI HTTP client or provider instance.""" + assert provider._build_provider() is None + + +@pytest.mark.asyncio +async def test_discover_models_returns_empty( + provider: SentenceTransformersInferenceProvider, +) -> None: + """Remote discovery is unsupported; catalog comes from allowed_models.""" + assert await provider._discover_models() == [] + + +@pytest.mark.asyncio +async def test_load_catalog_from_allowed_models() -> None: + """Allow-list only: infers types, prefers explicit model_type, omits llm.""" + provider = SentenceTransformersInferenceProvider( + _entry( + allowed_models=[ + "sentence-transformers/all-mpnet-base-v2", + "cross-encoder/ms-marco-MiniLM-L-12-v2", + { + "name": "forced-embed", + "provider_model_id": "cross-encoder/looks-like-rerank", + "model_type": "embedding", + }, + { + "name": "chat", + "provider_model_id": "some-chat-model", + "model_type": "llm", + }, + ] + ) + ) + + catalog = await provider._load_catalog() + + assert [ + (model.name, model.provider_model_id, model.model_type) for model in catalog + ] == [ + ( + "sentence-transformers/all-mpnet-base-v2", + "sentence-transformers/all-mpnet-base-v2", + "embedding", + ), + ( + "cross-encoder/ms-marco-MiniLM-L-12-v2", + "cross-encoder/ms-marco-MiniLM-L-12-v2", + "rerank", + ), + ( + "forced-embed", + "cross-encoder/looks-like-rerank", + "embedding", + ), + ] + + +@pytest.mark.asyncio +async def test_initialize_does_not_load_models(mocker: Any) -> None: + """Models stay lazy; initialize only builds the catalog.""" + sentence_transformer_cls = mocker.patch( + "providers.inference.sentence_transformers.provider.SentenceTransformer", + ) + cross_encoder_cls = mocker.patch( + "providers.inference.sentence_transformers.provider.CrossEncoder", + ) + provider = SentenceTransformersInferenceProvider(_entry()) + + await provider.initialize() + + sentence_transformer_cls.assert_not_called() + cross_encoder_cls.assert_not_called() + assert provider._provider is None + assert [ + (model.name, model.model_type) for model in provider.list_catalog_models() + ] == [ + ("embed", "embedding"), + ("rerank", "rerank"), + ] + + +@pytest.mark.parametrize( + ("device", "cache_dir"), + [ + ("cpu", "/tmp/hf-cache"), + (None, None), + ], +) +def test_create_embedding_passes_config_and_caches( + mocker: Any, + device: Optional[str], + cache_dir: Optional[str], +) -> None: + """Builds SentenceTransformer with config, wraps in PAI, and caches.""" + provider = SentenceTransformersInferenceProvider( + _entry(device=device, cache_dir=cache_dir) + ) + runtime = mocker.Mock(name="sentence-transformer") + wrapped = mocker.Mock(name="embedding-model") + sentence_transformer_cls = mocker.patch( + "providers.inference.sentence_transformers.provider.SentenceTransformer", + return_value=runtime, + ) + embedding_model_cls = mocker.patch( + "providers.inference.sentence_transformers.provider." + "SentenceTransformerEmbeddingModel", + return_value=wrapped, + ) + catalog = _catalog_model( + name="embed", + provider_model_id="sentence-transformers/all-MiniLM-L6-v2", + model_type="embedding", + ) + + first = provider._create_embedding(catalog) + second = provider._create_embedding(catalog) + + assert first is wrapped + assert second is wrapped + sentence_transformer_cls.assert_called_once_with( + "sentence-transformers/all-MiniLM-L6-v2", + device=device, + cache_folder=cache_dir, + ) + embedding_model_cls.assert_called_once_with(runtime) + + +@pytest.mark.parametrize( + ("device", "cache_dir"), + [ + ("cpu", "/tmp/hf-cache"), + (None, None), + ], +) +def test_create_rerank_passes_config_and_caches( + mocker: Any, + device: Optional[str], + cache_dir: Optional[str], +) -> None: + """Builds CrossEncoder with config and caches the instance.""" + provider = SentenceTransformersInferenceProvider( + _entry(device=device, cache_dir=cache_dir) + ) + created = mocker.Mock(name="cross-encoder", spec=CrossEncoder) + cross_encoder_cls = mocker.patch( + "providers.inference.sentence_transformers.provider.CrossEncoder", + return_value=created, + ) + catalog = _catalog_model( + name="rerank", + provider_model_id="cross-encoder/ms-marco-MiniLM-L-6-v2", + model_type="rerank", + ) + + first = provider._create_rerank(catalog) + second = provider._create_rerank(catalog) + + assert first is created + assert second is created + cross_encoder_cls.assert_called_once_with( + "cross-encoder/ms-marco-MiniLM-L-6-v2", + device=device, + cache_folder=cache_dir, + ) + + +def test_create_llm_is_unsupported( + provider: SentenceTransformersInferenceProvider, +) -> None: + """Rejects LLM construction because sentence-transformers never serves chat.""" + with pytest.raises(ModelTypeNotSupportedError) as exc_info: + provider._create_llm( + _catalog_model( + name="chat", + provider_model_id="should-not-load", + model_type="llm", + ) + ) + + assert exc_info.value.model_type == "llm" + assert exc_info.value.provider_id == "sentence_transformers" + + +@pytest.mark.asyncio +async def test_get_model_instance_resolves_embed_and_rerank( + provider: SentenceTransformersInferenceProvider, + mocker: Any, +) -> None: + """Catalog names resolve to cached embedding and rerank instances.""" + await provider.initialize() + embedding = mocker.Mock(name="embedding") + reranker = mocker.Mock(name="rerank", spec=CrossEncoder) + mocker.patch( + "providers.inference.sentence_transformers.provider.SentenceTransformer", + ) + mocker.patch( + "providers.inference.sentence_transformers.provider." + "SentenceTransformerEmbeddingModel", + return_value=embedding, + ) + mocker.patch( + "providers.inference.sentence_transformers.provider.CrossEncoder", + return_value=reranker, + ) + + assert provider.get_model_instance("embed", "embedding") is embedding + assert provider.get_model_instance("rerank", "rerank") is reranker + + +@pytest.mark.asyncio +async def test_get_llm_raises_model_not_available( + provider: SentenceTransformersInferenceProvider, +) -> None: + """LLM requests fail with ModelNotAvailableError; llm is never catalogued.""" + await provider.initialize() + + with pytest.raises(ModelNotAvailableError) as exc_info: + provider.get_model_instance("embed", "llm") + + assert exc_info.value.model_type == "llm" + assert exc_info.value.provider_id == "sentence_transformers" + assert exc_info.value.name == "embed" + + +@pytest.mark.parametrize( + ("model_id", "expected"), + [ + ("sentence-transformers/all-MiniLM-L6-v2", "embedding"), + ("sentence-transformers/all-mpnet-base-v2", "embedding"), + ("nomic-ai/nomic-embed-text-v1.5", "embedding"), + ("cross-encoder/ms-marco-MiniLM-L-6-v2", "rerank"), + ("BAAI/bge-reranker-base", "rerank"), + ("org/my_cross_encoder_model", "rerank"), + ("antoinelouis/crossencoder-mMiniLMv2-L6-mmarcoFR", "rerank"), + ("vendor/text-rerank-v1", "rerank"), + ], +) +def test_infer_model_type( + provider: SentenceTransformersInferenceProvider, + model_id: str, + expected: str, +) -> None: + """Classifies Hub ids as embedding or rerank; never as llm.""" + assert provider._infer_model_type(model_id) == expected + + +def test_provider_registered_on_provider_classes() -> None: + """Registry maps sentence_transformers to this provider class.""" + assert ( + PROVIDER_CLASSES["sentence_transformers"] + is SentenceTransformersInferenceProvider + )