some refactor

_2_splitting.py

@@ -3,7 +3,7 @@ import re
 
 
 # Конфиги для парсинга дат
-YEARS = {
+YEARS_ALIASES = {
 'O': 1918,
 'M': 2000,
 'N': 2026
@@ -11,17 +11,15 @@ YEARS = {
 
 
 def _parse_single_year(year_str: str) -> int:
-    """
-    Парсит один год.
-    
+    """   
     Args:
-        year_str: "1962", "O", "N"
+        year_str: "1962" or alias like "O", "M", "N"
         
     Returns:
         int: Год
     """
-    if year_str in YEARS:
-        return YEARS[year_str]
+    if year_str in YEARS_ALIASES:
+        return YEARS_ALIASES[year_str]
     else:
         try:
             return int(year_str)
@@ -30,16 +28,11 @@ def _parse_single_year(year_str: str) -> int:
 
 
 def _parse_date_range(date_str: str) -> tuple[int, int]:
-    """
-    Парсит строку с датой и возвращает (start_year, end_year).
+    """Парсит строку с датой и возвращает (start_year, end_year).
     
     Поддерживает:
     - "1962-2002" -> (1962, 2002)
     - "1962" -> (1962, 1962)
-    - "O" -> (1918, 1918)
-    - "N" -> (2026, 2026)
-    - "1962-N" -> (1962, 2026)
-    - "O-1962" -> (1918, 1962)
     
     Args:
         date_str: Строка с датой
@@ -56,10 +49,8 @@ def _parse_date_range(date_str: str) -> tuple[int, int]:
         end = _parse_single_year(parts[1].strip())
         assert start <= end, f"Год начала {start} должен быть меньше или равен году конца {end}"
         return (start, end)
-    else:
-        # Один год
-        year = _parse_single_year(date_str)
-        return (year, year)
+    year = _parse_single_year(date_str)
+    return (year, year)
 
 
 def parse_metadata_from_document(text: str) -> list[tuple[str, tuple[int, int], str]]:
@@ -108,25 +99,6 @@ def parse_metadata_from_document(text: str) -> list[tuple[str, tuple[int, int],
     return result
 
 
-def years_overlap(range1: tuple[int, int], range2: tuple[int, int]) -> bool:
-    """
-    Проверяет, имеют ли два года диапазонов ненулевое пересечение (включительно).
-    
-    range1 (start1, end1) пересекается с range2 (start2, end2) если:
-    start1 <= end2 AND end1 >= start2
-    
-    Args:
-        range1: (start_year, end_year)
-        range2: (start_year, end_year)
-        
-    Returns:
-        bool: True если есть пересечение
-    """
-    start1, end1 = range1
-    start2, end2 = range2
-    return start1 <= end2 and end1 >= start2
-
-
 def process_documents(documents) -> tuple[pd.DataFrame, pd.DataFrame]:
     """
     Обрабатывает документы с парсингом дат и создает два датафрейма.

generation.py

@@ -3,7 +3,7 @@ import re
 warnings.filterwarnings('ignore')
 
 
-def lemmatize(text, vocabulary):
+def lemmatize(text: str, vocabulary: dict[str, str]) -> list[str]:
     """
     Лемматизирует текст и возвращает список терминов из словаря.
     Поддерживает поиск полных термин и частичный поиск по словам.
@@ -21,8 +21,6 @@ def lemmatize(text, vocabulary):
     
     # Приводим текст в нижний регистр для поиска
     text_lower = text.lower()
-    # Извлекаем все слова из текста
-    words_in_text = set(re.findall(r'\b\w+(?:-\w+)*\b', text_lower, re.UNICODE))
     
     # Сначала - поиск полных терминов
     for term in vocabulary.keys():
@@ -33,38 +31,11 @@ def lemmatize(text, vocabulary):
         if re.search(pattern, text_lower) and term not in found_terms_set:
             found_terms.append(term)
             found_terms_set.add(term)
-    
-    # Затем - поиск по первым компонентам составных терминов
-    # Например, "Икарус-280" будет найден по слову "Икарусы" или "икарус"
-    for term in vocabulary.keys():
-        if term in found_terms_set:
-            continue  # Пропускаем уже найденные
-            
-        term_lower = term.lower()
         
-        # Если термин содержит дефис, ищем первую часть
-        for separator in ['-', ' ']:
-            if separator in term_lower:
-                first_part = term_lower.split(separator)[0]
-                
-                # Проверяем каждое слово в тексте
-                for word in words_in_text:
-                    # Статья совпадает с первой частью полностью или является словоформой
-                    if word == first_part or word.startswith(first_part + 'ы') or \
-                       word.startswith(first_part + 'а') or word.startswith(first_part + 'у') or \
-                       word.startswith(first_part + 'е') or word.startswith(first_part + 'и') or \
-                       word.startswith(first_part + 'о'):
-                        found_terms.append(term)
-                        found_terms_set.add(term)
-                        break
-                
-                if term in found_terms_set:
-                    break
-    
     return found_terms
 
 
-def wrap_prompt(retrieved_text, query_text, inp_vocabula: dict[str, str]):
+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)

lemmatizer.py

@@ -92,10 +92,5 @@ class RussianLemmatizer:
         text = text.replace('ё', 'е').lower()
         doc = self.nlp(text)
         
-        # Извлекаем леммы, пропускаем пунктуацию
-        lemmas = []
-        for token in doc:
-            if not token.is_punct and token.lemma_.strip():  # Пропускаем пунктуацию и пробелы
-                lemmas.append(token.lemma_)
-        
-        return lemmas
+        # Извлекаем леммы, пропускаем пунктуацию и пробелы
+        return [token.lemma_ for token in doc if not token.is_punct and token.lemma_.strip()]