Improve GraphRAG extraction
This commit is contained in:
parent
40dc2ba38e
commit
9b59994d2c
@ -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
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
422
test/test_graphrag.py
Normal file
422
test/test_graphrag.py
Normal file
@ -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",
|
||||
]
|
||||
Loading…
Reference in New Issue
Block a user