core/embeddings/text_embedding_model.py

import torch

from typing import List
from sentence_transformers import SentenceTransformer
from utils.logger import logger
from config.model_configs import TEXT_EMBEDDING_MODEL

class TextEmbeddingModel:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Loading Text Embedding Model '{TEXT_EMBEDDING_MODEL}' to device: {self.device}")
        
        self.model = SentenceTransformer(TEXT_EMBEDDING_MODEL, device=self.device)
        logger.info("Text Embedding Model loaded successfully.")
        
    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
        if not texts:
            return []
        
        embeddings = self.model.encode(texts, convert_to_numpy=True).tolist()
        logger.debug(f"Generated {len(embeddings)} embeddings for {len(texts)} texts.")
        return embeddings