| from typing import List
|
| from functools import lru_cache
|
|
|
| import numpy as np
|
| from sentence_transformers import SentenceTransformer
|
|
|
|
|
| class EmbeddingService:
|
| def __init__(self, model_name: str) -> None:
|
| self.model = SentenceTransformer(model_name)
|
|
|
| def encode(self, texts: List[str]) -> np.ndarray:
|
| vectors = self.model.encode(
|
| texts,
|
| normalize_embeddings=True,
|
| show_progress_bar=False,
|
| convert_to_numpy=True,
|
| )
|
| return np.asarray(vectors, dtype=np.float32)
|
|
|
| @lru_cache(maxsize=2048)
|
| def encode_query_cached(self, text: str) -> bytes:
|
| vec = self.model.encode(
|
| [text],
|
| normalize_embeddings=True,
|
| show_progress_bar=False,
|
| convert_to_numpy=True,
|
| )
|
| arr = np.asarray(vec, dtype=np.float32)
|
| return arr.tobytes()
|
|
|
| def encode_query(self, text: str) -> np.ndarray:
|
| buf = self.encode_query_cached(text)
|
| return np.frombuffer(buf, dtype=np.float32).reshape(1, -1)
|
|
|
