Improve GraphRAG extraction

scripts/graphrag.py

@@ -40,6 +40,17 @@ TITLE_LABELS = {
     "navrh na nazov diplomovej prace",
 }
 
+FALLBACK_TITLE_LABELS = {
+    "bachelor_thesis": {
+        "navrh na nazov bp",
+        "navrh na nazov bakalarskej prace",
+    },
+    "diploma_thesis": {
+        "navrh na nazov dp",
+        "navrh na nazov diplomovej prace",
+    },
+}
+
 SKIP_LINE_PREFIXES = (
     "dokument:",
     "sekcia:",
@@ -411,6 +422,117 @@ def extract_work_title(
     return None
 
 
+def extract_explicit_proposal_title(
+    work_type: str,
+    chunks: list[dict[str, Any]],
+) -> tuple[str | None, list[str]]:
+    labels = FALLBACK_TITLE_LABELS.get(
+        work_type
+    )
+
+    if not labels:
+        return None, []
+
+    candidates: dict[
+        str,
+        dict[str, Any],
+    ] = {}
+
+    for chunk in sorted(
+        chunks,
+        key=lambda item: int(
+            item.get(
+                "chunk_index",
+                0,
+            )
+        ),
+    ):
+        text = str(
+            chunk.get(
+                "text",
+                "",
+            )
+        )
+
+        if not text:
+            continue
+
+        lines = text.splitlines()
+
+        for index, raw_line in enumerate(lines):
+            line = clean_markdown_line(
+                raw_line
+            )
+
+            if not line:
+                continue
+
+            before, separator, after = (
+                line.partition(":")
+            )
+
+            if not separator:
+                continue
+
+            normalized_before = normalize_text(
+                before
+            )
+
+            if normalized_before not in labels:
+                continue
+
+            title = after.strip()
+
+            if not title:
+                title = next_title_line(
+                    lines,
+                    index + 1,
+                )
+
+            if not title:
+                continue
+
+            normalized_title = normalize_text(
+                title
+            )
+
+            if not normalized_title:
+                continue
+
+            candidate = candidates.setdefault(
+                normalized_title,
+                {
+                    "title": title,
+                    "chunk_ids": set(),
+                },
+            )
+
+            chunk_id = chunk.get(
+                "chunk_id"
+            )
+
+            if chunk_id:
+                candidate[
+                    "chunk_ids"
+                ].add(
+                    str(chunk_id)
+                )
+
+    if len(candidates) != 1:
+        return None, []
+
+    candidate = next(
+        iter(candidates.values())
+    )
+
+    return (
+        str(candidate["title"]),
+        sorted(
+            candidate["chunk_ids"]
+        ),
+    )
+
+
 def extract_start_year(
     path: str,
     chunks: list[dict[str, Any]],
@@ -587,6 +709,16 @@ def extract_works(
                     str(chunk_id)
                 )
 
+    work_type_counts: dict[
+        str,
+        int,
+    ] = defaultdict(int)
+
+    for entry in grouped.values():
+        work_type_counts[
+            entry["heading"].work_type
+        ] += 1
+
     result = []
 
     for entry in grouped.values():
@@ -601,6 +733,30 @@ def extract_works(
             entry["texts"]
         )
 
+        source_chunk_ids = set(
+            entry["chunk_ids"]
+        )
+
+        if (
+            title is None
+            and work_type_counts[
+                heading.work_type
+            ] == 1
+        ):
+            (
+                fallback_title,
+                fallback_chunk_ids,
+            ) = extract_explicit_proposal_title(
+                heading.work_type,
+                chunks,
+            )
+
+            if fallback_title:
+                title = fallback_title
+                source_chunk_ids.update(
+                    fallback_chunk_ids
+                )
+
         result.append(
             {
                 "work_type": heading.work_type,
@@ -613,11 +769,7 @@ def extract_works(
                 "year": year,
                 "title": title,
                 "source_chunk_ids": sorted(
-                    set(
+                    source_chunk_ids
-                        entry[
-                            "chunk_ids"
-                        ]
-                    )
                 ),
             }
         )

test/test_graphrag.py

@@ -0,0 +1,422 @@
+from __future__ import annotations
+
+from scripts.graphrag import (
+    build_graph_payload,
+    extract_explicit_proposal_title,
+    extract_works,
+)
+
+
+JAKUB_PATH = (
+    "pages/students/2016/"
+    "jakub_maruniak/README.md"
+)
+
+
+def make_document(
+    *,
+    path: str = JAKUB_PATH,
+    title: str = "Jakub Maruniak",
+    categories: list[str] | None = None,
+) -> dict:
+    return {
+        "path": path,
+        "title": title,
+        "categories": (
+            categories
+            if categories is not None
+            else ["dp2021", "bp2019"]
+        ),
+        "tags": [
+            "spacy",
+            "ner",
+        ],
+        "author": "Daniel Hladek",
+        "published": True,
+    }
+
+
+def make_chunk(
+    index: int,
+    headings: list[str],
+    text: str,
+    *,
+    path: str = JAKUB_PATH,
+) -> dict:
+    return {
+        "chunk_id": (
+            f"{path}::chunk-{index}"
+        ),
+        "document_path": path,
+        "title": "Jakub Maruniak",
+        "categories": [
+            "dp2021",
+            "bp2019",
+        ],
+        "tags": [
+            "spacy",
+            "ner",
+        ],
+        "author": "Daniel Hladek",
+        "published": True,
+        "chunk_index": index,
+        "heading_paths": [
+            headings
+        ],
+        "text": text,
+        "text_length": len(text),
+        "token_count": 1,
+        "content_hash": f"hash-{index}",
+    }
+
+
+def test_explicit_dp_proposal_is_found() -> None:
+    chunks = [
+        make_chunk(
+            0,
+            ["Jakub Maruniak"],
+            (
+                "Dokument: Jakub Maruniak\n\n"
+                "*Návrh na názov DP*:\n\n"
+                "Anotácia a rozpoznávanie "
+                "pomenovaných entít v "
+                "slovenskom jazyku."
+            ),
+        )
+    ]
+
+    title, source_chunk_ids = (
+        extract_explicit_proposal_title(
+            "diploma_thesis",
+            chunks,
+        )
+    )
+
+    assert title == (
+        "Anotácia a rozpoznávanie "
+        "pomenovaných entít v "
+        "slovenskom jazyku."
+    )
+
+    assert source_chunk_ids == [
+        f"{JAKUB_PATH}::chunk-0"
+    ]
+
+
+def test_extract_works_uses_dp_proposal_fallback() -> None:
+    document = make_document()
+
+    chunks = [
+        make_chunk(
+            0,
+            ["Jakub Maruniak"],
+            (
+                "Dokument: Jakub Maruniak\n\n"
+                "*Rok začiatku štúdia*: 2016\n\n"
+                "*Návrh na názov DP*:\n\n"
+                "Anotácia a rozpoznávanie "
+                "pomenovaných entít v "
+                "slovenskom jazyku."
+            ),
+        ),
+        make_chunk(
+            1,
+            [
+                "Jakub Maruniak",
+                "Diplomová práca 2021",
+            ],
+            (
+                "Dokument: Jakub Maruniak\n"
+                "Sekcia: Diplomová práca 2021\n\n"
+                "Stretnutie 12.3.\n\n"
+                "Stav:\n\n"
+                "- Anotovanie dát."
+            ),
+        ),
+    ]
+
+    works = extract_works(
+        document,
+        chunks,
+    )
+
+    diploma_works = [
+        work
+        for work in works
+        if work["work_type"]
+        == "diploma_thesis"
+    ]
+
+    assert len(diploma_works) == 1
+
+    work = diploma_works[0]
+
+    assert work["year"] == 2021
+    assert work["title"] == (
+        "Anotácia a rozpoznávanie "
+        "pomenovaných entít v "
+        "slovenskom jazyku."
+    )
+
+    assert work["source_chunk_ids"] == [
+        f"{JAKUB_PATH}::chunk-0",
+        f"{JAKUB_PATH}::chunk-1",
+    ]
+
+
+def test_existing_work_title_has_priority() -> None:
+    document = make_document()
+
+    chunks = [
+        make_chunk(
+            0,
+            ["Jakub Maruniak"],
+            (
+                "Dokument: Jakub Maruniak\n\n"
+                "Návrh na názov DP: "
+                "Starší pracovný názov"
+            ),
+        ),
+        make_chunk(
+            1,
+            [
+                "Jakub Maruniak",
+                "Diplomová práca 2021",
+            ],
+            (
+                "Dokument: Jakub Maruniak\n"
+                "Sekcia: Diplomová práca 2021\n\n"
+                "Názov: Finálny názov práce"
+            ),
+        ),
+    ]
+
+    works = extract_works(
+        document,
+        chunks,
+    )
+
+    work = next(
+        work
+        for work in works
+        if work["work_type"]
+        == "diploma_thesis"
+    )
+
+    assert work["title"] == (
+        "Finálny názov práce"
+    )
+
+    assert work["source_chunk_ids"] == [
+        f"{JAKUB_PATH}::chunk-1"
+    ]
+
+
+def test_ambiguous_dp_proposals_are_not_used() -> None:
+    document = make_document()
+
+    chunks = [
+        make_chunk(
+            0,
+            ["Jakub Maruniak"],
+            (
+                "Návrh na názov DP: "
+                "Prvý návrh"
+            ),
+        ),
+        make_chunk(
+            1,
+            ["Jakub Maruniak"],
+            (
+                "Návrh na názov DP: "
+                "Druhý návrh"
+            ),
+        ),
+        make_chunk(
+            2,
+            [
+                "Jakub Maruniak",
+                "Diplomová práca 2021",
+            ],
+            (
+                "Dokument: Jakub Maruniak\n"
+                "Sekcia: Diplomová práca 2021\n\n"
+                "Stretnutie 12.3."
+            ),
+        ),
+    ]
+
+    works = extract_works(
+        document,
+        chunks,
+    )
+
+    work = next(
+        work
+        for work in works
+        if work["work_type"]
+        == "diploma_thesis"
+    )
+
+    assert work["title"] is None
+
+
+def test_proposal_is_not_used_for_diploma_project() -> None:
+    document = make_document()
+
+    chunks = [
+        make_chunk(
+            0,
+            ["Jakub Maruniak"],
+            (
+                "Návrh na názov DP: "
+                "Názov diplomovej práce"
+            ),
+        ),
+        make_chunk(
+            1,
+            [
+                "Jakub Maruniak",
+                "Diplomový projekt 2020",
+            ],
+            (
+                "Dokument: Jakub Maruniak\n"
+                "Sekcia: Diplomový projekt 2020\n\n"
+                "Stretnutie 9.3.2020\n\n"
+                "Úlohy:\n\n"
+                "- Pripraviť dáta."
+            ),
+        ),
+    ]
+
+    works = extract_works(
+        document,
+        chunks,
+    )
+
+    work = next(
+        work
+        for work in works
+        if work["work_type"]
+        == "diploma_project"
+    )
+
+    assert work["title"] is None
+
+
+def test_proposal_is_not_shared_between_multiple_dp_works() -> None:
+    document = make_document(
+        categories=[
+            "dp2021",
+            "dp2022",
+        ]
+    )
+
+    chunks = [
+        make_chunk(
+            0,
+            ["Jakub Maruniak"],
+            (
+                "Návrh na názov DP: "
+                "Jeden návrh"
+            ),
+        ),
+        make_chunk(
+            1,
+            [
+                "Jakub Maruniak",
+                "Diplomová práca 2021",
+            ],
+            (
+                "Dokument: Jakub Maruniak\n"
+                "Sekcia: Diplomová práca 2021\n\n"
+                "Stretnutie 1.1.2021"
+            ),
+        ),
+        make_chunk(
+            2,
+            [
+                "Jakub Maruniak",
+                "Diplomová práca 2022",
+            ],
+            (
+                "Dokument: Jakub Maruniak\n"
+                "Sekcia: Diplomová práca 2022\n\n"
+                "Stretnutie 1.1.2022"
+            ),
+        ),
+    ]
+
+    works = extract_works(
+        document,
+        chunks,
+    )
+
+    diploma_works = [
+        work
+        for work in works
+        if work["work_type"]
+        == "diploma_thesis"
+    ]
+
+    assert len(diploma_works) == 2
+    assert all(
+        work["title"] is None
+        for work in diploma_works
+    )
+
+
+def test_graph_payload_contains_fallback_title() -> None:
+    document = make_document()
+
+    chunks = [
+        make_chunk(
+            0,
+            ["Jakub Maruniak"],
+            (
+                "Dokument: Jakub Maruniak\n\n"
+                "Rok začiatku štúdia: 2016\n\n"
+                "Návrh na názov DP:\n\n"
+                "Anotácia a rozpoznávanie "
+                "pomenovaných entít v "
+                "slovenskom jazyku."
+            ),
+        ),
+        make_chunk(
+            1,
+            [
+                "Jakub Maruniak",
+                "Diplomová práca 2021",
+            ],
+            (
+                "Dokument: Jakub Maruniak\n"
+                "Sekcia: Diplomová práca 2021\n\n"
+                "Stretnutie 12.3."
+            ),
+        ),
+    ]
+
+    payload = build_graph_payload(
+        [document],
+        chunks,
+    )
+
+    diploma_work = next(
+        work
+        for work in payload["works"]
+        if work["work_type"]
+        == "diploma_thesis"
+    )
+
+    assert diploma_work["title"] == (
+        "Anotácia a rozpoznávanie "
+        "pomenovaných entít v "
+        "slovenskom jazyku."
+    )
+
+    assert diploma_work[
+        "source_chunk_ids"
+    ] == [
+        f"{JAKUB_PATH}::chunk-0",
+        f"{JAKUB_PATH}::chunk-1",
+    ]