move cross encoder to embedder

retrieval.py

@@ -1,5 +1,4 @@
 import re
-import time
 import hashlib
 import pickle
 from pathlib import Path
@@ -8,24 +7,23 @@ import numpy as np
 import torch
 import pandas as pd
 from rank_bm25 import BM25Okapi
+from sentence_transformers import SentenceTransformer
 import warnings
 warnings.filterwarnings('ignore')
 
 from _1_get_documents import load_and_process_data
-from _2_splitting import years_overlap, parse_metadata_from_document
+from _2_splitting import parse_metadata_from_document
 from lemmatizer import RussianLemmatizer
-# from _3_chunking import RussianEmbedder
 
-from sentence_transformers import CrossEncoder
 
-# Модель будет загружена автоматически
-# model = CrossEncoder('DiTy/cross-encoder-russian-msmarco', max_length=512)
+def normalize_array(arr):
+    min_val = np.min(arr)
+    max_val = np.max(arr)
+    return (arr - min_val) / (max_val - min_val)
 
 
 class Retrieval:
-    """
-    RAG (Retrieval-Augmented Generation) система на русском языке.
-    
+    """ 
     Структура хранения данных:
     ============================
     
@@ -55,15 +53,9 @@ class Retrieval:
        Комбинирует оба датафрейма через JOIN по paragraph_id.
        Содержит все колонки обоих датафреймов.
        Используется для поиска и фильтрации.
-    
-    Ключевые преимущества:
-    - Избегаем дублирования метаданных параграфов
-    - Легко фильтровать по году, summary, документу
-    - Оптимизировано для работы с 5000+ чанками
-    - Простой merge для получения полной информации
     """
     
-    def __init__(self, use_gpu: bool = False, load_json: bool = True, use_cache: bool = True):
+    def __init__(self, use_gpu: bool = False, use_cache: bool = True):
         print("Инициализация RAG системы...")
         self.device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
         self.use_cache = use_cache
@@ -91,9 +83,10 @@ class Retrieval:
         self.chunks_df['lemmatized_text'] = self._lemmatize_with_cache(self.chunks_df['text'].tolist())
         
         # Инициализируем CrossEncoder
-        print("3. Загрузка CrossEncoder модели...")
-        self.cross_encoder = CrossEncoder('DiTy/cross-encoder-russian-msmarco')
-        
+        # self.cross_encoder = CrossEncoder('DiTy/cross-encoder-russian-msmarco')
+        self.embedder = SentenceTransformer('cointegrated/rubert-tiny2')
+        self.embeddings_of_summary = self.embedder.encode(self.paragraphs_df['summary'].tolist(), convert_to_tensor=True)
+
         print("✅ RAG система готова к использованию")
 
     def _process_documents_with_dates(self):
@@ -331,46 +324,46 @@ class Retrieval:
             (merged['end_year'] >= year_range[0])
         ]
     
-    def rerank_search(self, query: str) -> list[dict]:
-        """Ранжирует все чанки используя CrossEncoder модель.
+    # def rerank_search(self, query: str) -> list[dict]:
+    #     """Ранжирует все чанки используя CrossEncoder модель.
         
-        Args:
-            query: Текст запроса
+    #     Args:
+    #         query: Текст запроса
             
-        Returns:
-            list: Отсортированный список результатов с scores
-        """
-        pairs = [[query, text] for text in self.chunks_df['text'].tolist()]
-        scores = self.cross_encoder.predict(pairs)
-        breakpoint()
-        
-        # Добавляем scores в датафрейм и сортируем
-        results = self.chunks_df.copy()
-        results['score'] = scores
-        return results.sort_values('score', ascending=False).to_dict('records')
+    #     Returns:
+    #         list: Отсортированный список результатов с scores
+    #     """
+    #     pairs = [[query, text] for text in self.paragraphs_df['summary'].tolist()]
+    #     scores = self.cross_encoder.predict(pairs)
+    #     sorted_scores = dict(sorted(scores.items(), key=lambda item: item[0]))
 
-    def semantic_search(self, query: str) -> list:    
+    
+    def semantic_search(self, query: str) -> torch.Tensor:    
         # 1. Семантический поиск
-        query_embedding = self.embedder.encode_query(query)
-        semantic_scores = torch.nn.functional.cosine_similarity(self.embeddings, query_embedding, eps=1e-8).cpu()
+        query_embedding = torch.tensor(self.embedder.encode_query(query))
+        semantic_scores = torch.nn.functional.cosine_similarity(self.embeddings_of_summary, query_embedding, eps=1e-8).cpu()
         return semantic_scores
     
-    def bm25_search(self, query: str) -> list:
+    def bm25_search(self, query: str) -> np.ndarray:
         """BM25 поиск, используя лемматизированные чанки.
         
         Args:
             query: Текст запроса
             
         Returns:
-            list: Скоры для каждого чанка
+            np.ndarray: Скоры для каждого абзаца (не предложения!)
         """
         bm25 = BM25Okapi(self.chunks_df['lemmatized_text'].tolist())
         tokenized_query = self.lemmatizer.tokenize_text(query)
-        return bm25.get_scores(tokenized_query)
+        sentences_scores = bm25.get_scores(tokenized_query)
+        df = self.chunks_df['paragraph_id'].to_frame().copy()
+        df['score'] = sentences_scores
+        paragraph_scores = df.groupby('paragraph_id')['score'].max().reindex(self.paragraphs_df['paragraph_id']).fillna(0)
+        return paragraph_scores
     
     # ============ Для тестирования cross-encoder ============ 
     def test_query_with_cross_encoder(self, query: str,
-                                      target_summary: str):
+                                      target_summary: str, weight_bm25: float = 0.5, weight_semantic: float = 0.5) -> None:
         """ Тестирует запрос с cross-encoder и выводит результаты.
         
         Args:
@@ -379,81 +372,9 @@ class Retrieval:
         """
         print(f"{'='*90}")
         print(f"   ✓ Target summary: '{target_summary}'\n")
-
-        # Получаем объединённый датафрейм
-        merged_df = self.get_merged_data()
-        
-        # ================================================================
-        # 1. BM25 ПОИСК
-        # ================================================================        
-        print(f"   📊 BM25 ЛЕКСИЧЕСКИЙ ПОИСК:")
-        
-        # Инициализируем BM25
-        bm25_scores = self.bm25_search(query)
-        
-        # Добавляем scores в помощный датафрейм
-        search_df = merged_df.copy()
-        search_df['bm25_score'] = bm25_scores
-        
-        # Получаем топ-30 по BM25
-        top_bm25 = search_df.nlargest(30, 'bm25_score')
-        
-        print(f"      Топ-10 чанков, их summary-ы:")
-        
-        # Собираем уникальные summary из BM25 результатов
-        bm25_summaries = top_bm25['summary'].unique()
-        summary_scores_bm25 = dict(
-            top_bm25.groupby('summary')['bm25_score'].first()
-        )
-        
-        for rank, summary in enumerate(bm25_summaries[:10], 1):
-            score = summary_scores_bm25[summary]
-            print(f"      {rank:2}. BM25={score:6.2f} [{summary[:50]:50}]")
-        
-        print(f"      → Уникальных summary найдено: {len(bm25_summaries)}")
-        print(f"      → Целевой summary в результатах: {'✓ ДА' if target_summary in bm25_summaries else '✗ НЕТ'}")
-        
-        # ================================================================
-        # 2. КРОСС-ЭНКОДЕР РАНЖИРОВАНИЕ
-        # ================================================================
-        print(f"\n   🏆 КРОСС-ЭНКОДЕР РАНЖИРОВАНИЕ:")
-        
-        # Собираем ВСЕ уникальные summary
-        all_unique_summaries = merged_df['summary'].unique().tolist()
-        assert target_summary in all_unique_summaries, breakpoint()
-        
-        cross_encoder_start = time.time()
-        
-        # Подготавливаем пары query-summary
-        pairs = [[query, summary] for summary in all_unique_summaries]
-        
-        # Ранжируем через кросс-энкодер
-        cross_scores = self.cross_encoder.predict(pairs)
-        cross_encoder_time = time.time() - cross_encoder_start
-        
-        # Сортируем результаты
-        ranked_indices = sorted(
-            range(len(cross_scores)),
-            key=lambda i: cross_scores[i],
-            reverse=True
-        )
-        
-        print(f"      (время: {cross_encoder_time:.3f} сек)")
-        print(f"      Top-5 summary (из {len(all_unique_summaries)} всего):")
-        
-        cross_target_rank = None
-        
-        for rank, idx in enumerate(ranked_indices[:5], 1):
-            summary = all_unique_summaries[idx]
-            score = cross_scores[idx]
-            
-            is_target = summary == target_summary
-            mark = "⭐ TARGET ⭐" if is_target else " " * 13
-            
-            print(f"      {mark} {rank}. Cross={score:7.4f} [{summary[:50]:50}]")
-            
-            if is_target:
-                cross_target_rank = rank
-        
-        if not cross_target_rank:
-            print(f"      ❌ Целевой summary НЕ в топ-5")
+       
+        bm25_scores = self.bm25_search(query)        
+        semantic_scores = self.semantic_search(query).numpy()
+        bm25_scores = normalize_array(bm25_scores)
+        semantic_scores = normalize_array(semantic_scores)
+        return weight_semantic * semantic_scores + weight_bm25 * bm25_scores

test_cross_encoder.py

@@ -36,29 +36,15 @@ test_cases = [
 
 
 def test_cross_encoder_vs_bm25():
-    """Тестирует кросс-энкодер vs BM25 на всех документах."""
-    print("=" * 90)
-    print("СРАВНЕНИЕ: КРОСС-ЭНКОДЕР vs BM25 ЛЕММАТИЗИРОВАННЫЙ ПОИСК")
-    print("=" * 90)
-    
+    """Тестирует кросс-энкодер vs BM25 на всех документах."""   
     # Создаем объект Retrieval (загружает корпус автоматически)
     retrieval = Retrieval(use_gpu=False)
     
-    # Тестируем каждый тестовый случай
-    print("=" * 90)
-    print("ТЕСТИРОВАНИЕ ОТДЕЛЬНЫХ ЗАПРОСОВ")
-    print("=" * 90)
-    
-    for test_num, test_case in enumerate(test_cases, 1):
+    for test_case in test_cases:
         retrieval.test_query_with_cross_encoder(
             query=test_case.query,
             target_summary=test_case.good_answer,
-            test_num=test_num
         )
-    
-    print("\n" + "=" * 90)
-    print(f"✅ Тестирование завершено")
-    print("=" * 90)
 
 
 if __name__ == "__main__":