src/embeddings.py

from sentence_transformers import SentenceTransformer
from typing import List
import numpy as np
from src.config import config

class EmbeddingService:
    def __init__(self, model_name: str = None):
        self.model_name = model_name or config.EMBEDDING_MODEL
        self.model = None
    
    def load_model(self):
        if self.model is None:
            print(f"Loading embedding model: {self.model_name}")
            self.model = SentenceTransformer(self.model_name)
            print(f"Model loaded successfully")
    
    def embed_text(self, text: str) -> List[float]:
        self.load_model()
        embedding = self.model.encode(text, convert_to_numpy=True)
        return embedding.tolist()
    
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        self.load_model()
        embeddings = self.model.encode(texts, convert_to_numpy=True, show_progress_bar=True)
        return embeddings.tolist()
    
    def get_embedding_dimension(self) -> int:
        self.load_model()
        return self.model.get_sentence_embedding_dimension()

embedding_service = EmbeddingService()