parse_documents.py

import pandas as pd
import re


# Конфиги для парсинга дат
YEARS_ALIASES = {
'O': 1918,
'M': 2000,
'N': 2026
}


def _parse_single_year(year_str: str) -> int:
    """   
    Args:
        year_str: "1962" or alias like "O", "M", "N"
        
    Returns:
        int: Год
    """
    if year_str in YEARS_ALIASES:
        return YEARS_ALIASES[year_str]
    else:
        try:
            return int(year_str)
        except ValueError:
            raise ValueError(f"Невозможно распарсить год: {year_str}")


def _parse_date_range(date_str: str) -> tuple[int, int]:
    """Парсит строку с датой и возвращает (start_year, end_year).
    
    Поддерживает:
    - "1962-2002" -> (1962, 2002)
    - "1962" -> (1962, 1962)
    
    Args:
        date_str: Строка с датой
        
    Returns:
        tuple: (start_year, end_year)
    """
    date_str = date_str.strip()
    
    # Если содержит дефис
    if '-' in date_str:
        parts = date_str.split('-')
        start = _parse_single_year(parts[0].strip())
        end = _parse_single_year(parts[1].strip())
        assert start <= end, f"Год начала {start} должен быть меньше или равен году конца {end}"
        return (start, end)
    year = _parse_single_year(date_str)
    return (year, year)


def parse_metadata_from_document(text: str) -> list[tuple[str, tuple[int, int], str]]:
    """Парсит markdown текст и возвращает список (чанк_текста, (год_начала, год_конца), summary).
    
    Формат разметки ОБЯЗАТЕЛЕН:
    - ## Summary text - заголовок summary (двойной хэш + пробел)
    - ### 1962-2002 - заголовок с годом (тройной хэш + пробел)
    
    Правила:
    - Каждый документ ДОЛЖЕН начинаться с "## {summary}"
    - После summary ДОЛЖНЫ быть заголовки "### {годы}" с текстом
    - ## распространяется на все абзацы ниже до следующего ## или конца файла
    - ### распространяется на абзацы ниже до следующего ### или ##
    - Текст БЕЗ предшествующего ### Не добавляется в результат
    
    Args:
        text: Полный текст документа
        
    Returns:
        list: [(chunk_text, (start_year, end_year), summary), ...]
    """
    lines = text.split('\n')
    
    result = []
    current_summary = None
    current_year_range = None
    
    for line in lines:
        strip_line = line.strip()
        if not strip_line:
            continue
        
        if strip_line.startswith('## '):   
            current_summary = strip_line[3:].strip()  # Пропускаем "## "
        
        # Проверяем, является ли строка "### " (год с пробелом после)
        elif strip_line.startswith('### '):
            current_year_range = _parse_date_range(strip_line[4:])
        
        else:
            # Добавляем текст только если у нас есть год
            assert current_year_range and current_summary, breakpoint()
            result.append((line, current_year_range, current_summary))
    
    return result


def process_documents(documents) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Обрабатывает документы с парсингом дат и создает два датафрейма.
    
    Returns:
        tuple: (paragraphs_df, chunks_df)
        
        paragraphs_df:
            - paragraph_id: уникальный идентификатор абзаца
            - summary: название документа/раздела
            - start_year: год начала периода
            - end_year: год окончания периода
            - text: текст абзаца
            - document_id: ссылка на исходный документ
        
        chunks_df:
            - chunk_id: уникальный идентификатор чанка
            - paragraph_id: ссылка на абзац (foreign key)
            - text: текст чанка
            - lemmatized_text: лемматизированный текст (добавляется позже)
    """
    paragraphs_data = []
    chunks_data = []
    
    paragraph_id_counter = 0
    chunk_id_counter = 0
    
    for doc_id, document in enumerate(documents):
        dated_chunks = parse_metadata_from_document(document)
        
        for chunk_text, year_range, summary in dated_chunks:
            paragraphs = chunk_text.split('\n')
            
            for paragraph in paragraphs:
                paragraph = paragraph.strip()
                
                # Добавляем информацию о параграфе в датафрейм параграфов
                paragraphs_data.append({
                    'paragraph_id': paragraph_id_counter,
                    'summary': summary,
                    'start_year': year_range[0],
                    'end_year': year_range[1],
                    'text': paragraph,
                    'document_id': doc_id
                })
                
                # Разбиваем параграф на предложения и создаем чанки
                sentences = re.split(r'(?<=[.!?])\s+', paragraph)                    
                for sent in sentences:
                    chunks_data.append({
                        'chunk_id': chunk_id_counter,
                        'paragraph_id': paragraph_id_counter,
                        'text': sent.strip()
                    })
                    chunk_id_counter += 1
                
                paragraph_id_counter += 1
    
    # Создаем датафреймы
    paragraphs_df = pd.DataFrame(paragraphs_data)
    chunks_df = pd.DataFrame(chunks_data)
    
    print(f"Создано {len(chunks_df)} чанков")
    print(f"Из {len(paragraphs_df)} абзацев в {len(set(paragraphs_df['document_id']))} документах")
    
    return paragraphs_df, chunks_df