Aigenthix/Graph_RAG

1"""Embedding generation service"""2 3from typing import List, Dict, Any4from sentence_transformers import SentenceTransformer5import logging6 7logger = logging.getLogger(__name__)8 9 10class EmbeddingService:11    """Service for generating text embeddings"""12 13    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):14        self.model_name = model_name15        self.model = SentenceTransformer(model_name)16        logger.info(f"Loaded embedding model: {model_name}")17 18    def embed_text(self, text: str) -> List[float]:19        """Generate embedding for single text"""20        embedding = self.model.encode(text, convert_to_tensor=False)21        return embedding.tolist()22 23    def embed_batch(self, texts: List[str]) -> List[List[float]]:24        """Generate embeddings for multiple texts"""25        embeddings = self.model.encode(texts, convert_to_tensor=False)26        return embeddings.tolist()27 28    def embed_chunks(29        self,30        chunks: List[Dict[str, Any]],31    ) -> List[Dict[str, Any]]:32        """Generate embeddings for document chunks"""33        texts = [chunk["text"] for chunk in chunks]34        embeddings = self.embed_batch(texts)35 36        for chunk, embedding in zip(chunks, embeddings):37            chunk["embedding"] = embedding38 39        logger.info(f"Generated embeddings for {len(chunks)} chunks")40        return chunks41 42    def get_model_dimension(self) -> int:43        """Get embedding dimension"""44        return self.model.get_sentence_embedding_dimension()45