Spaces:
Runtime error
Runtime error
File size: 2,293 Bytes
d4ec2fa | 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 | """
Embeds a live search query into the same 384-dim space as the Chroma index (Ticket 7).
Implements the standard sentence-transformers "mean pooling + L2 normalize" approach
manually with plain `transformers`, rather than the `sentence-transformers` package's
special model format -- this lets us reuse the exact same "download 4 files from HF,
point a local path at them" pattern that already worked for DistilBERT (Tickets 4/5),
instead of dealing with the sentence-transformers library's different folder layout.
"""
import logging
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
from api.config import BASE_DIR, MODEL_DIR
logger = logging.getLogger(__name__)
EMBEDDING_MODEL_PATH = BASE_DIR / MODEL_DIR / "embedding_model"
def _mean_pooling(model_output, attention_mask):
token_embeddings = model_output[0]
mask = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)
class EmbeddingService:
def __init__(self):
self.tokenizer = None
self.model = None
self.loaded = False
self.load_error = None
def load(self):
try:
self.tokenizer = AutoTokenizer.from_pretrained(str(EMBEDDING_MODEL_PATH), local_files_only=True)
self.model = AutoModel.from_pretrained(str(EMBEDDING_MODEL_PATH), local_files_only=True, low_cpu_mem_usage=True)
self.model.eval()
self.loaded = True
logger.info("Embedding model loaded successfully.")
except Exception as e:
self.loaded = False
self.load_error = str(e)
logger.error(f"Embedding model loading failed: {e}")
def embed(self, text: str) -> list[float]:
if not self.loaded:
raise RuntimeError(self.load_error or "Embedding model not loaded")
inputs = self.tokenizer([text], padding=True, truncation=True, max_length=256, return_tensors="pt")
with torch.no_grad():
output = self.model(**inputs)
pooled = _mean_pooling(output, inputs["attention_mask"])
normalized = F.normalize(pooled, p=2, dim=1)
return normalized[0].tolist()
embedding_service = EmbeddingService()
|