Spaces:
Paused
Paused
File size: 2,717 Bytes
5919b13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import json
import networkx as nx
from pyvis.network import Network
def create_knowledge_graph(events_file="data/parsed_event_data.json",
characters_file="data/parsed_character_metadata.json",
output_file="data/knowledge_graph.html"):
"""
Create a knowledge graph from parsed events and character data
Args:
events_file (str): Path to parsed event data JSON
characters_file (str): Path to character metadata JSON
output_file (str): Path to save the HTML knowledge graph
"""
with open(events_file, "r") as f:
parsed_data = json.load(f)
with open(characters_file, "r") as f:
character_metadata = json.load(f)
story_title = list(parsed_data.keys())[0]
events = parsed_data[story_title]
character_metadata.pop("Story_name", None)
G = nx.MultiDiGraph()
for char, desc in character_metadata.items():
G.add_node(char, type="character", description=desc)
for event in events:
event_name = event.get("event_name", "")
event_desc = event.get("description", "")
G.add_node(event_name, type="event", description=event_desc)
for obj in event.get("objects", []):
G.add_node(obj, type="object")
G.add_edge(event_name, obj, relation="has_object")
for env in event.get("environment", []):
G.add_node(env, type="environment")
G.add_edge(event_name, env, relation="is_in")
for actor in event.get("actors", []):
matched = False
for char in character_metadata:
if char.lower() in actor.lower():
G.add_edge(char, event_name, relation="participates_in")
matched = True
break
if not matched:
G.add_node(actor, type="actor")
G.add_edge(actor, event_name, relation="participates_in")
for char, desc in character_metadata.items():
G.nodes[char]["description"] = desc
net = Network(height="800px", width="100%", notebook=True, directed=True, cdn_resources="in_line")
net.from_nx(G)
color_map = {
"character": "#ffb347",
"event": "#87ceeb",
"object": "#90ee90",
"environment": "#f08080",
"actor": "#dda0dd"
}
for node in net.nodes:
ntype = G.nodes[node["id"]].get("type", "actor")
node["color"] = color_map.get(ntype, "#cccccc")
node["title"] = G.nodes[node["id"]].get("description", "")
net.write_html(output_file)
print(f"Knowledge graph saved as {output_file}")
if __name__ == "__main__":
create_knowledge_graph() |