Spaces:
Sleeping
Sleeping
| """Local embeddings via Sentence-Transformers (no inference provider).""" | |
| from __future__ import annotations | |
| import os | |
| from functools import lru_cache | |
| from time import perf_counter | |
| from typing import cast | |
| from langchain_core.embeddings import Embeddings | |
| from sentence_transformers import SentenceTransformer | |
| from src.config import settings | |
| def _embedding_device() -> str | None: | |
| return "cpu" if os.getenv("SPACE_ID") else None | |
| def _model() -> SentenceTransformer: | |
| return SentenceTransformer(settings.embedding_model, device=_embedding_device()) | |
| class LocalSentenceTransformerEmbeddings(Embeddings): | |
| def embed_documents(self, texts: list[str]) -> list[list[float]]: # type: ignore[override] | |
| vecs = _model().encode(texts, normalize_embeddings=True) | |
| tolist = getattr(vecs, "tolist", None) | |
| if callable(tolist): | |
| return cast(list[list[float]], tolist()) | |
| return [list(map(float, v)) for v in vecs] | |
| def embed_query(self, text: str) -> list[float]: # type: ignore[override] | |
| vec = _model().encode([text], normalize_embeddings=True)[0] | |
| tolist = getattr(vec, "tolist", None) | |
| if callable(tolist): | |
| return cast(list[float], tolist()) | |
| return [float(x) for x in vec] | |
| def get_embeddings() -> Embeddings: | |
| return LocalSentenceTransformerEmbeddings() | |
| def warm_up_embeddings() -> float: | |
| started = perf_counter() | |
| _ = get_embeddings().embed_query("DocuLearn retrieval warmup") | |
| return (perf_counter() - started) * 1000 | |