update logic

_1_get_documents.py

@@ -8,16 +8,25 @@ def get_text(inst):
     if isinstance(inst, dict):
         return get_text(inst['text'])
  
+
+def process_file(file_path):
+    with open(file_path, 'r', encoding='utf-8-sig') as f:
+        text = f.read()
+        assert text
+    return str(file_path).split('.')[-1], text
+
+def process_folder_recursive(folder_path):
+    all_messages = []
+    for file in os.listdir(folder_path):
+        file_path = os.path.join(folder_path, file)
+        if os.path.isfile(file_path):
+            all_messages.append(process_file(file_path))
+        else:
+            all_messages += process_folder_recursive(file_path)
+    return all_messages
     
 def load_and_process_data() -> list[dict]:
     """Загрузка и предобработка данных из JSON файлов"""
-    all_messages = []
-    doc_names = os.listdir('texts')
-    txt_paths = ['texts/'+file for file in doc_names] 
-    for file_path in txt_paths:
-        with open(file_path, 'r', encoding='utf-8-sig') as f:
-            text = f.read()
-            assert text
-            all_messages.append(text)
+    all_messages = process_folder_recursive('texts')
 
-    return all_messages, [x[:-3] for x in doc_names]  # убираем расширение .md из имен документов
+    return [x[0] for x in all_messages], [x[1][:-3] for x in all_messages]  # возвращаем расширения и тексты документов

app.py

@@ -1,8 +1,5 @@
+import re
 import gradio as gr
-import numpy as np
-import plotly.express as px
-import plotly.graph_objects as go
-import pandas as pd
 from generation import wrap_prompt
 from llm import get_llm_answer
 from retrieval import Retrieval
@@ -14,90 +11,6 @@ vocabulary = parse_vocabulary('vocabulary/vocabulary.md')
 retrieval = Retrieval()
 
 
-def create_heatmap(scores, chunk_ids, top_k_indices=None):
-    """Создает heatmap релевантности документов по чанкам"""
-    if len(scores) == 0:
-        return go.Figure()
-   
-    # Группируем чанки по документам
-    docs_chunks = {}
-    chunk_to_doc_map = {}
-    
-    for idx, (chunk_id, score) in enumerate(zip(chunk_ids, scores)):
-        doc_id = retrieval.docs_metadata[chunk_id]
-        doc_name = retrieval.docs_names[doc_id]
-        chunk_to_doc_map[chunk_id] = doc_name
-                
-        if doc_name not in docs_chunks:
-            docs_chunks[doc_name] = []
-        
-        # Сохраняем информацию о чанке
-        docs_chunks[doc_name].append({
-            'absolute_idx': idx,
-            'chunk_id': chunk_id,
-            'score': score,
-            'in_top_k': top_k_indices is not None and idx in top_k_indices
-        })
-    
-    if not docs_chunks:
-        return go.Figure()
-    
-    # Сортируем чанки внутри каждого документа по chunk_id
-    for doc_name in docs_chunks:
-        docs_chunks[doc_name].sort(key=lambda x: x['chunk_id'])
-    
-    # Создаем DataFrame для heatmap с относительными номерами чанков
-    df_data = []
-    for doc_name, chunks in docs_chunks.items():
-        for relative_idx, chunk_info in enumerate(chunks):
-            df_data.append({
-                'Документ': doc_name,
-                'Чанк (внутри документа)': f'Чанк {relative_idx + 1}',
-                'Релевантность': chunk_info['score'],
-                'Абсолютный ID': chunk_info['chunk_id'],
-                'В top-k': chunk_info['in_top_k']
-            })
-    
-    df = pd.DataFrame(df_data)
-    
-    # Создаем heatmap
-    fig = px.density_heatmap(
-        df, 
-        x='Чанк (внутри документа)', 
-        y='Документ', 
-        z='Релевантность',
-        title='Heatmap релевантности (по документам, с относительными номерами чанков)',
-        color_continuous_scale='Viridis',
-        labels={'Релевантность': 'Score'}
-    )
-    
-    # Добавляем обводку для top-k чанков
-    top_k_df = df[df['В top-k'] == True]
-    if not top_k_df.empty:
-        fig.add_trace(go.Scatter(
-            x=top_k_df['Чанк (внутри документа)'],
-            y=top_k_df['Документ'],
-            mode='markers',
-            marker=dict(
-                symbol='circle-open',
-                size=20,
-                line=dict(color='red', width=2),
-                color='rgba(0,0,0,0)'
-            ),
-            name='Top-k чанки',
-            showlegend=True
-        ))
-    
-    fig.update_layout(
-        xaxis={'side': 'bottom', 'tickangle': -45},
-        height=max(400, len(docs_chunks) * 30),  # Адаптивная высота
-        width=800,
-        xaxis_title="Номер чанка в документе",
-        yaxis_title="Документ"
-    )
-    
-    return fig
-
 def perform_search(query, top_k, year_from, year_to):
     """Этап 1: Поиск и возврат результатов с фильтром по датам"""
    
@@ -125,6 +38,25 @@ def perform_search(query, top_k, year_from, year_to):
 
     # Выполняем поиск BM25
     scores = retrieval.bm25_search(query)
+    scores = list(scores)  # Преобразуем в список если это ndarray
+    
+    # Повышаем scores для документов, названия которых содержат паттерн маршрута из query
+    # Паттерн: [АМт]\d{2} (буква А, М или Т + две цифры, например А10, М30, Т2)
+    pattern = r'[АМт]\d{2}'
+    matches = re.findall(pattern, query)
+    
+    if matches:
+        max_score = max(scores) if scores else 0
+        boost_score = max_score + 1  # Максимальный score + 1
+        
+        for match in set(matches):  # Используем set чтобы избежать дубликатов
+            # Ищем документы, которые содержат этот паттерн в названии
+            for doc_id, doc_name in enumerate(retrieval.docs_names):
+                if match in doc_name:
+                    # Повышаем scores всех чанков из этого документа
+                    for chunk_id, chunk_doc_id in enumerate(retrieval.docs_metadata):
+                        if chunk_doc_id == doc_id:
+                            scores[chunk_id] = boost_score
     
     # Получаем индексы чанков
     chunk_ids = list(range(len(scores)))
@@ -147,29 +79,6 @@ def perform_search(query, top_k, year_from, year_to):
     
     return scores, chunk_ids, top_k_indices, status
 
-def filter_chunks_by_documents(top_k_indices, all_scores, selected_docs):
-    """Фильтрует чанки по выбранным документам"""
-    if len(top_k_indices)==0 or len(all_scores)==0:
-        return []
-    
-    filtered_indices = []
-    for idx in top_k_indices:
-        if idx >= len(retrieval.docs_metadata):
-            continue
-            
-        doc_id = retrieval.docs_metadata[idx]
-        doc_name = retrieval.docs_names[doc_id] if doc_id < len(retrieval.docs_names) else "Неизвестный документ"
-        
-        # Если документы выбраны, проверяем наличие в списке
-        if selected_docs and len(selected_docs) > 0:
-            if doc_name in selected_docs:
-                filtered_indices.append(idx)
-        else:
-            # Если ничего не выбрано, показываем все
-            filtered_indices.append(idx)
-    
-    return filtered_indices
-
 def format_selected_chunks(selected_indices):
     """Форматирует выбранные чанки в единый текст для вывода и LLM
     
@@ -352,8 +261,7 @@ def ask_llm(query, filtered_indices_state):
 
 # Создаем интерфейс Gradio
 with gr.Blocks(title="RAG Application", theme=gr.themes.Soft()) as iface:
-    gr.Markdown("# RAG Application для исторических документов")
-    gr.Markdown("## Двухэтапная работа с документами")
+    gr.Markdown("#№ Справочник по общественного истории транспорта Рязани")
     
     # Строка 1: поиск и фильтр по датам
     with gr.Row():
@@ -406,14 +314,6 @@ with gr.Blocks(title="RAG Application", theme=gr.themes.Soft()) as iface:
         )
     
     with gr.Row():
-        with gr.Column(scale=1):
-            # Фильтр ПОСЛЕ поиска для документов
-            docs_after = gr.CheckboxGroup(
-                choices=retrieval.docs_names,
-                label="Фильтр по документам",
-                info="Выберите документы (если ничего не выбрано - показываются все)"
-            )
-            
         with gr.Column(scale=2):
             # Большое текстовое поле для результатов retrieval
             retrieval_results = gr.Textbox(
@@ -454,27 +354,12 @@ with gr.Blocks(title="RAG Application", theme=gr.themes.Soft()) as iface:
         fn=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=filter_chunks_by_documents,
-        inputs=[top_k_indices_state, all_scores_state, docs_after],
-        outputs=[filtered_indices_state]
     ).then(
         fn=format_retrieval_results,
         inputs=[filtered_indices_state, top_k_slider],
         outputs=[retrieval_results]
     )
-    
-    # Обработчик изменения фильтра документов
-    docs_after.change(
-        fn=filter_chunks_by_documents,
-        inputs=[top_k_indices_state, all_scores_state, docs_after],
-        outputs=[filtered_indices_state]
-    ).then(
-        fn=format_retrieval_results,
-        inputs=[filtered_indices_state, top_k_slider],
-        outputs=[retrieval_results]
-    )
-    
+
     # Обработчик изменения слайдера top_k
     top_k_slider.change(
         fn=format_retrieval_results,

retrieval.py

@@ -7,7 +7,7 @@ import warnings
 warnings.filterwarnings('ignore')
 
 from _1_get_documents import load_and_process_data
-from _2_splitting import parse_year_metadata, years_overlap, YEAR_OLD, YEAR_NEW
+from _2_splitting import parse_year_metadata, years_overlap
 from lemmatizer import RussianLemmatizer
 # from _3_chunking import RussianEmbedder