Pankaj70's picture
Deploy AHCGR legal inference app
d8423e6 verified
Raw History Blame Contribute Delete
3.57 kB
"""
Embedding service — BAAI/bge-base-en-v1.5 via fastembed (local, no API key required).
This module handles query embedding for P1. The embedding vector drives
the dense search in P2.
Key design decisions:
• Uses BAAI/bge-base-en-v1.5 (768 dimensions) — local inference, no quota.
• fastembed handles model download, caching, and batching automatically.
• Normalize=True ensures cosine similarity is stable (unit vectors).
• Request-scoped caching avoids re-embedding the same query during retries.
• Full error logging — no silent failures.
"""
from __future__ import annotations
import time
import structlog
from fastembed import TextEmbedding
from ahcgr.config import get_settings
logger = structlog.get_logger("ahcgr.services.embedding")
# Module-level model instance — loaded once, shared across requests
_model: TextEmbedding | None = None
def _get_model() -> TextEmbedding:
"""Lazy-load the fastembed dense model (downloaded on first call)."""
global _model
if _model is None:
settings = get_settings()
logger.info("loading_embedding_model", model=settings.EMBEDDING_MODEL)
_model = TextEmbedding(
model_name=settings.EMBEDDING_MODEL,
)
logger.info("embedding_model_loaded", model=settings.EMBEDDING_MODEL)
return _model
class EmbeddingService:
"""
Embedding service using BAAI/bge-base-en-v1.5 via fastembed.
Runs entirely locally — no API key, no quota limits.
The model is downloaded once from HuggingFace and cached by fastembed.
Lifecycle:
Created once at app startup, shared across requests.
fastembed handles thread safety internally.
Caching:
A simple dict cache avoids re-embedding identical queries
within the same process. Useful during retry loops.
"""
def __init__(self) -> None:
settings = get_settings()
self._model_name = settings.EMBEDDING_MODEL
self._dimensions = settings.EMBEDDING_DIMENSIONS
self._cache: dict[str, list[float]] = {}
# Trigger lazy load at startup so the first request isn't slow
_get_model()
logger.info(
"embedding_service_initialized",
model=self._model_name,
dimensions=self._dimensions,
)
def embed(self, text: str) -> list[float]:
"""
Embed a text string into a 768-d unit vector.
Args:
text: The query or text to embed.
Returns:
768-dimensional float vector (L2-normalized).
Raises:
Exception: Propagated from fastembed on model errors.
"""
# Check cache first (exact match)
if text in self._cache:
logger.debug("embedding_cache_hit", text_preview=text[:60])
return self._cache[text]
t0 = time.perf_counter()
model = _get_model()
# fastembed.embed() returns a generator of numpy arrays
embeddings = list(model.embed([text]))
vector = embeddings[0].tolist()
duration_ms = (time.perf_counter() - t0) * 1000
# Cache for retry loops
self._cache[text] = vector
logger.info(
"embedding_created",
model=self._model_name,
text_preview=text[:60],
vector_dim=len(vector),
duration_ms=round(duration_ms, 2),
)
return vector
def clear_cache(self) -> None:
"""Clear the embedding cache (called between requests)."""
self._cache.clear()