Spaces:
Sleeping
Sleeping
Download src/ahcgr/services/embedding.py from Pankaj70/AHCGR-Legal-Inference: direct link, hf CLI and curl.
- Browser
- Download file 3.57 kB
-
https://huggingface.co/spaces/Pankaj70/AHCGR-Legal-Inference/resolve/main/src/ahcgr/services/embedding.py
- Command line
-
hf download hf://spaces/Pankaj70/AHCGR-Legal-Inference/src/ahcgr/services/embedding.py
-
curl -L -o embedding.py https://huggingface.co/spaces/Pankaj70/AHCGR-Legal-Inference/resolve/main/src/ahcgr/services/embedding.py
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() | |