Create core_model.py

core_model.py

@@ -0,0 +1,28 @@
+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]]