src/models/similarity_calculator.py

"""Semantic similarity calculation functionality"""

import torch
from sentence_transformers import SentenceTransformer


class SimilarityCalculator:
    """Handles semantic similarity calculations using sentence transformers"""
    
    def __init__(self):
        self.similarity_model = None
        
    def load_model(self, model_name='all-MiniLM-L6-v2'):
        """Load the sentence transformer model for similarity calculations"""
        self.similarity_model = SentenceTransformer(model_name)
    
    def compute_cosine_distance(self, text1, text2):
        """Compute cosine distance between two texts"""
        if not self.similarity_model:
            raise RuntimeError("Similarity model not loaded. Call load_model() first.")
        
        embeddings = self.similarity_model.encode([text1, text2])
        similarity = torch.cosine_similarity(
            torch.tensor(embeddings[0]).unsqueeze(0),
            torch.tensor(embeddings[1]).unsqueeze(0)
        )
        return 1.0 - similarity.item()