diff --git a/scripts/build_graphrag.py b/scripts/build_graphrag.py new file mode 100644 index 0000000..af12487 --- /dev/null +++ b/scripts/build_graphrag.py @@ -0,0 +1,463 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from scripts.graph_db import ( + create_neo4j_driver, + get_neo4j_settings, +) +from scripts.graphrag import ( + GRAPH_SCOPE, + build_graph_payload, +) + + +PROJECT_ROOT = Path( + __file__ +).resolve().parents[1] + +DOCUMENTS_FILE = ( + PROJECT_ROOT + / "data" + / "documents.json" +) + +CHUNKS_FILE = ( + PROJECT_ROOT + / "data" + / "chunks.json" +) + + +CONSTRAINTS = ( + ( + "zpwiki_document_path_unique", + "Document", + "path", + ), + ( + "zpwiki_person_id_unique", + "Person", + "id", + ), + ( + "zpwiki_author_id_unique", + "Author", + "id", + ), + ( + "zpwiki_category_name_unique", + "Category", + "name", + ), + ( + "zpwiki_topic_id_unique", + "Topic", + "id", + ), + ( + "zpwiki_work_id_unique", + "Work", + "id", + ), +) + + +NODE_QUERIES = { + "people": """ + UNWIND $rows AS row + MERGE (n:Person {id: row.id}) + SET n += row + """, + "authors": """ + UNWIND $rows AS row + MERGE (n:Author {id: row.id}) + SET n += row + """, + "documents": """ + UNWIND $rows AS row + MERGE (n:Document {path: row.path}) + SET n += row + """, + "categories": """ + UNWIND $rows AS row + MERGE (n:Category {name: row.name}) + SET n += row + """, + "topics": """ + UNWIND $rows AS row + MERGE (n:Topic {id: row.id}) + SET n += row + """, + "works": """ + UNWIND $rows AS row + MERGE (n:Work {id: row.id}) + SET n += row + """, +} + + +RELATIONSHIP_QUERIES = { + "person_documents": """ + UNWIND $rows AS row + MATCH (a:Person {id: row.person_id}) + MATCH (b:Document { + path: row.document_path + }) + MERGE (a)-[:HAS_DOCUMENT]->(b) + """, + "author_documents": """ + UNWIND $rows AS row + MATCH (a:Author {id: row.author_id}) + MATCH (b:Document { + path: row.document_path + }) + MERGE (a)-[:AUTHORED]->(b) + """, + "document_categories": """ + UNWIND $rows AS row + MATCH (a:Document { + path: row.document_path + }) + MATCH (b:Category { + name: row.category_name + }) + MERGE (a)-[:IN_CATEGORY]->(b) + """, + "document_topics": """ + UNWIND $rows AS row + MATCH (a:Document { + path: row.document_path + }) + MATCH (b:Topic { + id: row.topic_id + }) + MERGE (a)-[:HAS_TOPIC]->(b) + """, + "document_describes_topics": """ + UNWIND $rows AS row + MATCH (a:Document { + path: row.document_path + }) + MATCH (b:Topic { + id: row.topic_id + }) + MERGE (a)-[:DESCRIBES]->(b) + """, + "person_works": """ + UNWIND $rows AS row + MATCH (a:Person { + id: row.person_id + }) + MATCH (b:Work { + id: row.work_id + }) + MERGE (a)-[:HAS_WORK]->(b) + """, + "work_documents": """ + UNWIND $rows AS row + MATCH (a:Work { + id: row.work_id + }) + MATCH (b:Document { + path: row.document_path + }) + MERGE (a)-[:EVIDENCED_BY]->(b) + """, +} + + +def load_json( + path: Path, +) -> list[dict[str, Any]]: + if not path.exists(): + raise FileNotFoundError( + f"Missing file: {path}" + ) + + data = json.loads( + path.read_text( + encoding="utf-8" + ) + ) + + if not isinstance( + data, + list, + ): + raise ValueError( + f"Expected list in {path}" + ) + + return data + + +def create_constraints( + session, +) -> None: + for ( + name, + label, + property_name, + ) in CONSTRAINTS: + session.run( + ( + f"CREATE CONSTRAINT " + f"{name} IF NOT EXISTS " + f"FOR (n:{label}) " + f"REQUIRE n.{property_name} " + f"IS UNIQUE" + ) + ).consume() + + +def clear_managed_graph( + session, +) -> None: + session.run( + """ + MATCH (n) + WHERE n.graph_scope = $scope + DETACH DELETE n + """, + scope=GRAPH_SCOPE, + ).consume() + + +def write_rows( + session, + query: str, + rows: list[dict[str, Any]], +) -> None: + if not rows: + return + + session.run( + query, + rows=rows, + ).consume() + + +def graph_counts( + session, +) -> tuple[ + list[tuple[str, int]], + list[tuple[str, int]], +]: + node_rows = session.run( + """ + MATCH (n) + WHERE n.graph_scope = $scope + UNWIND labels(n) AS label + RETURN label, count(*) AS count + ORDER BY label + """, + scope=GRAPH_SCOPE, + ) + + nodes = [ + ( + str(record["label"]), + int(record["count"]), + ) + for record in node_rows + ] + + relationship_rows = session.run( + """ + MATCH (a)-[r]->(b) + WHERE + a.graph_scope = $scope + AND b.graph_scope = $scope + RETURN + type(r) AS relationship, + count(*) AS count + ORDER BY relationship + """, + scope=GRAPH_SCOPE, + ) + + relationships = [ + ( + str( + record[ + "relationship" + ] + ), + int(record["count"]), + ) + for record in relationship_rows + ] + + return nodes, relationships + + +def work_quality_counts( + session, +) -> tuple[int, int]: + record = session.run( + """ + MATCH (w:Work) + WHERE w.graph_scope = $scope + RETURN + count(w) AS total, + count(w.title) AS with_title + """, + scope=GRAPH_SCOPE, + ).single() + + if record is None: + return 0, 0 + + return ( + int(record["total"]), + int(record["with_title"]), + ) + + +def print_payload_summary( + payload: dict[ + str, + list[dict[str, Any]], + ], +) -> None: + print() + print("Prepared GraphRAG payload") + print("=" * 60) + + for key in ( + "people", + "authors", + "documents", + "categories", + "topics", + "works", + ): + print( + f"{key:<24}" + f"{len(payload[key]):>6}" + ) + + print("=" * 60) + + +def print_database_summary( + session, +) -> None: + nodes, relationships = ( + graph_counts( + session + ) + ) + + print() + print("Neo4j GraphRAG knowledge graph") + print("=" * 60) + + print("Nodes") + + for label, count in nodes: + print( + f" {label:<22}" + f"{count:>6}" + ) + + print() + print("Relationships") + + for relationship, count in relationships: + print( + f" {relationship:<22}" + f"{count:>6}" + ) + + total, with_title = ( + work_quality_counts( + session + ) + ) + + print() + print( + "Works with extracted title: " + f"{with_title}/{total}" + ) + + print("=" * 60) + + +def main() -> None: + parser = argparse.ArgumentParser() + + parser.add_argument( + "--replace", + action=argparse.BooleanOptionalAction, + default=True, + ) + + args = parser.parse_args() + + documents = load_json( + DOCUMENTS_FILE + ) + + chunks = load_json( + CHUNKS_FILE + ) + + payload = build_graph_payload( + documents, + chunks, + ) + + print_payload_summary( + payload + ) + + settings = get_neo4j_settings() + + with create_neo4j_driver( + settings + ) as driver: + driver.verify_connectivity() + + with driver.session( + database=settings.database + ) as session: + create_constraints( + session + ) + + if args.replace: + clear_managed_graph( + session + ) + + for key, query in ( + NODE_QUERIES.items() + ): + write_rows( + session, + query, + payload[key], + ) + + for key, query in ( + RELATIONSHIP_QUERIES.items() + ): + write_rows( + session, + query, + payload[key], + ) + + print_database_summary( + session + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/graphrag.py b/scripts/graphrag.py new file mode 100644 index 0000000..27f52ee --- /dev/null +++ b/scripts/graphrag.py @@ -0,0 +1,1048 @@ +from __future__ import annotations + +import re +import unicodedata +from collections import defaultdict +from dataclasses import dataclass +from typing import Any + + +GRAPH_SCOPE = "zpwiki" + +WORK_TYPE_LABELS = { + "bachelor_thesis": "Bakalárska práca", + "diploma_thesis": "Diplomová práca", + "diploma_project": "Diplomový projekt", + "team_project": "Tímový projekt", + "research_project": "Výskumný projekt", + "doctoral_thesis": "Dizertačná práca", +} + +WORK_HEADING_PATTERNS = ( + ("bachelor_thesis", "bakalarska praca", "bp"), + ("diploma_thesis", "diplomova praca", "dp"), + ("diploma_project", "diplomovy projekt", None), + ("team_project", "timovy projekt", "tp"), + ("research_project", "vyskumny projekt", "vp"), + ("doctoral_thesis", "dizertacna praca", "phd"), +) + +TITLE_LABELS = { + "nazov", + "nazov prace", + "nazov bakalarskej prace", + "nazov diplomovej prace", + "tema", + "navrh na nazov", + "navrh na nazov bp", + "navrh na nazov dp", + "navrh na nazov bakalarskej prace", + "navrh na nazov diplomovej prace", +} + +SKIP_LINE_PREFIXES = ( + "dokument:", + "sekcia:", + "rok zaciatku studia", + "meno veduceho", + "zadanie", + "navrh na zadanie", + "navrh na zadamie", + "ciel:", + "ciele:", + "ulohy:", + "zasobnik uloh", + "stretnutie", + "stav:", + "repozitar", + "git repozitar", + "crzp", +) + + +@dataclass(frozen=True) +class WorkHeading: + work_type: str + heading: str + category_prefix: str | None + + +def normalize_text(value: str) -> str: + decomposed = unicodedata.normalize( + "NFKD", + str(value), + ) + + without_marks = "".join( + char + for char in decomposed + if not unicodedata.combining(char) + ) + + return re.sub( + r"\s+", + " ", + without_marks.lower(), + ).strip() + + +def normalize_key(value: str) -> str: + normalized = normalize_text(value) + + normalized = re.sub( + r"[^a-z0-9]+", + "-", + normalized, + ) + + return normalized.strip("-") + + +def source_url_from_path(path: str) -> str: + normalized = path.replace("\\", "/") + + if normalized.startswith("pages/"): + normalized = normalized[len("pages/"):] + + if normalized.endswith("/README.md"): + normalized = normalized[:-len("/README.md")] + elif normalized.endswith(".md"): + normalized = normalized[:-3] + + return ( + "https://zp.kemt.fei.tuke.sk/" + + normalized.strip("/") + ) + + +def document_kind(path: str) -> str: + normalized = path.replace("\\", "/") + + if normalized.startswith("pages/students/"): + return "student" + + if normalized.startswith("pages/interns/"): + return "intern" + + if normalized.startswith("pages/topics/"): + return "topic" + + return "other" + + +def extract_full_years(value: str) -> list[int]: + years = [ + int(match) + for match in re.findall( + r"\b(20\d{2})\b", + value, + ) + ] + + for match in re.finditer( + r"\b(20\d{2})\s*/\s*(\d{2})\b", + value, + ): + first = int(match.group(1)) + second_short = int(match.group(2)) + + second = ( + (first // 100) * 100 + + second_short + ) + + if second not in years: + years.append(second) + + return sorted(set(years)) + + +def parse_category( + category: str, +) -> tuple[str, int | None]: + match = re.fullmatch( + r"(bp|dp|tp|vp|phd)(20\d{2})", + category.lower(), + ) + + if not match: + return "other", None + + prefix = match.group(1) + year = int(match.group(2)) + + category_types = { + "bp": "bachelor_thesis", + "dp": "diploma_thesis", + "tp": "team_project", + "vp": "research_project", + "phd": "doctoral_thesis", + } + + return category_types[prefix], year + + +def find_work_heading( + heading: str, +) -> WorkHeading | None: + normalized = normalize_text(heading) + + for ( + work_type, + phrase, + category_prefix, + ) in WORK_HEADING_PATTERNS: + if phrase in normalized: + return WorkHeading( + work_type=work_type, + heading=heading, + category_prefix=category_prefix, + ) + + return None + + +def category_years( + categories: list[str], + prefix: str | None, +) -> list[int]: + if prefix is None: + return [] + + pattern = re.compile( + rf"^{re.escape(prefix)}(20\d{{2}})$", + flags=re.IGNORECASE, + ) + + result = [] + + for category in categories: + match = pattern.fullmatch(category) + + if match: + result.append( + int(match.group(1)) + ) + + return sorted(set(result)) + + +def resolve_work_year( + heading: WorkHeading, + categories: list[str], +) -> int | None: + heading_years = extract_full_years( + heading.heading + ) + + category_candidates = category_years( + categories, + heading.category_prefix, + ) + + if heading.work_type == "diploma_project": + if heading_years: + return heading_years[-1] + + return None + + common = sorted( + set(heading_years) + & set(category_candidates) + ) + + if common: + return common[-1] + + if len(heading_years) > 1: + if len(category_candidates) == 1: + return category_candidates[0] + + return heading_years[-1] + + if heading_years: + return heading_years[-1] + + if category_candidates: + return category_candidates[-1] + + return None + + +def clean_markdown_line( + line: str, +) -> str: + value = line.strip() + + value = value.replace( + "*", + "", + ).replace( + "`", + "", + ) + + return value.strip() + + +def next_title_line( + lines: list[str], + start_index: int, +) -> str | None: + for raw_line in lines[start_index:]: + line = clean_markdown_line( + raw_line + ) + + if not line: + continue + + normalized = normalize_text(line) + + if normalized.startswith( + SKIP_LINE_PREFIXES + ): + continue + + if line.startswith( + ("http://", "https://", "[") + ): + continue + + if re.match( + r"^[-•]\s+", + line, + ): + continue + + if re.match( + r"^\d+[\.\)]\s*", + line, + ): + continue + + return line + + return None + + +def extract_work_title( + texts: list[str], +) -> str | None: + for text in texts: + 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(":") + ) + + normalized_before = normalize_text( + before + ) + + if ( + separator + and normalized_before + in TITLE_LABELS + ): + title = after.strip() + + if title: + return title + + candidate = next_title_line( + lines, + index + 1, + ) + + if candidate: + return candidate + + for text in texts: + lines = text.splitlines() + + for raw_line in lines: + line = clean_markdown_line( + raw_line + ) + + if not line: + continue + + normalized = normalize_text( + line + ) + + if normalized.startswith( + SKIP_LINE_PREFIXES + ): + continue + + if line.startswith( + ("http://", "https://", "[") + ): + continue + + if re.match( + r"^[-•]\s+", + line, + ): + continue + + if re.match( + r"^\d+[\.\)]\s*", + line, + ): + continue + + if len(line) < 8: + continue + + return line + + return None + + +def extract_start_year( + path: str, + chunks: list[dict[str, Any]], +) -> int | None: + pattern = re.compile( + r"rok\s+zaciatku\s+studia" + r"\s*\*?\s*:\s*\*?\s*" + r"(20\d{2})" + ) + + for chunk in chunks: + normalized = normalize_text( + str(chunk.get("text", "")) + ) + + match = pattern.search( + normalized + ) + + if match: + return int( + match.group(1) + ) + + match = re.search( + r"pages/students/(20\d{2})/", + path.replace("\\", "/"), + ) + + if match: + return int( + match.group(1) + ) + + return None + + +def topic_key_from_document_path( + path: str, +) -> str | None: + normalized = path.replace("\\", "/") + + if not normalized.startswith( + "pages/topics/" + ): + return None + + value = normalized[ + len("pages/topics/"): + ] + + if value.endswith("/README.md"): + value = value[:-len("/README.md")] + elif value.endswith(".md"): + value = value[:-3] + + if not value: + return None + + return normalize_key( + value.split("/")[0] + ) + + +def topic_display_name( + value: str, +) -> str: + aliases = { + "nlp": "NLP", + "rag": "RAG", + "ner": "NER", + "qa": "QA", + "ie": "IE", + "lm": "LM", + "nmt": "NMT", + "pos": "POS", + "gpu": "GPU", + "spacy": "Spacy", + } + + key = normalize_key(value) + + if key in aliases: + return aliases[key] + + return str(value).strip() + + +def extract_works( + document: dict[str, Any], + chunks: list[dict[str, Any]], +) -> list[dict[str, Any]]: + grouped: dict[ + tuple[str, str], + dict[str, Any], + ] = {} + + categories = list( + document.get( + "categories", + [], + ) + ) + + for chunk in sorted( + chunks, + key=lambda item: int( + item.get( + "chunk_index", + 0, + ) + ), + ): + heading_paths = chunk.get( + "heading_paths", + [], + ) + + candidates = [] + + for path in heading_paths: + if isinstance(path, list): + candidates.extend( + str(value) + for value in path + ) + elif path: + candidates.append( + str(path) + ) + + for raw_heading in candidates: + heading = find_work_heading( + raw_heading + ) + + if heading is None: + continue + + key = ( + heading.work_type, + normalize_text( + heading.heading + ), + ) + + entry = grouped.setdefault( + key, + { + "heading": heading, + "texts": [], + "chunk_ids": [], + }, + ) + + text = str( + chunk.get( + "text", + "", + ) + ) + + if text: + entry["texts"].append( + text + ) + + chunk_id = chunk.get( + "chunk_id" + ) + + if chunk_id: + entry["chunk_ids"].append( + str(chunk_id) + ) + + result = [] + + for entry in grouped.values(): + heading = entry["heading"] + + year = resolve_work_year( + heading, + categories, + ) + + title = extract_work_title( + entry["texts"] + ) + + result.append( + { + "work_type": heading.work_type, + "work_type_label": ( + WORK_TYPE_LABELS[ + heading.work_type + ] + ), + "heading": heading.heading, + "year": year, + "title": title, + "source_chunk_ids": sorted( + set( + entry[ + "chunk_ids" + ] + ) + ), + } + ) + + return sorted( + result, + key=lambda item: ( + item["year"] or 0, + item["work_type"], + item["heading"], + ), + ) + + +def build_graph_payload( + documents: list[dict[str, Any]], + chunks: list[dict[str, Any]], +) -> dict[str, list[dict[str, Any]]]: + chunks_by_document: dict[ + str, + list[dict[str, Any]], + ] = defaultdict(list) + + for chunk in chunks: + path = str( + chunk.get( + "document_path", + "", + ) + ) + + if path: + chunks_by_document[ + path + ].append(chunk) + + people: dict[ + str, + dict[str, Any], + ] = {} + + authors: dict[ + str, + dict[str, Any], + ] = {} + + document_nodes: dict[ + str, + dict[str, Any], + ] = {} + + categories: dict[ + str, + dict[str, Any], + ] = {} + + topics: dict[ + str, + dict[str, Any], + ] = {} + + works: dict[ + str, + dict[str, Any], + ] = {} + + person_documents = set() + author_documents = set() + document_categories = set() + document_topics = set() + document_describes_topics = set() + person_works = set() + work_documents = set() + + for document in documents: + path = str( + document["path"] + ) + + title = str( + document.get( + "title", + path, + ) + ) + + kind = document_kind( + path + ) + + doc_categories = [ + str(value) + for value in document.get( + "categories", + [], + ) + ] + + doc_tags = [ + str(value) + for value in document.get( + "tags", + [], + ) + ] + + author = document.get( + "author" + ) + + document_nodes[path] = { + "path": path, + "name": title, + "title": title, + "source_url": ( + source_url_from_path( + path + ) + ), + "document_kind": kind, + "published": bool( + document.get( + "published", + False, + ) + ), + "categories": ( + doc_categories + ), + "tags": doc_tags, + "author": ( + str(author) + if author + else None + ), + "graph_scope": ( + GRAPH_SCOPE + ), + } + + if author: + author_name = str( + author + ) + + author_id = normalize_key( + author_name + ) + + authors[author_id] = { + "id": author_id, + "name": author_name, + "graph_scope": ( + GRAPH_SCOPE + ), + } + + author_documents.add( + ( + author_id, + path, + ) + ) + + for category in doc_categories: + category_type, year = ( + parse_category( + category + ) + ) + + categories[category] = { + "name": category, + "category_type": ( + category_type + ), + "year": year, + "graph_scope": ( + GRAPH_SCOPE + ), + } + + document_categories.add( + ( + path, + category, + ) + ) + + for tag in doc_tags: + topic_id = normalize_key( + tag + ) + + current = topics.get( + topic_id + ) + + if current is None: + topics[topic_id] = { + "id": topic_id, + "name": ( + topic_display_name( + tag + ) + ), + "graph_scope": ( + GRAPH_SCOPE + ), + } + + document_topics.add( + ( + path, + topic_id, + ) + ) + + described_topic = ( + topic_key_from_document_path( + path + ) + ) + + if described_topic: + topics[ + described_topic + ] = { + "id": described_topic, + "name": title, + "graph_scope": ( + GRAPH_SCOPE + ), + } + + document_describes_topics.add( + ( + path, + described_topic, + ) + ) + + if kind not in { + "student", + "intern", + }: + continue + + person_id = ( + f"person:{path}" + ) + + document_chunks = ( + chunks_by_document.get( + path, + [], + ) + ) + + people[person_id] = { + "id": person_id, + "name": title, + "person_kind": kind, + "start_year": ( + extract_start_year( + path, + document_chunks, + ) + if kind == "student" + else None + ), + "source_document": path, + "graph_scope": ( + GRAPH_SCOPE + ), + } + + person_documents.add( + ( + person_id, + path, + ) + ) + + if kind != "student": + continue + + for work in extract_works( + document, + document_chunks, + ): + work_id = ( + "work:" + + path + + ":" + + normalize_key( + work["heading"] + ) + ) + + display_name = ( + work["title"] + or work["heading"] + ) + + works[work_id] = { + "id": work_id, + "name": display_name, + "title": work[ + "title" + ], + "heading": work[ + "heading" + ], + "work_type": work[ + "work_type" + ], + "work_type_label": work[ + "work_type_label" + ], + "year": work[ + "year" + ], + "source_document": path, + "source_chunk_ids": work[ + "source_chunk_ids" + ], + "graph_scope": ( + GRAPH_SCOPE + ), + } + + person_works.add( + ( + person_id, + work_id, + ) + ) + + work_documents.add( + ( + work_id, + path, + ) + ) + + def relation_rows( + pairs: set[tuple[str, str]], + source_name: str, + target_name: str, + ) -> list[dict[str, str]]: + return [ + { + source_name: source, + target_name: target, + } + for source, target in sorted( + pairs + ) + ] + + return { + "people": sorted( + people.values(), + key=lambda row: row["id"], + ), + "authors": sorted( + authors.values(), + key=lambda row: row["id"], + ), + "documents": sorted( + document_nodes.values(), + key=lambda row: row[ + "path" + ], + ), + "categories": sorted( + categories.values(), + key=lambda row: row[ + "name" + ], + ), + "topics": sorted( + topics.values(), + key=lambda row: row["id"], + ), + "works": sorted( + works.values(), + key=lambda row: row["id"], + ), + "person_documents": relation_rows( + person_documents, + "person_id", + "document_path", + ), + "author_documents": relation_rows( + author_documents, + "author_id", + "document_path", + ), + "document_categories": relation_rows( + document_categories, + "document_path", + "category_name", + ), + "document_topics": relation_rows( + document_topics, + "document_path", + "topic_id", + ), + "document_describes_topics": ( + relation_rows( + document_describes_topics, + "document_path", + "topic_id", + ) + ), + "person_works": relation_rows( + person_works, + "person_id", + "work_id", + ), + "work_documents": relation_rows( + work_documents, + "work_id", + "document_path", + ), + }