Spaces:
Sleeping
Sleeping
File size: 2,684 Bytes
b92b744 57a9a08 b92b744 57a9a08 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """Local MiniLM embeddings shared by offline indexing and online retrieval.
No hosted embedding endpoint is used. Loading/downloading a model is deferred
until embeddings are requested, so BM25 and metadata-only operation stay light.
"""
import asyncio
from functools import lru_cache
DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
class Embedder:
def __init__(self, model_name: str = DEFAULT_MODEL, device: str = "cpu", batch_size: int = 64):
self.model_name = model_name
self.device = device
self.batch_size = batch_size
self._model = None
def _load(self):
if self._model is None:
try:
from sentence_transformers import SentenceTransformer
except ImportError as exc:
raise RuntimeError("Install the ingest extra to enable local embeddings: uv sync --extra ingest") from exc
try:
# Avoid repeated network checks when the shared model cache is already populated.
self._model = SentenceTransformer(self.model_name, device=self.device,
local_files_only=True)
except OSError:
self._model = SentenceTransformer(self.model_name, device=self.device)
return self._model
def encode(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
return self._load().encode(texts, batch_size=self.batch_size, normalize_embeddings=True,
show_progress_bar=False, convert_to_numpy=True).tolist()
def embed_query(self, text: str) -> list[float]:
return self.encode([text])[0]
def embed_documents(self, texts: list[str]) -> list[list[float]]:
return self.encode(texts)
@lru_cache(maxsize=4)
def get_embedder(model_name: str = DEFAULT_MODEL, device: str = "cpu") -> Embedder:
return Embedder(model_name=model_name, device=device)
class QueryEmbedder:
"""One query embedding per call, loaded on first use and encoded off the loop.
``None`` rather than a zero vector when embeddings are disabled: a caller
must say it could not embed, not search with a direction it invented.
"""
def __init__(self, config):
self.config = config
self.embedder: Embedder | None = None
async def embed(self, text: str) -> list[float] | None:
if not self.config.embeddings_enabled:
return None
if self.embedder is None:
self.embedder = get_embedder(self.config.embedding_model, self.config.embedding_device)
return await asyncio.to_thread(self.embedder.embed_query, text)
|