add lemmatization

app.py

@@ -4,7 +4,7 @@ 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 llm import get_llm_answer
 from retrieval import Retrieval
 from _2_splitting import YEAR_OLD, YEAR_NEW
 from vocabulary.parse_vocabulary import parse_vocabulary
@@ -490,4 +490,6 @@ with gr.Blocks(title="RAG Application", theme=gr.themes.Soft()) as iface:
     )
 
 if __name__ == "__main__":
-    iface.launch(ssr_mode=False, share=True)
+    iface.launch(ssr_mode=False, 
+                #  share=True
+                 )

requirements.txt

@@ -5,4 +5,5 @@ plotly== 5.24.1
 gradio
 rank-bm25==0.2.2
 sentence-transformers==5.2.0
-huggingface_hub
+huggingface_hub
+spacy==3.8.11

retrieval.py

@@ -1,10 +1,13 @@
 import re
+from pathlib import Path
 
 import torch
 from rank_bm25 import BM25Okapi
 import warnings
 warnings.filterwarnings('ignore')
 
+import spacy
+
 from _1_get_documents import load_and_process_data
 from _2_splitting import Splitter, parse_year_metadata, years_overlap, YEAR_OLD, YEAR_NEW
 # from _3_chunking import RussianEmbedder
@@ -20,6 +23,19 @@ class Retrieval:
         print("Инициализация RAG системы...")
         self.device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
         
+        # Инициализация лемматизатора для русского языка (spaCy)
+        print("   Загрузка русской модели spaCy...")
+        try:
+            self.nlp = spacy.load("ru_core_news_sm")
+        except OSError:
+            print("   ⚠️  Модель ru_core_news_sm не найдена, скачиваю...")
+            import subprocess
+            subprocess.check_call(["python", "-m", "spacy", "download", "ru_core_news_sm"])
+            self.nlp = spacy.load("ru_core_news_sm")
+        
+        # Загружаем топонимы и регистрируем их как исключения для лемматизации
+        self._register_toponyms()
+        
         # Загружаем и обрабатываем данные
         print("1. Загрузка данных из JSON...")
         self.documents, self.docs_names = load_and_process_data()
@@ -94,6 +110,47 @@ class Retrieval:
     def get_all_chunks(self, doc_id):
         return [chunk for chunk, x in zip(self.chunks, self.docs_metadata, strict=True) if x == doc_id]
 
+    def _register_toponyms(self):
+        """
+        Загружает топонимы из vocabulary/toponims.txt и регистрирует их в spaCy
+        как custom component для исправления лемм.
+        """
+        toponims_file = Path(__file__).parent / "vocabulary" / "toponims.txt"
+        
+        if not toponims_file.exists():
+            print(f"   ⚠️  Файл {toponims_file} не найден, управление топонимами пропущено")
+            return
+        
+        # Читаем топонимы
+        with open(toponims_file, 'r', encoding='utf-8') as f:
+            self.toponims = {line.strip().lower(): line.strip().lower() for line in f if line.strip()}
+        
+        print(f"   Загружено {len(self.toponims)} топонимов")
+        
+        # Добавляем custom component для обработки топонимов после лемматизации
+        @self.nlp.component("fix_toponyms")
+        def fix_toponyms(doc):
+            """Компонент для исправления лемм топонимов и их форм"""
+            for token in doc:
+                lemma_lower = token.lemma_.lower()
+                
+                # Проверяем прямое совпадение
+                if lemma_lower in self.toponims:
+                    token.lemma_ = self.toponims[lemma_lower]
+                else:
+                    # Проверяем, может ли это быть формой топонима
+                    # Для каждого топонима проверяем, совпадает ли начало слова
+                    for toponim in self.toponims:
+                        if lemma_lower.startswith(toponim[:len(lemma_lower)-2]) and len(lemma_lower) > len(toponim) - 3:
+                            # Похоже, это форма топонима (например, "дягилев" -> "дягилево")
+                            token.lemma_ = toponim
+                            break
+            return doc
+        
+        # Добавляем компонент после лемматизатора
+        if "fix_toponyms" not in self.nlp.pipe_names:
+            self.nlp.add_pipe("fix_toponyms", after="lemmatizer")
+    
     def _prepare_bm25(self):
         """Подготавливаем BM25 индекс для ключевого поиска"""
         # Токенизация для BM25
@@ -101,10 +158,17 @@ class Retrieval:
         return BM25Okapi(tokenized_chunks)
     
     def _tokenize_text(self, text: str) -> list[str]:
-        """Простая токенизация для BM25"""
-        # Удаляем пунктуацию и приводим к нижнему регистру
-        text = re.sub(r'[^\w\s]', ' ', text.lower())
-        return text.split()
+        """Лемматизация текста для русского языка (spaCy)"""
+        # Используем spaCy для обработки текста и получения лемм
+        doc = self.nlp(text.lower())
+        
+        # Извлекаем леммы, пропускаем пунктуацию
+        lemmas = []
+        for token in doc:
+            if not token.is_punct and token.lemma_.strip():  # Пропускаем пунктуацию и пробелы
+                lemmas.append(token.lemma_)
+        
+        return lemmas
 
     def rerank_search(self, query: str) -> list[dict]:
         """

tests/test_lemmatization.py

@@ -0,0 +1,268 @@
+# -*- coding: utf-8 -*-
+"""
+Тест лемматизации для BM25 поиска на русском языке (spaCy с топонимами).
+Проверяет:
+1. Поиск по топонимам с разными падежами
+2. Лемматизацию обычных слов (включая беглые гласные)
+3. Обработку букв Е и Ё
+"""
+import re
+import sys
+from pathlib import Path
+from dataclasses import dataclass
+
+# Добавляем родительскую директорию в path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import spacy
+from rank_bm25 import BM25Okapi
+
+
+# ============================================================================
+# Тестовые данные
+# ============================================================================
+
+@dataclass
+class TestSearch:
+    word: str
+    sentence: str
+
+
+# тесты для топонимов (включая разные падежи и формы)
+TESTS_TOPONIMS = [
+    TestSearch(word="Ворошиловка", sentence="Он ездил на работу с Ворошиловки до телезавода, а обратно ходил пешком."),
+    TestSearch(word='Рязань-1', sentence='Примерно в 1969 году запускаются поезда от станции Рязани-1 до Дягилева и Леска - пара рейсов в день, но всё же.'),
+    TestSearch(word='Дягилево', sentence='Примерно в 1969 году запускаются поезда от станции Рязани-1 до Дягилева и Леска - пара рейсов в день, но всё же.'),
+    TestSearch(word='Лесок', sentence='Примерно в 1969 году запускаются поезда от станции Рязани-1 до Дягилева и Леска - пара рейсов в день, но всё же.'),
+    TestSearch(word="Дашково-Песочня", sentence="31 маршрутка ходила из Дашково-Песочни в Дягилево."),
+    TestSearch(word="Ворошиловку", sentence="Он ездил на работу с Ворошиловки до телезавода, а обратно ходил пешком."),
+    TestSearch(word='Рязани-1', sentence='Примерно в 1969 году запускаются поезда от станции Рязани-1 до Дягилева и Леска - пара рейсов в день, но всё же.'),
+    # TestSearch(word='Дягилева', sentence='Примерно в 1969 году запускаются поезда от станции Рязани-1 до Дягилева и Леска - пара рейсов в день, но всё же.'),
+    TestSearch(word='Леске', sentence='Примерно в 1969 году запускаются поезда от станции Рязани-1 до Дягилева и Леска - пара рейсов в день, но всё же.'),
+    TestSearch(word="Дашково-Песочню", sentence="31 маршрутка ходила из Дашково-Песочни в Дягилево."),
+    #TestSearch(word="Песочни", sentence="31 маршрутка ходила из Дашково-Песочни в Дягилево."),
+    #TestSearch(word='Дашково-Песочня', sentence='То есть его продлили дальше в Песочню, и он стал охватывать так называемый основной транспортный коридор города.'),    
+]
+
+# тесты для лемматизации обычных слов. Стемминг здесь сломается!
+TESTS_LEMMATIZATION = [
+    TestSearch(word='хорошая', sentence='И там в лучшем случае было 20-23 машины где-то так.'),
+    TestSearch(word='петли', sentence='И вместо трёх петель осталась одна маленькая в Горроще.') # беглая Е - чередование в корне
+]
+
+# буквы Е и Ё в топонимах и не только
+TESTS_LETTER_E = [
+    TestSearch(word='Ленпоселок', sentence='Ленпосёлок — это вообще ныне почти забытый топоним, на мой взгляд, зря.'),
+    TestSearch(word='Шлаковый посёлок', sentence='7 ноября 1977 года наконец-то был открыт путепровод на Куйбышевском шоссе, который позволил сократить время в пути автобусом и троллейбусом из центра в Шлаковый поселок.')
+]
+
+# ============================================================================
+# Класс для тестирования BM25 с лемматизацией
+# ============================================================================
+
+class RussianBM25:
+    """BM25 с лемматизац��ей для русского языка (spaCy с топонимами)"""
+    
+    def __init__(self, documents: list[str]):
+        # Загружаем русскую модель spaCy
+        try:
+            self.nlp = spacy.load("ru_core_news_sm")
+        except OSError:
+            print("Модель ru_core_news_sm не найдена, скачиваю...")
+            import subprocess
+            subprocess.check_call(["python", "-m", "spacy", "download", "ru_core_news_sm"])
+            self.nlp = spacy.load("ru_core_news_sm")
+        
+        # Загружаем и регистрируем топонимы
+        self._register_toponyms()
+        
+        tokenized_docs = [self._tokenize_text(doc) for doc in documents]
+        self.bm25 = BM25Okapi(tokenized_docs)
+        self.documents = documents
+    
+    def _register_toponyms(self):
+        """Загружает топонимы из vocabulary/toponims.txt и регистрирует их в spaCy"""
+        toponims_file = Path(__file__).parent.parent / "vocabulary" / "toponims.txt"
+        
+        if not toponims_file.exists():
+            print(f"  Файл {toponims_file} не найден, управление топонимами пропущено")
+            return
+        
+        # Читаем топонимы
+        with open(toponims_file, 'r', encoding='utf-8') as f:
+            self.toponims = {line.strip().lower(): line.strip().lower() for line in f if line.strip()}
+        
+        print(f"  Загружено {len(self.toponims)} топонимов из {toponims_file.name}")
+        
+        # Добавляем custom component для обработки топонимов после лемматизации
+        @self.nlp.component("fix_toponyms")
+        def fix_toponyms(doc):
+            """Компонент для исправления лемм топонимов и их форм"""
+            for token in doc:
+                lemma_lower = token.lemma_.lower()
+                
+                # Проверяем прямое совпадение
+                if lemma_lower in self.toponims:
+                    token.lemma_ = self.toponims[lemma_lower]
+                else:
+                    # Проверяем, может ли это быть формой топонима
+                    # Для каждого топонима проверяем, совпадает ли начало слова
+                    for toponim in self.toponims:
+                        if lemma_lower.startswith(toponim[:len(lemma_lower)-2]) and len(lemma_lower) > len(toponim) - 3:
+                            # Похоже, это форма топонима (например, "дягилев" -> "дягилево")
+                            token.lemma_ = toponim
+                            break
+            return doc
+        
+        # Добавляем компонент после лемматизатора
+        if "fix_toponyms" not in self.nlp.pipe_names:
+            self.nlp.add_pipe("fix_toponyms", after="lemmatizer")
+    
+    def _tokenize_text(self, text: str) -> list[str]:
+        """Лемматизация текста для русского языка (spaCy)"""
+        # Используем spaCy для обработки текста и получения лемм
+        doc = self.nlp(text.lower())
+        
+        # Извлекаем леммы, пропускаем пунктуацию
+        lemmas = []
+        for token in doc:
+            if not token.is_punct and token.lemma_.strip():  # Пропускаем пунктуацию и пробелы
+                lemmas.append(token.lemma_)
+        
+        return lemmas
+    
+    def search(self, query: str) -> list[tuple[int, float]]:
+        """
+        Поиск по запросу.
+        Returns: список (индекс_документа, релевантность_score)
+        """
+        tokenized_query = self._tokenize_text(query)
+        scores = self.bm25.get_scores(tokenized_query)
+        
+        # Возвращаем индексы документов, отсортированные по релевантности
+        results = [(i, score) for i, score in enumerate(scores)]
+        results.sort(key=lambda x: x[1], reverse=True)
+        return results
+
+
+# ============================================================================
+# Функция для запуска тестов
+# ============================================================================
+
+# ============================================================================
+# Функция для запуска тестов
+# ============================================================================
+
+def run_test_suite(test_set_name: str, test_set: list[TestSearch]) -> tuple[int, int]:
+    """Запускает набор тестов и возвращает (пройдено, провалено)"""
+    
+    print("\n" + "-"*70)
+    print(f"НАБОР ТЕСТОВ: {test_set_name}")
+    print("-"*70)
+    
+    # Собираем все предложения (документы)
+    sentences = [test.sentence for test in test_set]
+    
+    # Инициализируем BM25
+    bm25 = RussianBM25(sentences)
+    
+    passed = 0
+    failed = 0
+    failed_tests = []
+    
+    for i, test in enumerate(test_set, 1):
+        word = test.word
+        sentence = test.sentence
+        expected_index = sentences.index(sentence)
+        
+        # Выполняем поиск
+        results = bm25.search(word)
+        
+        # Проверяем, что правильное предложение в топе результатов
+        found_index = None
+        rank = None
+        
+        for rank_pos, (doc_index, score) in enumerate(results, 1):
+            if doc_index == expected_index:
+                found_index = doc_index
+                rank = rank_pos
+                break
+        
+        # Определяем результат теста
+        success = found_index is not None and rank <= 3  # Топ-3 результатов
+        
+        status = "✓ PASS" if success else "✗ FAIL"
+        passed += 1 if success else 0
+        failed += 0 if success else 1
+        
+        print(f"\nТест {i}: {status}")
+        print(f"  Поиск:     '{word}'")
+        print(f"  Ожидается: '{sentence[:60]}...'")
+        
+        if success:
+            print(f"  Результат: Найдено на позиции {rank} (score: {results[rank-1][1]:.4f})")
+        else:
+            # Выводим подробно, что пошло не так
+            print(f"  Результат: НЕ НАЙДЕНО В ТОП-3")
+            if found_index is not None:
+                print(f"  (Найдено на позиции {rank}, но это слишком далеко)")
+            
+            print(f"  Топ результаты (ТОП-5):")
+            for pos, (doc_idx, score) in enumerate(results[:5], 1):
+                print(f"    {pos}. score={score:.4f}: '{sentences[doc_idx][:55]}...'")
+            
+            # Анализ: какие леммы получились?
+            doc = bm25.nlp(word.lower())
+            query_tokens = [(token.text, token.lemma_) for token in doc if not token.is_punct]
+            print(f"  Анализ запроса '{word}':")
+            for token_text, token_lemma in query_tokens:
+                print(f"    '{token_text}' -> lemma='{token_lemma}'")
+            
+            # Анализ: какие леммы в предложении?
+            doc_sent = bm25.nlp(sentence.lower())
+            sent_tokens = [(token.text, token.lemma_) for token in doc_sent if not token.is_punct]
+            print(f"  Анализ предложения '{sentence[:50]}...':")
+            for token_text, token_lemma in sent_tokens[:15]:  # Показываем первые 15
+                print(f"    '{token_text}' -> lemma='{token_lemma}'")
+            
+            failed_tests.append((word, sentence, query_tokens))
+    
+    # Итоговый результат
+    print("\n" + "-"*70)
+    print(f"РЕЗУЛЬТАТ: {passed} пройдено, {failed} провалено из {len(test_set)}")
+    
+    return passed, failed
+
+
+def test_lemmatization():
+    """Запускает все наборы тестов"""
+    
+    print("\n" + "="*70)
+    print("ПОЛНЫЙ ТЕСТ ЛЕММАТИЗАЦИИ ДЛЯ BM25 (spaCy)")
+    print("="*70)
+    
+    total_passed = 0
+    total_failed = 0
+    
+    # Запускаем все наборы тестов
+    for test_name, test_set in [
+        ("ТОПОНИМЫ (падежи и словоформы)", TESTS_TOPONIMS),
+        ("ЛЕММАТИЗАЦИЯ (беглые гласные)", TESTS_LEMMATIZATION),
+        ("ОБРАБОТКА Е/Ё", TESTS_LETTER_E),
+    ]:
+        if test_set:  # Только если есть тесты
+            passed, failed = run_test_suite(test_name, test_set)
+            total_passed += passed
+            total_failed += failed
+    
+    # Финальный результат
+    print("\n" + "="*70)
+    print(f"ИТОГО: {total_passed} пройдено, {total_failed} провалено из {total_passed + total_failed}")
+    print("="*70 + "\n")
+    
+    return total_failed == 0
+
+
+if __name__ == "__main__":
+    success = test_lemmatization()
+    sys.exit(0 if success else 1)

vocabulary/toponims.txt

@@ -0,0 +1,27 @@
+Рязань
+Семчино
+Недостоево
+Канищево
+Приокский
+Дягилево
+Московский
+Мервино
+Ворошиловка
+Троицкая слобода
+Ситники
+Октябрьский городок
+Дашки
+Ленпосёлок
+Горроща
+Посёлок Южный
+Борки
+Солотча
+Песочня
+Дашково-Песочня
+Шлаковый посёлок
+Соколовка
+Мирный
+Посёлок Строитель
+Рыбное
+Турлатово
+Дядьково