src/embedding/embedding_service.py

"""Embedding service: lazy-loading sentence-transformers wrapper."""

import logging
import os
from typing import Dict, List, Optional, Tuple

import numpy as np
import onnxruntime as ort
from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer, PreTrainedTokenizer

from src.utils.memory_utils import log_memory_checkpoint, memory_monitor


def mean_pooling(model_output, attention_mask: np.ndarray) -> np.ndarray:
    """Mean Pooling - Take attention mask into account for correct averaging."""
    token_embeddings = model_output.last_hidden_state
    input_mask_expanded = (
        np.expand_dims(attention_mask, axis=-1).repeat(token_embeddings.shape[-1], axis=-1).astype(float)
    )
    sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
    sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
    return sum_embeddings / sum_mask


class EmbeddingService:
    """HuggingFace sentence-transformers wrapper for generating embeddings.

    Uses lazy loading and a class-level cache to avoid repeated expensive model
    loads and to minimize memory footprint at startup.

    This version is optimized to use a quantized ONNX model for lower memory
    footprint.
    """

    _model_cache: Dict[str, Tuple[ORTModelForFeatureExtraction, PreTrainedTokenizer]] = {}
    _quantized_model_name = "optimum/all-MiniLM-L6-v2"

    def __init__(
        self,
        model_name: Optional[str] = None,
        device: Optional[str] = None,
        batch_size: Optional[int] = None,
    ):
        # Import config values as defaults
        from src.config import (
            EMBEDDING_BATCH_SIZE,
            EMBEDDING_DEVICE,
            EMBEDDING_MODEL_NAME,
        )

        # The original model name is kept for reference. Use quantized model only
        # when explicitly enabled via configuration (to avoid breaking tests).
        self.original_model_name = model_name or EMBEDDING_MODEL_NAME
        from src.config import EMBEDDING_USE_QUANTIZED

        if EMBEDDING_USE_QUANTIZED:
            self.model_name = self._quantized_model_name
        else:
            # Keep the model name as originally requested for compatibility
            self.model_name = self.original_model_name
        self.device = device or EMBEDDING_DEVICE or "cpu"
        self.batch_size = batch_size or EMBEDDING_BATCH_SIZE
        # Max tokens (sequence length) to bound memory; configurable via env
        # EMBEDDING_MAX_TOKENS (default 512)
        try:
            self.max_tokens = int(os.getenv("EMBEDDING_MAX_TOKENS", "512"))
        except ValueError:
            self.max_tokens = 512

        # Lazy loading - don't load model at initialization
        self.model: Optional[ORTModelForFeatureExtraction] = None
        self.tokenizer: Optional[PreTrainedTokenizer] = None

        logging.info(
            "Initialized EmbeddingService: model=%s base=%s device=%s max_tokens=%s",
            self.model_name,
            self.original_model_name,
            self.device,
            getattr(self, "max_tokens", "unset"),
        )

    def _ensure_model_loaded(
        self,
    ) -> Tuple[ORTModelForFeatureExtraction, PreTrainedTokenizer]:
        """Ensure the quantized ONNX model and tokenizer are loaded."""
        if self.model is None or self.tokenizer is None:
            import gc

            gc.collect()

            cache_key = f"{self.model_name}_{self.device}"

            if cache_key not in self._model_cache:
                log_memory_checkpoint("before_model_load")
                logging.info(
                    "Loading quantized model '%s' and tokenizer...",
                    self.model_name,
                )
                # Use the original model's tokenizer
                tokenizer = AutoTokenizer.from_pretrained(self.original_model_name)
                # Load the quantized model from Optimum Hugging Face Hub.
                # Some model repos contain multiple ONNX export files; we select a default explicitly.
                provider = "CPUExecutionProvider" if self.device == "cpu" else "CUDAExecutionProvider"
                file_name = os.getenv("EMBEDDING_ONNX_FILE", "model.onnx")
                local_dir = os.getenv("EMBEDDING_ONNX_LOCAL_DIR")
                if local_dir and os.path.isdir(local_dir):
                    # Attempt to load from a local exported directory first.
                    try:
                        logging.info(
                            "Attempting local ONNX load from %s (file=%s)",
                            local_dir,
                            file_name,
                        )
                        model = ORTModelForFeatureExtraction.from_pretrained(
                            local_dir,
                            provider=provider,
                            file_name=file_name,
                        )
                        logging.info("Loaded ONNX model from local directory '%s'", local_dir)
                    except Exception as e:
                        logging.warning(
                            "Local ONNX load failed (%s); " "falling back to hub repo '%s'",
                            e,
                            self.model_name,
                        )
                        local_dir = None  # disable local path for subsequent attempts
                if not local_dir:
                    # Configure ONNX Runtime threading for constrained CPU
                    intra = int(os.getenv("ORT_INTRA_OP_NUM_THREADS", "1"))
                    inter = int(os.getenv("ORT_INTER_OP_NUM_THREADS", "1"))
                    so = ort.SessionOptions()
                    so.intra_op_num_threads = intra
                    so.inter_op_num_threads = inter
                    try:
                        model = ORTModelForFeatureExtraction.from_pretrained(
                            self.model_name,
                            provider=provider,
                            file_name=file_name,
                            session_options=so,
                        )
                        logging.info(
                            "Loaded ONNX model file '%s' (intra=%d, inter=%d)",
                            file_name,
                            intra,
                            inter,
                        )
                    except Exception as e:
                        logging.warning(
                            "Explicit ONNX file '%s' failed (%s); " "retrying with auto-selection.",
                            file_name,
                            e,
                        )
                        # The key change: we now pass the file_name to the fallback as well
                        model = ORTModelForFeatureExtraction.from_pretrained(
                            self.model_name,
                            provider=provider,
                            file_name=file_name,  # Added this line
                            session_options=so,
                        )
                        logging.info(
                            "Loaded ONNX model using auto-selection fallback " "(intra=%d, inter=%d)",
                            intra,
                            inter,
                        )
                self._model_cache[cache_key] = (model, tokenizer)
                logging.info("Quantized model and tokenizer loaded successfully")
                log_memory_checkpoint("after_model_load")
            else:
                logging.info("Using cached quantized model '%s'", self.model_name)

            self.model, self.tokenizer = self._model_cache[cache_key]

        return self.model, self.tokenizer

    @memory_monitor
    def embed_text(self, text: str) -> List[float]:
        """Generate embedding for a single text."""
        embeddings = self.embed_texts([text])
        return embeddings[0]

    @memory_monitor
    def embed_texts(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for multiple texts in batches using ONNX model."""
        if not texts:
            return []

        try:
            model, tokenizer = self._ensure_model_loaded()

            log_memory_checkpoint("before_batch_embedding")

            processed_texts: List[str] = [t if t.strip() else " " for t in texts]

            all_embeddings: List[List[float]] = []
            for i in range(0, len(processed_texts), self.batch_size):
                batch_texts = processed_texts[i : i + self.batch_size]
                log_memory_checkpoint(f"batch_start_{i}//{self.batch_size}")

                # Tokenize sentences
                encoded_input = tokenizer(
                    batch_texts,
                    padding=True,
                    truncation=True,
                    max_length=self.max_tokens,
                    return_tensors="np",
                )

                # Compute token embeddings
                model_output = model(**encoded_input)

                # Perform pooling
                sentence_embeddings = mean_pooling(model_output, encoded_input["attention_mask"])

                # Normalize embeddings (L2) using pure NumPy to avoid torch dependency
                norms = np.linalg.norm(sentence_embeddings, axis=1, keepdims=True)
                norms = np.clip(norms, 1e-12, None)
                batch_embeddings = sentence_embeddings / norms

                log_memory_checkpoint(f"batch_end_{i}//{self.batch_size}")

                for emb in batch_embeddings:
                    all_embeddings.append(emb.tolist())

                import gc

                del batch_embeddings
                del batch_texts
                del encoded_input
                del model_output
                gc.collect()

            if os.getenv("LOG_DETAIL", "verbose") == "verbose":
                logging.info("Generated embeddings for %d texts", len(texts))
            return all_embeddings
        except Exception as e:
            logging.error("Failed to generate embeddings for texts: %s", e)
            raise

    def get_embedding_dimension(self) -> int:
        """Get the dimension of embeddings produced by this model."""
        try:
            model, _ = self._ensure_model_loaded()
            # The dimension can be found in the model's config
            return int(model.config.hidden_size)
        except Exception:
            logging.debug("Failed to get embedding dimension; returning 0")
            return 0

    def encode_batch(self, texts: List[str]) -> List[List[float]]:
        """Convenience wrapper that returns embeddings for a list of texts."""
        return self.embed_texts(texts)

    def similarity(self, text1: str, text2: str) -> float:
        """Cosine similarity between embeddings of two texts."""
        try:
            embeddings = self.embed_texts([text1, text2])
            embed1 = np.array(embeddings[0])
            embed2 = np.array(embeddings[1])
            similarity = np.dot(embed1, embed2) / (np.linalg.norm(embed1) * np.linalg.norm(embed2))
            return float(similarity)
        except Exception as e:
            logging.error("Failed to calculate similarity: %s", e)
            return 0.0