diff --git a/scripts/graphrag.py b/scripts/graphrag.py index 27f52ee..d32d3df 100644 --- a/scripts/graphrag.py +++ b/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( - entry[ - "chunk_ids" - ] - ) + source_chunk_ids ), } ) diff --git a/test/test_graphrag.py b/test/test_graphrag.py new file mode 100644 index 0000000..d03a23d --- /dev/null +++ b/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", + ]