more refactor

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import numpy as np
 
-from generation import wrap_prompt
 from llm import get_llm_answer
 from retrieval import Retrieval
 from _2_splitting import _parse_single_year
@@ -95,7 +94,7 @@ def ask_llm(query, filtered_indices_state):
         return
     
     # Формируем промпт и отправляем в LLM
-    prompt = wrap_prompt(context, query, VOCABULARY_MANAGER.vocabulary)
+    prompt = VOCABULARY_MANAGER.wrap_prompt(context, query)
     
     # Потоковая выдача ответа
     full_answer = ""

generation.py

@@ -1,65 +0,0 @@
-import warnings
-import re
-warnings.filterwarnings('ignore')
-
-
-def lemmatize(text: str, vocabulary: dict[str, str]) -> list[str]:
-    """
-    Лемматизирует текст и возвращает список терминов из словаря.
-    Поддерживает поиск полных термин и частичный поиск по словам.
-    Также ищет словоформы (например, "икарусы" для "икарус").
-    
-    Args:
-        text (str): Текст для анализа
-        vocabulary (dict): Словарь терминов {термин: определение}
-    
-    Returns:
-        list: Список найденных терминов из vocabulary в порядке их появления
-    """
-    found_terms = []
-    found_terms_set = set()  # Для отслеживания дубликатов
-    
-    # Приводим текст в нижний регистр для поиска
-    text_lower = text.lower()
-    
-    # Сначала - поиск полных терминов
-    for term in vocabulary.keys():
-        term_lower = term.lower()
-        # Ищем термин как целое слово
-        pattern = r'\b' + re.escape(term_lower) + r'\b'
-        
-        if re.search(pattern, text_lower) and term not in found_terms_set:
-            found_terms.append(term)
-            found_terms_set.add(term)
-        
-    return found_terms
-
-
-def wrap_prompt(retrieved_text: str, query_text: str, 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 = ''
-    for token in tokens_from_query + tokens_from_retrieved:
-        if token in vocabula:
-            info_for_llm += f'{token}: {vocabula[token]}\n'
-            del vocabula[token]  # удаляем, чтобы не повторять
-    if info_for_llm:
-        info_for_llm = f"\nСправка по терминам (пользователю эти термины очевидны)\n{info_for_llm}\n" + info_for_llm
-    return f"""Ты - эксперт по истории общественного транспорта в городе Рязань, отвечающий на вопросы на основе предоставленных архивных данных.
-
-ВОПРОС ПОЛЬЗОВАТЕЛЯ:
-{query_text}    
-{info_for_llm}
-
-КОНТЕКСТ ИЗ АРХИВОВ:
-{retrieved_text}
-
-ИНСТРУКЦИИ:
-1. Отвечай ТОЛЬКО на основе предоставленного контекста
-2. Если в контексте нет информации для ответа, скажи об этом честно
-3. Сохраняй историческую точность, упоминай даты, если они есть
-4. Будь кратким и информативным
-5. Цитируй конкретные факты из контекста
-
-ОТВЕТ:"""

vocabulary/parse_vocabulary.py

@@ -103,6 +103,50 @@ class ParserVocabulary:
                     return info['canonical']
         return None
 
+    def find_terms(self, text: str) -> set[str]:
+        """Возвращает список терминов в тексте (учитывая словоформы)
+        
+        Args:
+            text (str): Текст для анализа
+        
+        Returns:
+            set: Множество найденных терминов из vocabulary в порядке их появления
+        """
+        found_terms = set()  # Для отслеживания дубликатов
+        
+        text_lower = text.lower()
+        
+        for word in text_lower:
+            res = self.extract_stem(word)
+            if res:
+                found_terms.add(word)           
+        return found_terms
+
+    def wrap_prompt(self, retrieved_text: str, query_text: str):
+        tokens_from_query = self.find_terms(query_text)
+        tokens_from_retrieved = self.find_terms(retrieved_text)
+        info_for_llm = ''
+        for token in tokens_from_query + tokens_from_retrieved:
+            info_for_llm += f'{token}: {self.vocabulary[token]}\n'
+        if info_for_llm:
+            info_for_llm = f"\nСправка по терминам (пользователю эти термины очевидны)\n{info_for_llm}\n" + info_for_llm
+        return f"""Ты - эксперт по истории общественного транспорта в городе Рязань, отвечающий на вопросы на основе предоставленных архивных данных.
+
+ВОПРОС ПОЛЬЗОВАТЕЛЯ:
+{query_text}    
+{info_for_llm}
+
+КОНТЕКСТ ИЗ АРХИВОВ:
+{retrieved_text}
+
+ИНСТРУКЦИИ:
+1. Отвечай ТОЛЬКО на основе предоставленного контекста
+2. Если в контексте нет информации для ответа, скажи об этом честно
+3. Сохраняй историческую точность, упоминай даты, если они есть
+4. Цитируй конкретные факты из контекста
+
+ОТВЕТ:"""
+
 
 import os
 current_dir = os.path.dirname(os.path.abspath(__file__))