src/embeddings/huggingface_embedding.py

# src/embeddings/huggingface_embedding.py
from typing import List
from sentence_transformers import SentenceTransformer

from .base_embedding import BaseEmbedding

class HuggingFaceEmbedding(BaseEmbedding):
    def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
        """
        Initialize HuggingFace embedding model
        
        Args:
            model_name (str): Name of the embedding model
        """
        self.model = SentenceTransformer(model_name)
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """
        Embed a list of documents
        
        Args:
            texts (List[str]): List of texts to embed
        
        Returns:
            List[List[float]]: List of embeddings
        """
        return self.model.encode(texts).tolist()
    
    def embed_query(self, text: str) -> List[float]:
        """
        Embed a single query
        
        Args:
            text (str): Text to embed
        
        Returns:
            List[float]: Embedding vector
        """
        return self.model.encode(text).tolist()