pipeline/hf_embedding.py

import logging
from typing import List, Any, Optional, Tuple
import numpy as np
from sentence_transformers import SentenceTransformer

logger = logging.getLogger(__name__)

# Cache for loaded models
_model_cache = {}

def get_model(model_id: str) -> Tuple[Optional[SentenceTransformer], Optional[str]]:
    """
    Loads a SentenceTransformer model from the Hugging Face Hub.

    Args:
        model_id (str): The identifier for the model to load (e.g., 'sentence-transformers/LaBSE').

    Returns:
        Tuple[Optional[SentenceTransformer], Optional[str]]: A tuple containing the loaded model and its type ('sentence-transformer'),
                                                              or (None, None) if loading fails.
    """
    if model_id in _model_cache:
        logger.info(f"Returning cached model: {model_id}")
        return _model_cache[model_id], "sentence-transformer"

    logger.info(f"Loading SentenceTransformer model: {model_id}")
    try:
        model = SentenceTransformer(model_id)
        _model_cache[model_id] = model
        logger.info(f"Model '{model_id}' loaded successfully.")
        return model, "sentence-transformer"
    except Exception as e:
        logger.error(f"Failed to load SentenceTransformer model '{model_id}': {e}", exc_info=True)
        return None, None

def generate_embeddings(
    texts: List[str], 
    model: SentenceTransformer, 
    batch_size: int = 32, 
    show_progress_bar: bool = False
) -> np.ndarray:
    """
    Generates embeddings for a list of texts using a SentenceTransformer model.

    Args:
        texts (list[str]): A list of texts to embed.
        model (SentenceTransformer): The loaded SentenceTransformer model.
        batch_size (int): The batch size for encoding.
        show_progress_bar (bool): Whether to display a progress bar.

    Returns:
        np.ndarray: A numpy array containing the embeddings. Returns an empty array of the correct shape on failure.
    """
    if not texts or not isinstance(model, SentenceTransformer):
        logger.warning("Invalid input for generating embeddings. Returning empty array.")
        # Return a correctly shaped empty array
        embedding_dim = model.get_sentence_embedding_dimension() if isinstance(model, SentenceTransformer) else 768 # Fallback
        return np.zeros((len(texts), embedding_dim))

    logger.info(f"Generating embeddings for {len(texts)} texts with {type(model).__name__}...")
    try:
        embeddings = model.encode(
            texts, 
            batch_size=batch_size,
            convert_to_numpy=True, 
            show_progress_bar=show_progress_bar
        )
        logger.info(f"Embeddings generated with shape: {embeddings.shape}")
        return embeddings
    except Exception as e:
        logger.error(f"An unexpected error occurred during embedding generation: {e}", exc_info=True)
        embedding_dim = model.get_sentence_embedding_dimension()
        return np.zeros((len(texts), embedding_dim))