add cross-encoder

_1_get_documents.py

@@ -13,7 +13,7 @@ def process_file(file_path):
     with open(file_path, 'r', encoding='utf-8-sig') as f:
         text = f.read()
         assert text
-    return text, str(file_path).split('.')[-1]
+    return text, str(file_path).split('.')[0]
 
 def process_folder_recursive(folder_path):
     all_messages = []

_2_splitting.py

@@ -1,6 +1,3 @@
-from typing import List, Dict
-import re
-
 # Конфиги для парсинга дат
 YEARS = {
 'O': 1918,
@@ -53,64 +50,56 @@ def _parse_date_range(date_str: str) -> tuple[int, int]:
         parts = date_str.split('-')
         start = _parse_single_year(parts[0].strip())
         end = _parse_single_year(parts[1].strip())
-        return (start, end) if start <= end else (end, start)
+        assert start <= end, f"Год начала {start} должен быть меньше или равен году конца {end}"
+        return (start, end)
     else:
         # Один год
         year = _parse_single_year(date_str)
         return (year, year)
 
 
-def parse_year_metadata(text: str) -> list[tuple[str, tuple[int, int]]]:
-    """
-    Парсит markdown текст и возвращает список (чанк_текста, (год_начала, год_конца)).
+def parse_metadata_from_document(text: str) -> list[tuple[str, tuple[int, int], str]]:
+    """Парсит markdown текст и возвращает список (чанк_текста, (год_начала, год_конца), summary).
+    
+    Формат разметки ОБЯЗАТЕЛЕН:
+    - ## Summary text - заголовок summary (двойной хэш + пробел)
+    - ### 1962-2002 - заголовок с годом (тройной хэш + пробел)
     
-    Формат разметки: ## 1962-2002 или ### 1962-2002 или просто 1962 или O или N
-    Разметка распространяется на абзац ниже и все последующие до новой разметки.
+    Правила:
+    - Каждый документ ДОЛЖЕН начинаться с "## {summary}"
+    - После summary ДОЛЖНЫ быть заголовки "### {годы}" с текстом
+    - ## распространяется на все абзацы ниже до следующего ## или конца файла
+    - ### распространяется на абзацы ниже до следующего ### или ##
+    - Текст БЕЗ предшествующего ### Не добавляется в результат
     
     Args:
         text: Полный текст документа
         
     Returns:
-        list: [(chunk_text, (start_year, end_year)), ...]
-        
-    Raises:
-        ValueError: Если документ не начинается с разметки
+        list: [(chunk_text, (start_year, end_year), summary), ...]
     """
     lines = text.split('\n')
     
-    # Проверяем, начинается ли документ с разметки (# или ## или ###)
-    if not lines or not re.match(r'^#+\s*', lines[0].strip()):
-        raise ValueError(f"Документ не начинается с разметки! Первая строка: {lines[0] if lines else 'ПУСТО'}")
-    
     result = []
+    current_summary = None
     current_year_range = None
-    current_text = []
     
     for line in lines:
-        # Проверяем, является ли строка разметкой (начинается с #, ##, ### и т.д.)
-        match = re.match(r'^#+\s+(.+?)$', line.strip())
-        if match:
-            # Сохраняем предыдущий абзац если он есть
-            if current_text and current_year_range:
-                chunk = '\n'.join(current_text).strip()
-                if chunk:
-                    result.append((chunk, current_year_range))
-            
-            # Парсим новую разметку
-            date_str = match.group(1).strip()
-            current_year_range = _parse_date_range(date_str)
-            
-            current_text = []
+        strip_line = line.strip()
+        if not strip_line:
+            continue
+        
+        if strip_line.startswith('## '):   
+            current_summary = strip_line[3:].strip()  # Пропускаем "## "
+        
+        # Проверяем, является ли строка "### " (год с пробелом после)
+        elif strip_line.startswith('### '):
+            current_year_range = _parse_date_range(strip_line[4:])
+        
         else:
-            # Если это не разметка, добавляем в текущий абзац
-            if line.strip():  # Только непустые строки
-                current_text.append(line)
-    
-    # Сохраняем последний абзац
-    if current_text and current_year_range:
-        chunk = '\n'.join(current_text).strip()
-        if chunk:
-            result.append((chunk, current_year_range))
+            # Добавляем текст только если у нас есть год
+            assert current_year_range and current_summary, breakpoint()
+            result.append((line, current_year_range, current_summary))
     
     return result
 
@@ -132,68 +121,3 @@ def years_overlap(range1: tuple[int, int], range2: tuple[int, int]) -> bool:
     start1, end1 = range1
     start2, end2 = range2
     return start1 <= end2 and end1 >= start2
-
-
-class Splitter:
-    """
-    Класс для работы с русскоязычными эмбеддингами в RAG пайплайне.
-    Поддерживает дообучение Word2Vec/FastText и использование RuBERT.
-    """
-    
-    def __init__(self, 
-                 chunk_size: int = 350,
-                 chunk_overlap: int = 70):
-
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-        
-        # Инициализация компонентов
-        self.chunks = []
-        self.chunk_metadata = []
-        self.documents = []
-    
-    def load_documents(self, documents: List[Dict]):
-        """
-        Загрузка документов и создание чанков.
-        
-        Args:
-            documents: Список словарей с полем 'text'
-        """
-        self.documents = documents
-        print(f"📄 Загрузка {len(documents)} документов...")
-        
-        chunks = []
-        docs_metadata = []      # ID документов для каждого чанка
-        paragraph_metadata = [] # ID абзацев для каждого чанка
-        paragraph_id_counter = 0
-
-        for doc_id, document in enumerate(documents):
-            # Разбиваем документ на абзацы по \n\n
-            paragraphs = document.split('\n')
-            
-            for paragraph in paragraphs:
-                paragraph = paragraph.strip()
-
-                if paragraph == '':
-                    continue
-                sentences = re.split(r'(?<=[.!?])\s+', paragraph)
-                
-                # Если абзац слишком длинный, используем сплиттер для его разбиения
-                if len(sentences) > 1:
-                    for chunk in sentences:
-                        if len(chunk.strip()) >= 30:
-                            chunks.append(chunk)
-                            docs_metadata.append(doc_id)
-                            paragraph_metadata.append(paragraph_id_counter)
-                else:
-                    # Добавляем абзац как целый чанк
-                    chunks.append(paragraph)
-                    docs_metadata.append(doc_id)
-                    paragraph_metadata.append(paragraph_id_counter)
-                
-                paragraph_id_counter += 1
-
-        print(f"✅ Создано {len(chunks)} чанков")
-        print(f"   Из {paragraph_id_counter} абзацев в {len(documents)} документах")
-        return chunks, docs_metadata, paragraph_metadata
-  

retrieval.py

@@ -1,13 +1,18 @@
 import re
+import time
+import hashlib
+import pickle
 from pathlib import Path
 
+import numpy as np
 import torch
+import pandas as pd
 from rank_bm25 import BM25Okapi
 import warnings
 warnings.filterwarnings('ignore')
 
 from _1_get_documents import load_and_process_data
-from _2_splitting import parse_year_metadata, years_overlap
+from _2_splitting import years_overlap, parse_metadata_from_document
 from lemmatizer import RussianLemmatizer
 # from _3_chunking import RussianEmbedder
 
@@ -18,9 +23,55 @@ from sentence_transformers import CrossEncoder
 
 
 class Retrieval:
-    def __init__(self, use_gpu: bool = False):
+    """
+    RAG (Retrieval-Augmented Generation) система на русском языке.
+    
+    Структура хранения данных:
+    ============================
+    
+    1. ДАТАФРЕЙМ ПАРАГРАФОВ (self.paragraphs_df):
+       ┌──────────────────────┬─────────────────────────────────┐
+       │ Колонка              │ Описание                         │
+       ├──────────────────────┼─────────────────────────────────┤
+       │ paragraph_id         │ Уникальный ID параграфа         │
+       │ summary              │ Название документа/раздела      │
+       │ start_year           │ Год начала периода              │
+       │ end_year             │ Год окончания периода           │
+       │ document_id          │ Ссылка на исходный документ     │
+       └──────────────────────┴─────────────────────────────────┘
+    
+    2. ДАТАФРЕЙМ ЧАНКОВ (self.chunks_df):
+       ┌──────────────────────┬─────────────────────────────────┐
+       │ Колонка              │ Описание                         │
+       ├──────────────────────┼─────────────────────────────────┤
+       │ chunk_id             │ Уникальный ID чанка             │
+       │ paragraph_id         │ Foreign key на параграф         │
+       │ text                 │ Исходный текст чанка            │
+       │ lemmatized_text      │ Лемматизированный текст         │
+       │ (embeddings)         │ (будет добавлено в будущем)     │
+       └──────────────────────┴─────────────────────────────────┘
+    
+    3. ОБЪЕДИНЁННЫЙ ДАТАФРЕЙМ (get_merged_data()):
+       Комбинирует оба датафрейма через JOIN по paragraph_id.
+       Содержит все колонки обоих датафреймов.
+       Используется для поиска и фильтрации.
+    
+    Ключевые преимущества:
+    - Избегаем дублирования метаданных параграфов
+    - Легко фильтровать по году, summary, документу
+    - Оптимизировано для работы с 5000+ чанками
+    - Простой merge для получения полной информации
+    """
+    
+    def __init__(self, use_gpu: bool = False, load_json: bool = True, use_cache: bool = True):
         print("Инициализация RAG системы...")
         self.device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
+        self.use_cache = use_cache
+        
+        # Путь к кэшу
+        self.cache_dir = Path('.cache')
+        if self.use_cache:
+            self.cache_dir.mkdir(exist_ok=True)
         
         # Инициализация лемматизатора для русского языка
         print("   Инициализация лемматизатора...")
@@ -32,85 +83,271 @@ class Retrieval:
         # self.documents after this phase: list of {'text': str, 'date': str}
         print(f"   Загружено {len(self.documents)} сообщений")
         
-        # Парсим даты из документов и создаем чанки
-        self.chunks, self.docs_metadata, self.paragraph_metadata, self.chunk_dates = self._process_documents_with_dates()
-        self.bm25 = self._prepare_bm25()
-        # self.embedder = RussianEmbedder(self.chunks, model_type="rubert")
-
-        # Создание индекса
-        # self.embeddings = self.embedder.create_index()
+        # Парсим даты из документов и создаем датафреймы
+        self.paragraphs_df, self.chunks_df = self._process_documents_with_dates()
+        
+        # Добавляем лемматизированный текст в датафрейм чанков с кэшем
+        print("2. Лемматизация текстов (с кэшированием)...")
+        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')
+        
+        print("✅ RAG система готова к использованию")
 
-        print("✓ RAG система готова к работе!")
-    
     def _process_documents_with_dates(self):
         """
-        Обрабатывает документы с парсингом дат и создает чанки.
+        Обрабатывает документы с парсингом дат и создает два датафрейма.
         
         Returns:
-            tuple: (chunks, docs_metadata, paragraph_metadata, chunk_dates)
-            где chunk_dates - список ((start_year, end_year), ...) для каждого чанка
+            tuple: (paragraphs_df, chunks_df)
+            
+            paragraphs_df:
+                - paragraph_id: уникальный идентификатор абзаца
+                - summary: название документа/раздела
+                - start_year: год начала периода
+                - end_year: год окончания периода
+                - document_id: ссылка на исходный документ
+            
+            chunks_df:
+                - chunk_id: уникальный идентификатор чанка
+                - paragraph_id: ссылка на абзац (foreign key)
+                - text: текст чанка
+                - lemmatized_text: лемматизированный текст (добавляется позже)
         """
-        chunks = []
-        docs_metadata = []
-        paragraph_metadata = []
-        chunk_dates = []
+        paragraphs_data = []
+        chunks_data = []
+        
         paragraph_id_counter = 0
+        chunk_id_counter = 0
         
         for doc_id, document in enumerate(self.documents):
-            try:
-                # Парсим даты из документа
-                dated_chunks = parse_year_metadata(document)
+            dated_chunks = parse_metadata_from_document(document)
+            
+            for chunk_text, year_range, summary in dated_chunks:
+                paragraphs = chunk_text.split('\n')
                 
-                for chunk_text, year_range in dated_chunks:
-                    # Разбиваем на предложения, как в Splitter
-                    paragraphs = chunk_text.split('\n')
+                for paragraph in paragraphs:
+                    paragraph = paragraph.strip()
+                    
+                    # Добавляем информацию о параграфе в датафрейм параграфов
+                    paragraphs_data.append({
+                        'paragraph_id': paragraph_id_counter,
+                        'summary': summary,
+                        'start_year': year_range[0],
+                        'end_year': year_range[1],
+                        'document_id': doc_id
+                    })
                     
-                    for paragraph in paragraphs:
-                        paragraph = paragraph.strip()
-                        if not paragraph:
-                            continue
-                        
-                        sentences = re.split(r'(?<=[.!?])\s+', paragraph)
-                        
-                        if len(sentences) > 1:
-                            for sent in sentences:
-                                if len(sent.strip()) >= 30:
-                                    chunks.append(sent)
-                                    docs_metadata.append(doc_id)
-                                    paragraph_metadata.append(paragraph_id_counter)
-                                    chunk_dates.append(year_range)
-                        else:
-                            chunks.append(paragraph)
-                            docs_metadata.append(doc_id)
-                            paragraph_metadata.append(paragraph_id_counter)
-                            chunk_dates.append(year_range)
-                        
-                        paragraph_id_counter += 1
-                        
-            except ValueError as e:
-                print(f"   ⚠️ Ошибка при парсинге документа {doc_id} ({self.docs_names[doc_id]}): {e}")
-                # Пропускаем документ если он не имеет правильной разметки
-                continue
-        
-        print(f"✅ Создано {len(chunks)} чанков")
-        print(f"   Из {len(set(paragraph_metadata))} абзацев в {doc_id + 1} документах")
-        return chunks, docs_metadata, paragraph_metadata, chunk_dates
-    
-    def _prepare_bm25(self):
-        """Подготавливаем BM25 индекс для ключевого поиска"""
-        # Токенизация для BM25
-        tokenized_chunks = [self.lemmatizer.tokenize_text(chunk) for chunk in self.chunks]
-        return BM25Okapi(tokenized_chunks)
+                    # Разбиваем параграф на предложения и создаем чанки
+                    sentences = re.split(r'(?<=[.!?])\s+', paragraph)                    
+                    for sent in sentences:
+                        chunks_data.append({
+                            'chunk_id': chunk_id_counter,
+                            'paragraph_id': paragraph_id_counter,
+                            'text': sent.strip()
+                        })
+                        chunk_id_counter += 1
+                    
+                    paragraph_id_counter += 1
+        
+        # Создаем датафреймы
+        paragraphs_df = pd.DataFrame(paragraphs_data)
+        chunks_df = pd.DataFrame(chunks_data)
+        
+        print(f"Создано {len(chunks_df)} чанков")
+        print(f"Из {len(paragraphs_df)} абзацев в {len(set(paragraphs_df['document_id']))} документах")
+        
+        return paragraphs_df, chunks_df
+
+    # ============ Методы кэширования лемматизации ============
+    
+    @staticmethod
+    def _compute_text_hash(text: str) -> str:
+        """
+        Вычисляет SHA256 хэш текста.
+        
+        Args:
+            text: Текст для хэширования
+            
+        Returns:
+            str: Хэш в hex формате
+        """
+        return hashlib.sha256(text.encode('utf-8')).hexdigest()
+    
+    def _load_cache(self) -> dict:
+        """
+        Загружает кэш лемматизации из файловой системы.
+        
+        Returns:
+            dict: {text_hash -> lemmatized_tokens}
+        """
+        cache_file = self.cache_dir / 'lemmatization_cache.pkl'
+        
+        if cache_file.exists():
+            try:
+                with open(cache_file, 'rb') as f:
+                    cache = pickle.load(f)
+                print(f"   ✓ Кэш загружен ({len(cache)} записей)")
+                return cache
+            except Exception as e:
+                print(f"   ⚠ Ошибка при загрузке кэша: {e}")
+                return {}
+        return {}
+    
+    def _save_cache(self, cache: dict) -> None:
+        """
+        Сохраняет кэш лемматизации в файловую систему.
+        
+        Args:
+            cache: {text_hash -> lemmatized_tokens}
+        """
+        cache_file = self.cache_dir / 'lemmatization_cache.pkl'
+        
+        try:
+            with open(cache_file, 'wb') as f:
+                pickle.dump(cache, f)
+            print(f"   ✓ Кэш сохранён ({len(cache)} записей)")
+        except Exception as e:
+            print(f"   ⚠ Ошибка при сохранении кэша: {e}")
+    
+    def _lemmatize_with_cache(self, texts: list[str]) -> list:
+        """
+        Лемматизирует тексты с использованием кэша.
+        Проверяет хэши текстов - если хэш совпадает с кэшированным,
+        использует кэшированный результат. Иначе перелемматизирует.
+        
+        Args:
+            texts: Список текстов для лемматизации
+            
+        Returns:
+            list: Лемматизированные тексты
+        """
+        if not self.use_cache:
+            # Если кэш отключен, просто лемматизировать
+            return [self.lemmatizer.tokenize_text(text) for text in texts]
+        
+        # Загружаем существующий кэш
+        cache = self._load_cache()
+        text_hashes = {}
+        results = []
+        needs_save = False
+        
+        for text in texts:
+            text_hash = self._compute_text_hash(text)
+            text_hashes[text] = text_hash
+            
+            if text_hash in cache:
+                # Используем кэшированный результат
+                results.append(cache[text_hash])
+            else:
+                # Лемматизируем и добавляем в кэш
+                lemmatized = self.lemmatizer.tokenize_text(text)
+                results.append(lemmatized)
+                cache[text_hash] = lemmatized
+                needs_save = True
+        
+        # Сохраняем кэш если были новые записи
+        if needs_save:
+            self._save_cache(cache)
+        
+        return results
+    
+    def clear_cache(self) -> None:
+        """
+        Очищает кэш лемматизации.
+        """
+        cache_file = self.cache_dir / 'lemmatization_cache.pkl'
+        
+        try:
+            if cache_file.exists():
+                cache_file.unlink()
+                print("✓ Кэш очищен")
+            else:
+                print("⚠ Файл кэша не найден")
+        except Exception as e:
+            print(f"⚠ Ошибка при очистке кэша: {e}")
+    
+    def get_cache_stats(self) -> dict:
+        """
+        Возвращает статистику кэша.
+        
+        Returns:
+            dict: Информация о кэше
+        """
+        cache_file = self.cache_dir / 'lemmatization_cache.pkl'
+        
+        if cache_file.exists():
+            cache = self._load_cache() if self.use_cache else {}
+            file_size_mb = cache_file.stat().st_size / (1024 * 1024)
+            
+            return {
+                'cache_enabled': self.use_cache,
+                'cache_file': str(cache_file),
+                'cached_entries': len(cache),
+                'file_size_mb': round(file_size_mb, 2),
+                'exists': True
+            }
+        else:
+            return {
+                'cache_enabled': self.use_cache,
+                'cache_file': str(cache_file),
+                'cached_entries': 0,
+                'file_size_mb': 0,
+                'exists': False
+            }
 
+    # ============ Вспомогательные методы для работы с датафреймами ============
+    
+    def get_merged_data(self):
+        """Возвращает объединённый датафрейм чанков с метаданными параграфов.
+        
+        Returns:
+            pd.DataFrame: Датафрейм с полями:
+                chunk_id, paragraph_id, text, lemmatized_text,
+                summary, start_year, end_year, document_id
+        """
+        return self.chunks_df.merge(
+            self.paragraphs_df,
+            on='paragraph_id',
+            how='left'
+        )
+    
+    def filter_by_year_range(self, year_range: tuple[int, int]) -> pd.DataFrame:
+        """Возвращает чанки, которые пересекаются с заданным диапазоном лет.
+        
+        Args:
+            year_range: (start_year, end_year)
+            
+        Returns:
+            pd.DataFrame: Отфильтрованные чанки с метаданными
+        """
+        merged = self.get_merged_data()
+        
+        # Проверяем пересечение диапазонов
+        return merged[
+            (merged['start_year'] <= year_range[1]) &
+            (merged['end_year'] >= year_range[0])
+        ]
+    
     def rerank_search(self, query: str) -> list[dict]:
+        """Ранжирует все чанки используя CrossEncoder модель.
+        
+        Args:
+            query: Текст запроса
+            
+        Returns:
+            list: Отсортированный список результатов с scores
         """
-        [{'corpus_id': 0, 'score': 0.88126713},
-#  {'corpus_id': 2, 'score': 0.001042091},
-#  {'corpus_id': 3, 'score': 0.0010417715},
-#  {'corpus_id': 1, 'score': 0.0010344835},
-#  {'corpus_id': 4, 'score': 0.0010244923}]`"""
-        reranker_model = CrossEncoder('DiTy/cross-encoder-russian-msmarco')
-        return reranker_model.rank(query[0], self.chunks)
+        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')
 
     def semantic_search(self, query: str) -> list:    
         # 1. Семантический поиск
@@ -118,26 +355,105 @@ class Retrieval:
         semantic_scores = torch.nn.functional.cosine_similarity(self.embeddings, query_embedding, eps=1e-8).cpu()
         return semantic_scores
     
-    def bm25_search(self, query: str) -> list: 
-        # 2. Ключевой поиск (BM25)
+    def bm25_search(self, query: str) -> list:
+        """BM25 поиск, используя лемматизированные чанки.
+        
+        Args:
+            query: Текст запроса
+            
+        Returns:
+            list: Скоры для каждого чанка
+        """
+        bm25 = BM25Okapi(self.chunks_df['lemmatized_text'].tolist())
         tokenized_query = self.lemmatizer.tokenize_text(query)
-        return self.bm25.get_scores(tokenized_query)
+        return bm25.get_scores(tokenized_query)
     
-    def filter_by_year_range(self, indices: list[int], year_range: tuple[int, int]) -> list[int]:
-        """
-        Фильтрует индексы чанков по диапазону лет (с пересечением).
+    # ============ Для тестирования cross-encoder ============ 
+    def test_query_with_cross_encoder(self, query: str,
+                                      target_summary: str):
+        """ Тестирует запрос с cross-encoder и выводит результаты.
         
         Args:
-            indices: Список индексов чанков для фильтрации
-            year_range: (start_year, end_year) для поиска
+            query: Текст запроса
+            target_summary: Ожидаемый summary
+        """
+        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]
             
-        Returns:
-            list: Отфильтрованный список индексов
-        """
-        filtered = []
-        for idx in indices:
-            if idx < len(self.chunk_dates):
-                chunk_range = self.chunk_dates[idx]
-                if years_overlap(chunk_range, year_range):
-                    filtered.append(idx)
-        return filtered
+            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")

test_cross_encoder.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass
+from retrieval import Retrieval
+
+
+@dataclass
+class TestCaseForCrossEncoder:
+    query: str
+    good_answer: str
+
+
+test_cases = [
+    TestCaseForCrossEncoder(
+        'Какие изменения в транспорте Рязани были бы полезны на текущий момент?',
+        'Актуальные проекты новых троллейбусных линий, которые полезно бы построить',
+        ),
+    TestCaseForCrossEncoder(
+        'Какие продления троллейбусной сети были бы полезны на текущий момент?',
+        'Актуальные проекты новых троллейбусных линий, которые полезно бы построить',
+        ),
+    TestCaseForCrossEncoder(
+        'Расскажи о провалившихся экспериментах в Рязани',
+        'Попытки (все из которых неудачные) запустить городскую электричку в истории',
+    ),
+    TestCaseForCrossEncoder(
+        'Расскажи историю маршрута маршрутки № 92 в Рязани',
+        'история ныне закрытой маршрутки № 92',
+),
+    TestCaseForCrossEncoder(
+        'Какой маршрут в Рязани закрылся из-за плохой трассировки?',
+        'У троллейбусного маршрута №2 была неудачная трасса - в объезд основных узлов города',
+    ),
+    TestCaseForCrossEncoder(
+    'Когда маршрут троллейбуса №10 продлили до площади Попова?',
+    'история троллейбусного маршрута № 10')
+]
+
+
+def test_cross_encoder_vs_bm25():
+    """Тестирует кросс-энкодер vs BM25 на всех документах."""
+    print("=" * 90)
+    print("СРАВНЕНИЕ: КРОСС-ЭНКОДЕР vs BM25 ЛЕММАТИЗИРОВАННЫЙ ПОИСК")
+    print("=" * 90)
+    
+    # Создаем объект Retrieval (загружает корпус автоматически)
+    retrieval = Retrieval(use_gpu=False)
+    
+    # Тестируем каждый тестовый случай
+    print("=" * 90)
+    print("ТЕСТИРОВАНИЕ ОТДЕЛЬНЫХ ЗАПРОСОВ")
+    print("=" * 90)
+    
+    for test_num, test_case in enumerate(test_cases, 1):
+        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__":
+    test_cross_encoder_vs_bm25()
+

tests/test_retirieval.py

@@ -0,0 +1,5 @@
+from retrieval import Retrieval
+
+retr = Retrieval(use_gpu=False)
+res = retr.bm25_search('канищево', top_k=5)
+print(res)