Harshavard21/FinRAG

1"""2src/embeddings/embedder.py3===========================4Embedding engine using BAAI/bge-large-en-v1.5.5 6Why BGE-large?7- Free, runs locally on your RTX 30508- Top of MTEB leaderboard for retrieval tasks (among free models)9- 768-dimensional output — good balance of quality vs. memory10- Trained with specific instruction prefix for queries vs. documents11 12BGE-specific usage:13    Documents  → embed as-is (no prefix)14    Queries    → add "Represent this sentence: " prefix15    This asymmetric approach improves retrieval accuracy significantly.16 17L2 Normalization:18    All embeddings are L2-normalized before storage.19    This makes cosine similarity = dot product → faster ANN search in Qdrant.20"""21 22from __future__ import annotations23 24import numpy as np25from typing import Union26 27from tqdm import tqdm28 29# config import first — sets USE_TF=0 before sentence_transformers loads30from config.settings import settings31from src.utils.logger import logger32 33# Now safe to import34from sentence_transformers import SentenceTransformer35 36 37class Embedder:38    """39    Singleton-pattern embedding engine.40    Loads BAAI/bge-large-en-v1.5 once and reuses it.41 42    GPU-accelerated on RTX 3050 automatically (via settings.get_device()).43    """44 45    _instance: "Embedder | None" = None46 47    def __new__(cls):48        """Enforce singleton — only load the model once."""49        if cls._instance is None:50            cls._instance = super().__new__(cls)51            cls._instance._initialized = False52        return cls._instance53 54    def __init__(self):55        if self._initialized:56            return57        self._load_model()58        self._initialized = True59 60    def _load_model(self):61        """Load the embedding model onto GPU/CPU."""62        model_name = settings.embedding_model63        device = settings.get_device()64 65        logger.info(f"Loading embedding model: {model_name} → device: {device}")66        logger.info("(First load downloads ~1.3GB — subsequent loads use cache)")67 68        self.model = SentenceTransformer(model_name, device=device)69        self.model.max_seq_length = 512    # BGE-large max input length70        self.dim = settings.embedding_dim  # 76871 72        logger.info(f"Embedding model loaded | dim={self.dim} | device={device}")73 74    # ---------------------------------------------------------------- #75    # Public API76    # ---------------------------------------------------------------- #77 78    def embed_documents(self, texts: list[str]) -> np.ndarray:79        """80        Embed a list of document chunks (corpus side).81        No instruction prefix needed for documents with BGE.82 83        Args:84            texts: List of text strings to embed85 86        Returns:87            np.ndarray of shape (len(texts), 768), L2-normalized88        """89        if not texts:90            return np.array([])91 92        logger.info(f"Embedding {len(texts)} documents in batches of {settings.embed_batch_size}...")93 94        embeddings = []95        batch_size = settings.embed_batch_size96 97        for i in tqdm(range(0, len(texts), batch_size), desc="Embedding", unit="batch"):98            batch = texts[i: i + batch_size]99            batch_emb = self.model.encode(100                batch,101                batch_size=batch_size,102                show_progress_bar=False,103                normalize_embeddings=True,   # L2 normalize in-place104                convert_to_numpy=True,105            )106            embeddings.append(batch_emb)107 108        result = np.vstack(embeddings)109        logger.info(f"Embedding complete: shape={result.shape}")110        return result111 112    def embed_query(self, query: str) -> np.ndarray:113        """114        Embed a single search query (query side).115        BGE models benefit from an instruction prefix on queries.116 117        Args:118            query: The user's question119 120        Returns:121            np.ndarray of shape (768,), L2-normalized122        """123        # BGE instruction prefix — improves retrieval accuracy124        prefixed_query = f"Represent this sentence for searching relevant passages: {query}"125 126        embedding = self.model.encode(127            prefixed_query,128            normalize_embeddings=True,129            convert_to_numpy=True,130            show_progress_bar=False,131        )132        return embedding133 134    def embed_queries(self, queries: list[str]) -> np.ndarray:135        """136        Embed multiple queries (for multi-query retrieval expansion).137 138        Args:139            queries: List of question strings140 141        Returns:142            np.ndarray of shape (len(queries), 768)143        """144        prefixed = [145            f"Represent this sentence for searching relevant passages: {q}"146            for q in queries147        ]148        embeddings = self.model.encode(149            prefixed,150            normalize_embeddings=True,151            convert_to_numpy=True,152            show_progress_bar=False,153        )154        return embeddings155 156    @property157    def embedding_dim(self) -> int:158        return self.dim159