From e555acb4b04f1638b200cc3b7a22e04029f4284c Mon Sep 17 00:00:00 2001 From: matheu-s Date: Thu, 20 Mar 2025 23:52:16 +0200 Subject: [PATCH] Add graph visualization feature --- README.md | 10 ++++++++++ visualize_graph.py | 25 +++++++++++++++++++++++++ 2 files changed, 35 insertions(+) create mode 100644 visualize_graph.py diff --git a/README.md b/README.md index e7ad841..9c02fb6 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,16 @@ for file_name in txt_files: rag.insert(file.read()) ``` +## Graph Visualization +Once your files were processed, you may visualize how the graph looks like. +1. Ensure the file `graph_chunk_entity_relation.graphml` has been written during the main indexation. +2. Run `visualize_graph.py`. + +It will create an HTML of your graph. Feel free to open it on your favorite browser. +* The edge labels are the relationship strength. +* All nodes have tooltips with their category, description and chunk source. + + ## Cite Please cite our paper if you use this code in your own work: ```python diff --git a/visualize_graph.py b/visualize_graph.py new file mode 100644 index 0000000..84a1674 --- /dev/null +++ b/visualize_graph.py @@ -0,0 +1,25 @@ +import networkx as nx +from pyvis.network import Network + +# Load the GraphML file +graphml_file = "graph_chunk_entity_relation.graphml" +try: + G = nx.read_graphml(graphml_file) +except FileNotFoundError as e: + friendly_error_msg = f"The file {graphml_file} was not found on this path. Make sure it was generated during indexation." + raise Exception(friendly_error_msg) + +# Create a Pyvis network +net = Network(notebook=True, width="100vw", height="100vh", directed=False) + +# Convert the NetworkX graph to Pyvis with node data as tooltip +for node, data in G.nodes(data=True): + net.add_node(node, label=node, title=str(data)) + +# Display the weight as edge label +for source, target, data in G.edges(data=True): + weight = data.get("weight", "") + net.add_edge(source, target, label=str(weight)) + +# Show the network +net.show('knowledge_graph.html')