| 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]] |
|
|
