Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from typing import Optional, List | |
| from src.config import settings | |
| from src.logger import logger | |
| class _STEmbeddings: | |
| """Minimal LangChain-compatible embeddings via sentence-transformers.""" | |
| def __init__(self, model_name: str): | |
| from sentence_transformers import SentenceTransformer | |
| self.model_name = model_name | |
| self._model = SentenceTransformer(model_name) | |
| def embed_documents(self, texts: List[str]) -> List[List[float]]: | |
| return self._model.encode(texts, show_progress_bar=False).tolist() | |
| def embed_query(self, text: str) -> List[float]: | |
| return self._model.encode([text], show_progress_bar=False)[0].tolist() | |
| # LangChain / Chroma sometimes call this | |
| def __call__(self, text: str) -> List[float]: | |
| return self.embed_query(text) | |
| def get_embeddings(model: Optional[str] = None): | |
| model_name = model or settings.embedding_model or "sentence-transformers/all-MiniLM-L6-v2" | |
| logger.info("Using SentenceTransformer embeddings", context={"model": model_name}) | |
| return _STEmbeddings(model_name) |