Update app.py

app.py

@@ -92,10 +92,11 @@ def create_custom_entity_viz(data, full_text):
                 # Check for overlapping spans
                 overlapping = any(s.start < end and start < s.end for s in spans)
                 if not overlapping:                
-                    span = Span(doc, start, end, label=node["type"])
+                    node_type = node.get("type", "Entity")
+                    span = Span(doc, start, end, label=node_type)
                     spans.append(span)
-                    if node["type"] not in colors:
-                        colors[node["type"]] = get_random_light_color()
+                    if node_type not in colors:
+                        colors[node_type] = get_random_light_color()
 
     doc.set_ents(spans, default="unmodified")
     doc.spans["sc"] = spans
@@ -119,13 +120,22 @@ def create_custom_entity_viz(data, full_text):
 def create_graph(json_data):
     G = nx.Graph()
 
-    # Add nodes with tooltips
+    # Add nodes with tooltips - with error handling for missing keys
     for node in json_data['nodes']:
-        G.add_node(node['id'], title=f"{node['type']}: {node['detailed_type']}")
+        # Get node type with fallback
+        node_type = node.get("type", "Entity")
+        # Get detailed type with fallback
+        detailed_type = node.get("detailed_type", node_type)
+        
+        # Use node ID and type info for the tooltip
+        G.add_node(node['id'], title=f"{node_type}: {detailed_type}")
 
     # Add edges with labels
     for edge in json_data['edges']:
-        G.add_edge(edge['from'], edge['to'], title=edge['label'], label=edge['label'])
+        # Check if the required keys exist
+        if 'from' in edge and 'to' in edge:
+            label = edge.get('label', 'related')
+            G.add_edge(edge['from'], edge['to'], title=label, label=label)
 
     # Create network visualization
     nt = Network(
@@ -339,14 +349,13 @@ def create_ui():
         
         # Full width visualization area at the bottom
         with gr.Row():
-            with gr.Column():
-                # Tab container for visualizations
-                with gr.Tabs():
-                    with gr.TabItem("🧩 Knowledge Graph"):
-                        output_graph = gr.HTML(label="")
-                    
-                    with gr.TabItem("🏷️ Entity Recognition"):
-                        output_entity_viz = gr.HTML(label="")
+            # Full width visualization area
+            with gr.Tabs():
+                with gr.TabItem("🧩 Knowledge Graph"):
+                    output_graph = gr.HTML(label="")
+                
+                with gr.TabItem("🏷️ Entity Recognition"):
+                    output_entity_viz = gr.HTML(label="")
         
         # Functionality
         submit_button.click(