embedding_model.py

"""
Local embedding model implementation using all-MiniLM-L6-v2.
"""
from typing import List
from langchain_huggingface import HuggingFaceEmbeddings

class EmbeddingModel:
    """
    Local embedding model using all-MiniLM-L6-v2 from HuggingFace.
    """
    
    def __init__(self,
                 api_key: str = None,
                 model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
        """
        Initialize the embedding model.
        
        Args:
            api_key: Not used for local models, kept for compatibility
            model_name: The embedding model to use (default: "sentence-transformers/all-MiniLM-L6-v2")
        """
        self.model_name = model_name
        
        # Initialize the local HuggingFace embedding model
        self.model = HuggingFaceEmbeddings(
            model_name=self.model_name,
            model_kwargs={'device': 'cpu'},  # Use CPU for compatibility
            encode_kwargs={'normalize_embeddings': True}  # Normalize for better similarity
        )
    
    async def embed(self, text: str) -> List[float]:
        """
        Generate embedding for a single text.
        
        Args:
            text: The text to embed
            
        Returns:
            A float array representing the embedding
        """
        # Run in thread pool to avoid blocking the event loop
        import asyncio
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.model.embed_query, text)
    
    async def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """
        Generate embeddings for multiple texts using local model.
        
        Args:
            texts: List of texts to embed
            
        Returns:
            List of float arrays representing the embeddings
        """
        if not texts:
            return []
        
        # Run in thread pool to avoid blocking the event loop
        import asyncio
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.model.embed_documents, texts)