fix app

app.py

@@ -1,5 +1,6 @@
-import re
 import gradio as gr
+import numpy as np
+
 from generation import wrap_prompt
 from llm import get_llm_answer
 from retrieval import Retrieval
@@ -8,114 +9,69 @@ from vocabulary.parse_vocabulary import parse_vocabulary
 
 
 vocabulary, _ = parse_vocabulary('vocabulary/vocabulary.md')
-retrieval = Retrieval()
 
 
-def perform_search(query, top_k, year_from, year_to):
-    """Этап 1: Поиск и возврат результатов с фильтром по датам"""
-   
-    if not query:
-        return None, [], [], "Введите вопрос для поиска"
-    
-    # Преобразуем входные значения
-    try:
-        year_from = _parse_single_year(year_from)
-        year_to = _parse_single_year(year_to)
-            
-        # Проверяем корректность диапазона
-        if year_from > year_to:
-            year_from, year_to = year_to, year_from
-            
-    except (ValueError, TypeError):
-        return None, [], [], f"⚠️ Ошибка: некорректный диапазон лет ({year_from} - {year_to})"
+class Perform:
 
-    # Выполняем поиск BM25
-    scores = retrieval.bm25_search(query)
-    scores = list(scores)  # Преобразуем в список если это ndarray    
-    
-    # Получаем индексы чанков
-    chunk_ids = list(range(len(scores)))
-    
-    # Применяем ЖЕСТКИЙ фильтр по датам ДО выбора top-k
-    year_search_range = (year_from, year_to)
-    filtered_by_date = retrieval.filter_by_year_range(chunk_ids, year_search_range)
-    
-    # Если нет результатов после фильтра по датам
-    if not filtered_by_date:
-        return scores, chunk_ids, [], f"⚠️ Нет результатов в диапазоне {year_from}-{year_to}"
-    
-    # Находим top-k среди отфильтрованных по датам (сортируем по релевантности BM25)
-    top_k = min(top_k, len(filtered_by_date))
-    filtered_scores = [(idx, scores[idx]) for idx in filtered_by_date]
-    filtered_scores.sort(key=lambda x: x[1], reverse=True)
-    top_k_indices = [idx for idx, _ in filtered_scores[:top_k]]
-    
-    status = f"Найдено {len(scores)} чанков, {len(filtered_by_date)} в диапазоне {year_from}-{year_to}. Top-{top_k} выбраны."
-    
-    return scores, chunk_ids, top_k_indices, status
+    def __init__(self):
+        self.retrieval = Retrieval()
+        lengthh = len(self.retrieval.paragraphs_df)
+        self.scores = None
+        self.sorted_idx = None
+        self.years_mask = np.ones(lengthh, dtype=bool)
 
-def format_selected_chunks(selected_indices):
-    """Форматирует выбранные чанки в единый текст для вывода и LLM
-    
-    Выводит целые абзацы по выбранным чанкам с названиями документов:
-    Документ {название}:
-    {полный текст абзаца}
-    """
-    if not selected_indices:
-        return ""
-    
-    # Найдем все уникальные абзацы из выбранных чанков
-    paragraphs_to_show = {}  # paragraph_id -> doc_id
-    
-    for idx in selected_indices:
-        if idx >= len(retrieval.paragraph_metadata) or idx >= len(retrieval.docs_metadata):
-            continue
-        
-        paragraph_id = retrieval.paragraph_metadata[idx]
-        doc_id = retrieval.docs_metadata[idx]
-        paragraphs_to_show[paragraph_id] = doc_id
-    
-    # Для каждого отмеченного абзаца найдем ВСЕ его чанки
-    full_paragraph_chunks = {}  # paragraph_id -> [chunk_ids]
-    for chunk_id, paragraph_id in enumerate(retrieval.paragraph_metadata):
-        if paragraph_id in paragraphs_to_show:
-            if paragraph_id not in full_paragraph_chunks:
-                full_paragraph_chunks[paragraph_id] = []
-            full_paragraph_chunks[paragraph_id].append(chunk_id)
-    
-    # Форматируем вывод
-    result_lines = []
-    for paragraph_id in sorted(paragraphs_to_show.keys()):
-        doc_id = paragraphs_to_show[paragraph_id]
-        chunk_indices = sorted(full_paragraph_chunks[paragraph_id])
-        
-        doc_name = retrieval.docs_names[doc_id] if doc_id < len(retrieval.docs_names) else "Неизвестный документ"
-        
-        # Объединяем все чанки абзаца в полный текст
-        paragraph_text = " ".join([retrieval.chunks[idx] for idx in chunk_indices])
-        
-        # Форматируем вывод с названием документа
-        result_lines.append(f"Документ {doc_name}:")
-        result_lines.append(paragraph_text)
-        result_lines.append("")  # Пустая строка между абзацами
-    
-    return "\n".join(result_lines)
+    def get_years_range_mask(self, year_from, year_to):
+        try:
+            year_from = _parse_single_year(year_from)
+            year_to = _parse_single_year(year_to)
+            if year_from > year_to:
+                year_from, year_to = year_to, year_from
+        except (ValueError, TypeError):
+            raise ValueError(f"Некорректный диапазон лет: {year_from} - {year_to}")
 
+        self.years_mask = (
+            (self.retrieval.paragraphs_df["end_year"] >= year_from) &
+            (self.retrieval.paragraphs_df["start_year"] <= year_to)
+        ).values
+
+    def perform_search(self, query, top_k, year_from, year_to):
+        self.get_years_range_mask(year_from, year_to)
+
+        # если есть query → считаем scores
+        if query:
+            self.scores = self.retrieval.search(query)
+            self.sorted_idx = np.argsort(self.scores)[::-1]
+
+        # если нет query и scores нет → используем только фильтр
+        if self.scores is None:
+            filtered = np.where(self.years_mask)[0]
+            if len(filtered) <= top_k:
+                return None, None, filtered, "Показаны все записи по фильтру лет"
+            return None, None, filtered[-top_k:], "Показаны записи по фильтру лет"
+
+        # применяем mask к отсортированным индексам
+        filtered_sorted = self.sorted_idx[self.years_mask[self.sorted_idx]]
+
+        if len(filtered_sorted) == 0:
+            return self.scores, None, [], "⚠️ Нет результатов в выбранном диапазоне лет"
+
+        top_k_indices = filtered_sorted[:top_k]
+
+        return self.scores, None, top_k_indices, f"Найдено {len(filtered_sorted)} результатов"
+
+    def format_retrieval_results(self, top_k_indices):
+        if len(top_k_indices) == 0:
+            return "Нет результатов"
+
+        texts = self.retrieval.paragraphs_df["texts"].iloc[top_k_indices]
+        return "\n\n".join(texts)
+
+    def format_selected_chunks(self, indices):
+        texts = self.retrieval.paragraphs_df["texts"].iloc[indices]
+        return "\n\n".join(texts)
+
+perform = Perform()
 
-def format_retrieval_results(filtered_indices, top_k_results):
-    """Форматирует результаты retrieval для отображения в текстовом поле
-    
-    Берет top_k результатов и выводит целые абзацы с названиями документов
-    """
-    if len(filtered_indices) == 0:
-        return "Нет результатов"
-    
-    top_k_results = min(top_k_results, len(filtered_indices))
-    
-    # Берем top-k индексов (уже отсортированы по релевантности)
-    top_k_indices = filtered_indices[:top_k_results]
-    
-    return format_selected_chunks(top_k_indices)
 
 def ask_llm(query, filtered_indices_state):
     """Этап 2: Отправка отфильтрованных чанков в LLM с потоковой выдачей"""
@@ -123,22 +79,19 @@ def ask_llm(query, filtered_indices_state):
         yield "Введите вопрос"
         return
     
-    # Используем все отфильтрованные чанки
-    chunks_to_use = filtered_indices_state if filtered_indices_state else []
-    
-    if not chunks_to_use:
+    if not filtered_indices_state:
         yield "Нет выбранных чанков для отправки в LLM"
         return
     
     # Форматируем контекст используя ту же функцию, что и в интерфейсе
-    context = format_selected_chunks(list(chunks_to_use))
+    context = perform.format_selected_chunks(filtered_indices_state)
     
     if not context or context == "Нет валидных чанков":
         yield "Нет валидных чанков для отправки"
         return
     
     # Формируем промпт и отправляем в LLM
-    prompt = wrap_prompt(context, query, vocabula=vocabulary.copy())
+    prompt = wrap_prompt(context, query, vocabulary)
     
     # Потоковая выдача ответа
     full_answer = ""
@@ -237,18 +190,18 @@ with gr.Blocks(title="RAG Application", theme=gr.themes.Soft()) as iface:
     
     # Обработчик поиска
     search_btn.click(
-        fn=perform_search,
+        fn=perform.perform_search,
         inputs=[search_query_input, top_k_slider, year_from_input, year_to_input],
         outputs=[all_scores_state, all_chunk_ids_state, top_k_indices_state, search_status]
     ).then(
-        fn=format_retrieval_results,
+        fn=perform.format_retrieval_results,
         inputs=[top_k_indices_state, top_k_slider],
         outputs=[retrieval_results]
     )
 
     # Обработчик изменения слайдера top_k
     top_k_slider.change(
-        fn=format_retrieval_results,
+        fn=perform.format_retrieval_results,
         inputs=[top_k_indices_state, top_k_slider],
         outputs=[retrieval_results]
     )

generation.py

@@ -64,7 +64,8 @@ def lemmatize(text, vocabulary):
     return found_terms
 
 
-def wrap_prompt(retrieved_text, query_text, vocabula: dict[str, str]):
+def wrap_prompt(retrieved_text, query_text, inp_vocabula: dict[str, str]):
+    vocabula = inp_vocabula.copy()  # Создаем копию словаря, чтобы не изменять оригинал
     tokens_from_query = lemmatize(query_text, vocabula)
     tokens_from_retrieved = lemmatize(retrieved_text, vocabula)
     info_for_llm = ''

retrieval.py

@@ -29,18 +29,19 @@ class Retrieval:
     
     1. ДАТАФРЕЙМ ПАРАГРАФОВ (self.paragraphs_df):
        ┌──────────────────────┬─────────────────────────────────┐
-       │ Колонка              │ Описание                         │
+       │ Колонка              │ Описание                        │
        ├──────────────────────┼─────────────────────────────────┤
        │ paragraph_id         │ Уникальный ID параграфа         │
        │ summary              │ Название документа/раздела      │
        │ start_year           │ Год начала периода              │
        │ end_year             │ Год окончания периода           │
+       │ text                 │ Текст                           │
        │ document_id          │ Ссылка на исходный документ     │
        └──────────────────────┴─────────────────────────────────┘
     
     2. ДАТАФРЕЙМ ЧАНКОВ (self.chunks_df):
        ┌──────────────────────┬─────────────────────────────────┐
-       │ Колонка              │ Описание                         │
+       │ Колонка              │ Описание                        │
        ├──────────────────────┼─────────────────────────────────┤
        │ chunk_id             │ Уникальный ID чанка             │
        │ paragraph_id         │ Foreign key на параграф         │
@@ -101,6 +102,7 @@ class Retrieval:
                 - summary: название документа/раздела
                 - start_year: год начала периода
                 - end_year: год окончания периода
+                - text: текст абзаца
                 - document_id: ссылка на исходный документ
             
             chunks_df:
@@ -130,6 +132,7 @@ class Retrieval:
                         'summary': summary,
                         'start_year': year_range[0],
                         'end_year': year_range[1],
+                        'text': paragraph,
                         'document_id': doc_id
                     })
                     
@@ -361,11 +364,8 @@ class Retrieval:
         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, weight_bm25: float = 0.5, weight_semantic: float = 0.5) -> None:
-        """ Тестирует запрос с cross-encoder и выводит результаты.
-        
+    def search(self, query: str, target_summary: str, weight_bm25: float = 0.5, weight_semantic: float = 0.5) -> None:
+        """       
         Args:
             query: Текст запроса
             target_summary: Ожидаемый summary