Обновить

Комментарии 1

Нашёл конкретную поломку из списка «сырого»: MarkdownHeaderTextSplitter запускается для каждой страницы отдельно, поэтому на продолжении параграфа заголовка уже нет и фильтр по нему отбрасывает фрагмент.

В статье нет ссылки на репозиторий, поэтому проверил только этот фрагментный фикс, не весь бот. Заголовки можно переносить между страницами так:

HEADER_LEVELS = ("Глава", "Параграф", "Подраздел")


def carry_header_metadata(chunks, active=None):
    current = dict(active or {})
    for chunk in chunks:
        metadata = chunk.metadata
        if "Глава" in metadata:
            current = {"Глава": metadata["Глава"]}
        if "Параграф" in metadata:
            current["Параграф"] = metadata["Параграф"]
            current.pop("Подраздел", None)
        if "Подраздел" in metadata:
            current["Подраздел"] = metadata["Подраздел"]
        for key in HEADER_LEVELS:
            if key not in metadata and key in current:
                metadata[key] = current[key]
    return current

active_headers = {}
for page_dict in pages_data:
    page_text = page_dict["text"]
    page_num = page_dict["metadata"].get("page", 0) + 1
    md_chunks = markdown_splitter.split_text(page_text)
    active_headers = carry_header_metadata(md_chunks, active_headers)
    safe_chunks = text_splitter.split_documents(md_chunks)
    for chunk in safe_chunks:
        chunk.metadata["Страница"] = page_num
        chunk.metadata["Класс"] = grade
        all_safe_chunks.append(chunk)

Три теста на продолжение страницы, смену главы и смену параграфа прошли (3/3):

from dataclasses import dataclass, field
from page_metadata import carry_header_metadata

@dataclass
class Chunk:
    metadata: dict = field(default_factory=dict)

def test_page_continuation_keeps_paragraph_and_own_page():
    first = [Chunk({"Глава": "Трудовое право", "Параграф": "§ 23"})]
    active = carry_header_metadata(first)
    second = [Chunk({"Страница": 140})]
    carry_header_metadata(second, active)
    assert second[0].metadata == {
        "Страница": 140, "Глава": "Трудовое право", "Параграф": "§ 23"
    }

def test_new_chapter_clears_old_paragraph():
    active = {"Глава": "Трудовое право", "Параграф": "§ 23", "Подраздел": "Стороны"}
    chunks = [Chunk({"Глава": "Семейное право"})]
    assert carry_header_metadata(chunks, active) == {"Глава": "Семейное право"}

def test_new_paragraph_clears_old_subsection():
    active = {"Глава": "Трудовое право", "Параграф": "§ 23", "Подраздел": "Стороны"}
    chunks = [Chunk({"Параграф": "§ 24"})]
    assert carry_header_metadata(chunks, active) == {
        "Глава": "Трудовое право", "Параграф": "§ 24"
    }
Зарегистрируйтесь на Хабре, чтобы оставить комментарий

Публикации