Shashiguduri/github-code-explainer

1"""2backend/services/embeddings.py3--------------------------------4Generates dense vector embeddings using BAAI/bge-small-en5from the sentence-transformers library.6 7Model: BAAI/bge-small-en8  - 384 dimensions9  - ~130 MB download, cached locally on first run10  - Free and runs entirely locally (no API calls)11  - Top-ranked on MTEB leaderboard for retrieval tasks12  - BGE models require a special query prefix for best results13"""14 15import logging16import numpy as np17from langchain_core.documents import Document18from sentence_transformers import SentenceTransformer19 20logger = logging.getLogger(__name__)21 22MODEL_NAME = "BAAI/bge-small-en"23 24# BGE models achieve best retrieval performance when the query25# is prefixed with this instruction string.26# Document embeddings do NOT use the prefix.27BGE_QUERY_PREFIX = "Represent this sentence for searching relevant passages: "28 29 30class CodeEmbedder:31    """32    Wraps sentence-transformers for generating L2-normalised embeddings.33    """34 35    def __init__(self, model_name: str = MODEL_NAME, device: str = "cpu"):36        logger.info("Loading embedding model: %s …", model_name)37        self.model      = SentenceTransformer(model_name, device=device)38        self.model_name = model_name39        self.dim        = self.model.get_sentence_embedding_dimension()40        logger.info("Model ready. Dim: %d", self.dim)41 42    def embed_documents(43        self,44        documents: list[Document],45        batch_size: int = 128,46        show_progress: bool = False,47    ) -> np.ndarray:48        """49        Encode a list of Documents into a float32 matrix.50 51        Returns:52            np.ndarray of shape (n, dim), dtype=float32, L2-normalised.53        """54        if not documents:55            return np.empty((0, self.dim), dtype=np.float32)56 57        texts = [doc.page_content for doc in documents]58        logger.info("Embedding %d chunks …", len(texts))59 60        vecs = self.model.encode(61            texts,62            batch_size=batch_size,63            show_progress_bar=show_progress,64            convert_to_numpy=True,65            normalize_embeddings=True,  # L2-normalise for cosine via dot product66        )67        return vecs.astype(np.float32)68 69    def embed_query(self, query: str) -> np.ndarray:70        """71        Encode a user question into a 1-D float32 vector.72        BGE models use a query prefix for better retrieval quality.73        """74        prefixed = BGE_QUERY_PREFIX + query75        vec = self.model.encode(76            prefixed,77            convert_to_numpy=True,78            normalize_embeddings=True,79        )80        return vec.astype(np.float32)81 82    @property83    def embedding_dim(self) -> int:84        return self.dim85