src/embeddings.py

import torch
import numpy as np
from typing import List
from sentence_transformers import SentenceTransformer
from langchain_core.embeddings import Embeddings

from config import EMBEDDING_MODEL_NAME


class TurkishEmbedder(Embeddings):
    """Embedding wrapper for E5 multilingual models."""
    
    def __init__(self, model_name: str = EMBEDDING_MODEL_NAME):
        self.model_name = model_name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self._model = None
    
    @property
    def model(self) -> SentenceTransformer:
        """Lazy load the model on first access."""
        if self._model is None:
            self._model = SentenceTransformer(self.model_name, device=self.device)
        return self._model
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed documents with 'passage:' prefix for E5 models."""
        if not texts:
            return []
        
        prefixed = [f"passage: {t}" for t in texts]
        embeddings = self.model.encode(
            prefixed,
            convert_to_numpy=True,
            show_progress_bar=len(texts) > 10,
            normalize_embeddings=True
        )
        return embeddings.tolist()
    
    def embed_query(self, text: str) -> List[float]:
        """Embed query with 'query:' prefix for E5 models."""
        prefixed = f"query: {text}"
        embedding = self.model.encode(
            prefixed,
            convert_to_numpy=True,
            normalize_embeddings=True
        )
        return embedding.tolist()
    
    def embed_passages(self, texts: List[str]) -> np.ndarray:
        """Convenience method returning numpy array."""
        return np.array(self.embed_documents(texts))
    
    @property
    def embedding_dimension(self) -> int:
        """Return the model's embedding dimension."""
        return self.model.get_sentence_embedding_dimension()


_embedder_instance = None


def get_embedder() -> TurkishEmbedder:
    """Return singleton embedder instance."""
    global _embedder_instance
    if _embedder_instance is None:
        _embedder_instance = TurkishEmbedder()
    return _embedder_instance


def reset_embedder():
    """Reset singleton for testing."""
    global _embedder_instance
    _embedder_instance = None


if __name__ == "__main__":
    embedder = get_embedder()
    
    docs = ["Bu bir test belgesidir.", "Atlas ERP sistemi hakkinda bilgi."]
    doc_embeddings = embedder.embed_documents(docs)
    print(f"Document embeddings shape: {len(doc_embeddings)}x{len(doc_embeddings[0])}")
    
    query = "Atlas nedir?"
    query_embedding = embedder.embed_query(query)
    print(f"Query embedding dimension: {len(query_embedding)}")
    print(f"Model dimension: {embedder.embedding_dimension}")