a lot of refactor

tests/test_lemmatization.py

@@ -11,7 +11,7 @@ from dataclasses import dataclass
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
 from lemmatizer import RussianLemmatizer
-
+lemmatizer = RussianLemmatizer()
 
 @dataclass
 class TestSearch:
@@ -48,14 +48,13 @@ TESTS_LETTER_E = [
 ]
 
 
-def lemmas_in_sentence(query: str, sentence: str, lemmatizer: RussianLemmatizer) -> bool:
+def lemmas_in_sentence(query: str, sentence: str) -> bool:
     """
     Проверяет, есть ли леммы из query в sentence после лемматизации.
     
     Args:
         query: поисковое слово
         sentence: предложение для проверки
-        lemmatizer: инстанс лемматизатора
         
     Returns:
         True если есть пересечение лемм
@@ -67,7 +66,7 @@ def lemmas_in_sentence(query: str, sentence: str, lemmatizer: RussianLemmatizer)
     return len(query_lemmas & sentence_lemmas) > 0
 
 
-def run_test_suite(test_set_name: str, test_set: list[TestSearch], lemmatizer: RussianLemmatizer) -> tuple[int, int]:
+def run_test_suite(test_set_name: str, test_set: list[TestSearch]) -> tuple[int, int]:
     """Запускает набор тестов и возвращает (пройдено, провалено)"""
     
     print("\n" + "-"*70)
@@ -113,10 +112,7 @@ def test_lemmatization():
     print("\n" + "="*70)
     print("ТЕСТ ЛЕММАТИЗАЦИИ для русского языка")
     print("="*70)
-    
-    # Инициализируем лемматизатор один раз
-    lemmatizer = RussianLemmatizer()
-    
+       
     total_passed = 0
     total_failed = 0
     
@@ -127,7 +123,7 @@ def test_lemmatization():
         ("ОБРАБОТКА Е/Ё", TESTS_LETTER_E),
     ]:
         if test_set:  # Только если есть тесты
-            passed, failed = run_test_suite(test_name, test_set, lemmatizer)
+            passed, failed = run_test_suite(test_name, test_set)
             total_passed += passed
             total_failed += failed
     
@@ -140,13 +136,4 @@ def test_lemmatization():
 
 
 if __name__ == "__main__":
-    try:
-        success = test_lemmatization()
-        sys.exit(0 if success else 1)
-    except Exception as e:
-        print("\n" + "="*70)
-        print("✗ ОШИБКА ПРИ ВЫПОЛНЕНИИ ТЕСТОВ")
-        print("="*70)
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)
+   test_lemmatization()

tests/test_vocabular.py

@@ -0,0 +1,20 @@
+"""Тест поиска терминов с учётом падежей"""
+
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from vocabulary.parse_vocabulary import VOCABULARY_MANAGER
+
+TESTS_CASES = {
+    'Газель газель Газель газели Газелей': set(('ГАЗель', )),
+    'Икаруса-280 Икарус-260 Икарусом-260 Икарусов-280': set(('Икарус-260', 'Икарус-280')),
+    'Купили ещё автобусов': set(),
+    'Икарус-2600': set(),
+    'Иккарус-260': set(),
+    'Ока': set(('Ока', )),
+    'Окой': set(('Ока', )),
+    'Окружная дорога': set(),
+    }
+
+for text, terms in TESTS_CASES.items():
+    assert VOCABULARY_MANAGER.find_terms(text) == set((el.lower() for el in terms)), breakpoint()

vocabulary/parse_vocabulary.py

@@ -1,18 +1,50 @@
+import re
 from typing import Optional
-
-
-def parse_term(term_raw: str) -> tuple[str, str, Optional[str]]:
-    """Извлекает лемму термина, корень и окончание из строки термина"""
-    if "[" in term_raw and "]" in term_raw:
-        bracket_pos = term_raw.index("[")
-        stem = term_raw[:bracket_pos]
-        ending = term_raw[bracket_pos + 1: term_raw.index("]")]
-    else:
-        stem = term_raw
-        ending = None 
-        # важно, что не пустую строку! Есть пустое окончание, есть отсутствие окончания у несклоняемых мслов
-    term_lemma = term_raw.replace("[", "").replace("]", "")
-    return term_lemma, stem, ending
+from dataclasses import dataclass
+
+
+@dataclass
+class TermPattern:
+    """Класс для представления шаблона термина с окончанием"""
+    raw_pattern: str
+    
+    def __post_init__(self):
+        """Разбор шаблон при создании"""
+        self.base_part, self.ending_part, self.suffix = self._parse_pattern(self.raw_pattern)
+    
+    def lemma(self):
+        return self.raw_pattern.replace('[', '').replace(']', '')
+    
+    def _parse_pattern(self, pattern: str) -> tuple[str, Optional[str], str, bool]:
+        """Разбирает шаблон на составные части.
+        Возвращает (основа, окончание, суффикс, был_ли_дефис)
+        """
+        # Ищем окончание в квадратных скобках
+        match = re.search(r'\[(.*?)\]', pattern)
+        
+        if match:
+            ending = match.group(1)
+            start, end = match.span()
+            prefix = pattern[:start]
+            suffix = pattern[end:]
+            return prefix.lower(), ending.lower(), suffix.lower()
+        return pattern, None, ''
+       
+    def matches(self, word: str) -> bool:
+        """Проверяет, соответствует ли слово шаблону"""
+        word = word.lower()
+        if self.ending_part is None:
+            # Если окончания нет, слово должно точно совпадать
+            return word == self.lemma().lower()
+        if not word.startswith(self.base_part):
+            return False
+        word_without_base = word[len(self.base_part):]
+        if not self.suffix:
+            return abs((len(word_without_base)-len(self.ending_part))) <= 2
+        if not word_without_base.endswith(self.suffix):
+            return False
+        word_without_base_and_suffix = word_without_base[:-len(self.suffix)]
+        return abs((len(word_without_base_and_suffix)-len(self.ending_part))) <= 2
 
 
 class ParserVocabulary:
@@ -26,10 +58,10 @@ class ParserVocabulary:
         Определение термина
         
         - vocabulary: термин (без скобок) -> определение
-        - stems: корень -> информация об окончании
+        - patterns: список объектов TermPattern для проверки словоформ
         """
         self.vocabulary: dict[str, str] = {}
-        self.stems: dict[str, dict[str, str]] = {}  # корень -> информация об окончании
+        self.patterns: list[TermPattern] = []  # список шаблонов для проверки
         
         with open(filepath, 'r', encoding='utf-8') as f:
             lines = f.readlines()
@@ -71,38 +103,10 @@ class ParserVocabulary:
             i += 1
 
     def __parse_term_raw(self, term_raw, definition):
-        term_lemma, stem, ending = parse_term(term_raw)
-        self.vocabulary[term_lemma] = definition
+        pattern = TermPattern(term_raw)
+        self.vocabulary[pattern.lemma().lower()] = definition
+        self.patterns.append(pattern)
         
-        # Сохраняем информацию о корне для гибкого поиска
-        if stem and ending is not None:
-            if stem.lower() not in self.stems:
-                self.stems[stem.lower()] = {
-                    'canonical': term_lemma,
-                    'ending': ending
-                }
-
-    def extract_stem(self, word: str):
-        """
-        Проверяет, совпадает ли слово с каким-то корнем из stems_dict
-        с допустимым окончанием (может отличаться на 0-2 буквы)
-        
-        Args:
-            word: слово для проверки (в нижнем регистре)
-        
-        Returns:
-            Канонический вид термина или None
-        """
-        for stem, info in self.stems.items():
-            if word.startswith(stem):
-                # Вычисляем окончание в слове
-                word_ending = word[len(stem):]
-                expected_ending = info['ending']
-                
-                if abs(len(word_ending) - len(expected_ending)) <= 2:
-                    return info['canonical']
-        return None
-
     def find_terms(self, text: str) -> set[str]:
         """Возвращает список терминов в тексте (учитывая словоформы)
         
@@ -112,15 +116,20 @@ class ParserVocabulary:
         Returns:
             set: Множество найденных терминов из vocabulary в порядке их появления
         """
-        found_terms = set()  # Для отслеживания дубликатов
+        found_lemmas = set()  # Для отслеживания дубликатов
         
-        text_lower = text.lower()
+        # Очищаем текст от знаков препинания и разбиваем на слова
+        words = re.findall(r'\b\w+(?:[.-]\w+)*\b', text.lower())
         
-        for word in text_lower:
-            res = self.extract_stem(word)
-            if res:
-                found_terms.add(word)           
-        return found_terms
+        for word in words:
+            # проверяем по паттернам
+            for pattern in self.patterns:
+                if pattern.matches(word):
+                    lemma = pattern.lemma().lower()
+                    assert lemma in self.vocabulary, breakpoint()
+                    found_lemmas.add(lemma)
+                    break  # Слово найдено, переходим к следующему
+        return found_lemmas
 
     def wrap_prompt(self, retrieved_text: str, query_text: str):
         tokens_from_query = self.find_terms(query_text)
@@ -148,12 +157,8 @@ class ParserVocabulary:
 ОТВЕТ:"""
 
 
+# Создаём VOCABULARY_MANAGER, чтобы импортировать его при импортах модуля
 import os
 current_dir = os.path.dirname(os.path.abspath(__file__))
 vocab_file = os.path.join(current_dir, 'vocabulary.md')
 VOCABULARY_MANAGER = ParserVocabulary(vocab_file)
-
-
-if __name__ == '__main__':
-    print(VOCABULARY_MANAGER.vocabulary)
-    print(VOCABULARY_MANAGER.stems)

vocabulary/vocabulary.md

@@ -118,10 +118,10 @@
 ### Икарус[]-180
 Автобус особо большого класса (с гармошкой).
 
-### Икарус[]-260, Икарус[] 260
+### Икарус[]-260
 Автобус производства венгерской фирмы Ikarus второго поколения.
 
-### Икарус[]-280, Икарус[] 280, Икарус[]
+### Икарус[]-280, Икарус[]
 Автобус особо большого класса (с гармошкой), долгое время эксплуатировавшийся в Рязани и ставший символом рязанского автобуса.
 
 ### ЛиАЗ[]-5292, ЛиАЗ[]