notebook/embedding.py

"""embedding.py — wraps a SentenceTransformer model to turn text into vectors."""
import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List


class EmbeddingManager:
    """Handles embedding generation using a SentenceTransformer model."""

    def __init__(self, model_name: str = "BAAI/bge-base-en-v1.5"):
        self.model_name = model_name
        self.model = None
        self._load_model()

    def _load_model(self):
        try:
            print(f"Loading embedding model: {self.model_name}")
            self.model = SentenceTransformer(self.model_name)
            print(f"Model loaded. Embedding dimension: "
                  f"{self.model.get_sentence_embedding_dimension()}")
        except Exception as e:
            print(f"Error loading model {self.model_name}: {e}")
            raise

    # BAAI/bge-* retrieval is asymmetric: passages are embedded plain, but a short
    # search query gets this instruction prepended (s2p retrieval). The corpus was
    # built with generate_embeddings() (no instruction), so queries must use
    # embed_query() to land in the matching subspace.
    QUERY_INSTRUCTION = "Represent this sentence for searching relevant passages: "

    def generate_embeddings(self, texts: List[str]) -> np.ndarray:
        """Return embeddings of shape (len(texts), embedding_dim). Used for the
        CORPUS (passages) — no query instruction."""
        if not self.model:
            raise ValueError("Model not loaded")
        return self.model.encode(texts, show_progress_bar=False)

    def embed_query(self, text: str) -> List[float]:
        """Embed a single SEARCH query with the bge query instruction."""
        if not self.model:
            raise ValueError("Model not loaded")
        vec = self.model.encode([self.QUERY_INSTRUCTION + (text or "")],
                                show_progress_bar=False)[0]
        return vec.tolist()