| from sentence_transformers import CrossEncoder | |
| from sentence_transformers import SentenceTransformer | |
| import faiss | |
| import json | |
| class CustomRetriever: | |
| def __init__(self, chunks_path, embeddings_path, metadata_path, top_k=50): | |
| self.model_bi = SentenceTransformer("deepvk/USER-bge-m3") | |
| self.model_cross = CrossEncoder("DiTy/cross-encoder-russian-msmarco") | |
| with open(chunks_path, "r") as f: | |
| self.chunks = json.load(f) | |
| self.index = faiss.read_index(embeddings_path) | |
| self.top_k = top_k | |
| with open(metadata_path, "r") as f: | |
| self.metadata = json.load(f) | |
| def retrieve(self, query): | |
| query_vector = self.model_bi.encode([query]) | |
| faiss.normalize_L2(query_vector) | |
| distances, indices = self.index.search(query_vector, self.top_k) | |
| possible_answers = list() | |
| for i in range(len(indices[0])): | |
| possible_answers.append(self.chunks[indices[0][i]]) | |
| s = self.model_cross.rank(query, possible_answers) | |
| context = '' | |
| for i in range(5): | |
| meta = self.metadata[str(indices[0][s[i]["corpus_id"]])] | |
| context += f"Факт {str(i + 1)}: {possible_answers[s[i]['corpus_id']]}. Источник:\nкнига - {meta['book']}\nномер статьи - {meta['article_num']}\nссылка на книгу - {meta['link']}\n" | |
| return context | |