models/embedder.py

# models/embedder.py

import re
from typing import Any, Dict, List

import torch
from FlagEmbedding import BGEM3FlagModel
from pydantic import BaseModel

from core.exceptions import ModelLoadError
from core.logger import setup_logger

logger = setup_logger("embedder")

# Data structure for return (Type Hinting)
class EmbedderResult(BaseModel):
    dense_vector: List[float]
    sparse_indices: List[int]
    sparse_values: List[float]

class TextEmbedder:
    """
    Converts the input text into Dense Vectors and Sparse Vectors (Lexical Weights) using the BGE-M3 model.
    """
    def __init__(self, model_name: str = "BAAI/bge-m3", use_fp16: bool = False,
                 revision: str | None = None):
        self.model_name = model_name
        self.device = self._get_device()
        precision = "fp16" if use_fp16 and self.device.startswith("cuda") else "fp32"
        self.compatibility_id = model_name + "@" + (revision or "unpinned") + ":" + precision + ":unit-cls-v1"
        
        try:
            logger.info(f"⏳ Loading Embedder Model: {self.model_name} on {self.device}")
            model_location = self.model_name
            if revision is not None:
                if not re.fullmatch(r"[0-9a-f]{40}", revision):
                    raise ValueError("Embedding revision must be an exact commit hash")
                from huggingface_hub import snapshot_download
                # FlagEmbedding does not forward a revision keyword to all model
                # loaders. Resolve the exact snapshot before constructing it.
                model_location = snapshot_download(repo_id=model_name, revision=revision,
                    allow_patterns=["*.json", "*.model", "*.txt", "*.bin", "*.pt",
                                    "*.safetensors", "1_Pooling/*"])
            self.model = BGEM3FlagModel(
                model_location,
                use_fp16=(use_fp16 and self.device.startswith("cuda")),
                trust_remote_code=False,
                normalize_embeddings=True,
                pooling_method="cls"
            )
            self._warmup()
            logger.info("✅ Embedder Model loaded successfully.")
        except Exception as e:
            logger.critical(f"❌ Failed to load Embedder Model: {e}", exc_info=True)
            raise ModelLoadError(f"Embedder initialization failed: {e}")

    def _get_device(self) -> str:
        if torch.cuda.is_available():
            return "cuda"
        elif torch.backends.mps.is_available():
            return "mps" # Apple Silicon
        return "cpu"
    
    def _warmup(self):
        logger.info("Warming up embedder model with a dummy input.")
        self.encode_query("Hello world")

    def encode_query(self, text: str) -> EmbedderResult:
        """
        Converts a single query text into Qdrant hybrid search format.
        """
        try:
            # 1. model inference to get dense vector and sparse lexical weights
            output = self.model.encode(
                text, 
                return_dense=True, 
                return_sparse=True, 
                return_colbert_vecs=False
            )
            
            dense_vec = output['dense_vecs'].tolist()
            lexical_weights: Dict[str, float] = output['lexical_weights']
            
            # 2. Sparse Vector Transformation and Duplicate Index Removal
            # Prevent duplicates by storing in the format {token_id: max_weight}
            unique_sparse_data = {}
            
            # FlagEmbedding returns vocabulary IDs as strings; converting these
            # through convert_tokens_to_ids collapses numeric IDs to unknown IDs.
            for token_id_string, weight in lexical_weights.items():
                token_id = int(token_id_string)
                unique_sparse_data[token_id] = max(
                    unique_sparse_data.get(token_id, 0.0), float(weight)
                )
                    
            sparse_indices = list(unique_sparse_data.keys())
            sparse_values = list(unique_sparse_data.values())

            return EmbedderResult(
                dense_vector=dense_vec,
                sparse_indices=sparse_indices,
                sparse_values=sparse_values
            )
            
        except Exception as e:
            logger.error("Embedding generation failed")
            raise RuntimeError("Embedding generation failed") from e

    # options for batch encoding if needed in the future
    def encode_documents(self, texts: List[str], batch_size: int = 12,
                         max_length: int = 8192) -> Dict[str, Any]:
        if max_length < 8192:
            for text in texts:
                if len(self.model.tokenizer(text, truncation=False)["input_ids"]) > max_length:
                    raise ValueError("Embedding segment exceeds model token limit")
        return self.model.encode(
            texts, 
            batch_size=batch_size, 
            max_length=max_length,
            return_dense=True, 
            return_sparse=True, 
            return_colbert_vecs=False
        )