Spaces:
Sleeping
Sleeping
| """Sentence-transformer embeddings for metadata vector search.""" | |
| from __future__ import annotations | |
| from functools import lru_cache | |
| from metadata_search.config import EMBEDDING_MODEL | |
| def _load_model(model_name: str): | |
| try: | |
| from sentence_transformers import SentenceTransformer | |
| except ImportError as exc: | |
| raise SystemExit( | |
| "Missing dependency: sentence-transformers. Install with `pip install -r requirements.txt`." | |
| ) from exc | |
| try: | |
| return SentenceTransformer(model_name, local_files_only=True) | |
| except Exception: | |
| return SentenceTransformer(model_name) | |
| class MetadataEmbedder: | |
| def __init__(self, model_name: str = EMBEDDING_MODEL, *, batch_size: int = 64): | |
| self.model_name = model_name | |
| self.batch_size = batch_size | |
| self.model = _load_model(model_name) | |
| def embed_documents(self, texts: list[str]) -> list[list[float]]: | |
| if not texts: | |
| return [] | |
| vectors = self.model.encode( | |
| texts, | |
| batch_size=self.batch_size, | |
| normalize_embeddings=True, | |
| show_progress_bar=len(texts) > 128, | |
| ) | |
| return vectors.tolist() | |
| def embed_query(self, query: str) -> list[float]: | |
| vector = self.model.encode( | |
| query, | |
| normalize_embeddings=True, | |
| show_progress_bar=False, | |
| ) | |
| return vector.tolist() | |