core_model.py

import torch
from sentence_transformers import SentenceTransformer, util


class CoreModel:
    def __init__(self, model_name="all-MiniLM-L6-v2"):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = SentenceTransformer(model_name).to(self.device)

    def compute_similarity(self, input_text: str, candidates: list[str]) -> list[tuple[str, float]]:
        """
        Расчитывает сходство входного текста с каждым из кандидатов и возвращает отсортированный список.
        """
        input_emb = self.model.encode(input_text, convert_to_tensor=True).to(self.device)
        candidates_emb = self.model.encode(candidates, convert_to_tensor=True).to(self.device)

        similarities = util.cos_sim(input_emb, candidates_emb)[0]
        scored_candidates = [(candidates[i], float(similarities[i])) for i in range(len(candidates))]
        scored_candidates.sort(key=lambda x: x[1], reverse=True)

        return scored_candidates

    def extract_top_k(self, input_text: str, candidates: list[str], top_k: int = 3) -> list[str]:
        """
        Возвращает топ-K наиболее релевантных кандидатов на основе сходства.
        """
        similarities = self.compute_similarity(input_text, candidates)
        return [candidate for candidate, _ in similarities[:top_k]]