pipeline/embedder.py

"""
DocMind — Dense Embedding with BGE-M3

Loads the BAAI/bge-m3 sentence-transformer model and provides
batch encoding for chunks and single encoding for queries.
"""

import logging
from typing import List, Optional, Callable

import numpy as np
import streamlit as st

logger = logging.getLogger(__name__)


@st.cache_resource(show_spinner="Loading embedding model...")
def load_bge_model():
    """Load embedding model once and cache across Streamlit sessions."""
    from sentence_transformers import SentenceTransformer
    from config import CONFIG

    logger.info("Loading embedding model: %s", CONFIG.embedding.model_name)
    model = SentenceTransformer(CONFIG.embedding.model_name)
    logger.info("%s loaded successfully (dim=%d)", CONFIG.embedding.model_name, model.get_sentence_embedding_dimension())
    return model


def embed_chunks(
    model,
    texts: List[str],
    batch_size: int = 32,
    progress_callback: Optional[Callable[[float], None]] = None,
) -> np.ndarray:
    """
    Encode a list of chunk texts into dense vectors.

    Returns:
        np.ndarray of shape (len(texts), vector_dim)
    """
    all_embeddings = []
    total_batches = (len(texts) + batch_size - 1) // batch_size

    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        embeddings = model.encode(
            batch,
            show_progress_bar=False,
            normalize_embeddings=True,
        )
        all_embeddings.append(embeddings)

        if progress_callback:
            batch_idx = i // batch_size + 1
            progress_callback(batch_idx / total_batches)

    result = np.vstack(all_embeddings)
    logger.info("Embedded %d texts → shape %s", len(texts), result.shape)
    return result


def embed_query(model, query: str) -> np.ndarray:
    """Encode a single query string into a dense vector."""
    vector = model.encode(
        [query],
        show_progress_bar=False,
        normalize_embeddings=True,
    )
    return vector[0]