Update app.py

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
-import plotly.graph_objects as go
 import networkx as nx
-import pandas as pd
-from collections import defaultdict
+from pyvis.network import Network
+import tempfile
+import os
 
 # Entity type colors
 ENTITY_COLORS = {
@@ -75,141 +75,68 @@ class NetworkGraphBuilder:
         
         return G
     
-    def create_plotly_graph(self, G, layout_type='spring'):
-        """Create interactive Plotly visualization"""
+    def create_pyvis_graph(self, G):
+        """Create interactive pyvis visualization"""
         if len(G.nodes) == 0:
             return None
         
-        # Choose layout
-        if layout_type == 'spring':
-            pos = nx.spring_layout(G, k=2, iterations=50)
-        elif layout_type == 'circular':
-            pos = nx.circular_layout(G)
-        elif layout_type == 'kamada_kawai':
-            pos = nx.kamada_kawai_layout(G)
-        else:
-            pos = nx.shell_layout(G)
-        
-        # Create edge traces
-        edge_traces = []
-        edge_labels = []
-        
-        for edge in G.edges(data=True):
-            x0, y0 = pos[edge[0]]
-            x1, y1 = pos[edge[1]]
-            
-            # Edge line
-            edge_trace = go.Scatter(
-                x=[x0, x1, None],
-                y=[y0, y1, None],
-                mode='lines',
-                line=dict(width=2, color='#888'),
-                hoverinfo='none',
-                showlegend=False
-            )
-            edge_traces.append(edge_trace)
-            
-            # Edge label (relationship type)
-            rel_type = edge[2].get('relationship', '')
-            edge_label = go.Scatter(
-                x=[(x0 + x1) / 2],
-                y=[(y0 + y1) / 2],
-                mode='text',
-                text=[rel_type],
-                textfont=dict(size=10, color='#555'),
-                hoverinfo='text',
-                hovertext=f"{edge[0]} → {rel_type} → {edge[1]}",
-                showlegend=False
-            )
-            edge_labels.append(edge_label)
+        # Create pyvis network
+        net = Network(height="600px", width="100%", bgcolor="#fafafa", font_color="#333")
+        net.set_options("""
+        {
+          "physics": {
+            "enabled": true,
+            "barnesHut": {
+              "gravitationalConstant": -8000,
+              "springLength": 150,
+              "springConstant": 0.04
+            }
+          },
+          "nodes": {
+            "font": {
+              "size": 16
+            }
+          }
+        }
+        """)
         
-        # Create node traces (one per entity type for legend)
-        node_traces = {}
+        # Add nodes
         for node, data in G.nodes(data=True):
             entity_type = data.get('entity_type', 'UNKNOWN')
+            color = ENTITY_COLORS.get(entity_type, '#CCCCCC')
             
-            if entity_type not in node_traces:
-                node_traces[entity_type] = {
-                    'x': [],
-                    'y': [],
-                    'text': [],
-                    'hovertext': [],
-                    'degree': []
-                }
-            
-            x, y = pos[node]
-            node_traces[entity_type]['x'].append(x)
-            node_traces[entity_type]['y'].append(y)
-            node_traces[entity_type]['text'].append(node)
+            # Node size based on degree
+            degree = G.degree(node)
+            size = 20 + (degree * 5)
             
-            # Create hover text with connections
+            # Create title (tooltip)
             connections = list(G.neighbors(node))
-            hover_info = f"<b>{node}</b><br>"
-            hover_info += f"Type: {entity_type}<br>"
-            hover_info += f"Connections: {len(connections)}<br>"
+            title = f"{node}\nType: {entity_type}\nConnections: {len(connections)}"
             if connections:
-                hover_info += f"Connected to: {', '.join(connections[:5])}"
+                title += f"\nConnected to: {', '.join(connections[:5])}"
                 if len(connections) > 5:
-                    hover_info += f"... and {len(connections) - 5} more"
+                    title += f"... +{len(connections) - 5} more"
             
-            node_traces[entity_type]['hovertext'].append(hover_info)
-            node_traces[entity_type]['degree'].append(G.degree(node))
+            net.add_node(node, label=node, color=color, size=size, title=title)
         
-        # Create Plotly traces for each entity type
-        data = edge_traces + edge_labels
+        # Add edges
+        for edge in G.edges(data=True):
+            rel_type = edge[2].get('relationship', '')
+            net.add_edge(edge[0], edge[1], title=rel_type, label=rel_type)
         
-        for entity_type, trace_data in node_traces.items():
-            # Calculate node sizes based on degree
-            max_degree = max(trace_data['degree']) if trace_data['degree'] else 1
-            sizes = [20 + (degree / max_degree) * 30 for degree in trace_data['degree']]
-            
-            node_trace = go.Scatter(
-                x=trace_data['x'],
-                y=trace_data['y'],
-                mode='markers+text',
-                marker=dict(
-                    size=sizes,
-                    color=ENTITY_COLORS.get(entity_type, '#CCCCCC'),
-                    line=dict(width=2, color='white')
-                ),
-                text=trace_data['text'],
-                textposition='top center',
-                textfont=dict(size=10, color='#333'),
-                hovertext=trace_data['hovertext'],
-                hoverinfo='text',
-                name=entity_type,
-                showlegend=True
-            )
-            data.append(node_trace)
-        
-        # Create figure
-        fig = go.Figure(
-            data=data,
-            layout=go.Layout(
-                title=dict(
-                    text='<b>Entity Network Graph</b><br><sub>Node size indicates number of connections</sub>',
-                    x=0.5,
-                    xanchor='center'
-                ),
-                showlegend=True,
-                hovermode='closest',
-                margin=dict(b=20, l=5, r=5, t=80),
-                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-                plot_bgcolor='#fafafa',
-                height=700,
-                legend=dict(
-                    title=dict(text='<b>Entity Types</b>'),
-                    orientation='v',
-                    yanchor='top',
-                    y=1,
-                    xanchor='left',
-                    x=1.02
-                )
-            )
-        )
+        # Save to temporary file
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.html', mode='w')
+        net.save_graph(temp_file.name)
+        temp_file.close()
+        
+        # Read the HTML content
+        with open(temp_file.name, 'r', encoding='utf-8') as f:
+            html_content = f.read()
         
-        return fig
+        # Clean up
+        os.unlink(temp_file.name)
+        
+        return html_content
 
 def collect_entities_from_records(*args):
     """Collect all entities from the input fields"""
@@ -246,18 +173,18 @@ def collect_entities_from_records(*args):
     
     # Create summary
     summary = f"""
-    ### 📊 Identified Entities
-    - **Total entities:** {len(builder.entities)}
-    - **People:** {sum(1 for e in builder.entities if e['type'] == 'PERSON')}
-    - **Locations:** {sum(1 for e in builder.entities if e['type'] == 'LOCATION')}
-    - **Events:** {sum(1 for e in builder.entities if e['type'] == 'EVENT')}
-    - **Organizations:** {sum(1 for e in builder.entities if e['type'] == 'ORGANIZATION')}
-    - **Dates:** {sum(1 for e in builder.entities if e['type'] == 'DATE')}
-    
-    Now define relationships between these entities on the right →
-    """
+### 📊 Identified Entities
+- **Total entities:** {len(builder.entities)}
+- **People:** {sum(1 for e in builder.entities if e['type'] == 'PERSON')}
+- **Locations:** {sum(1 for e in builder.entities if e['type'] == 'LOCATION')}
+- **Events:** {sum(1 for e in builder.entities if e['type'] == 'EVENT')}
+- **Organizations:** {sum(1 for e in builder.entities if e['type'] == 'ORGANIZATION')}
+- **Dates:** {sum(1 for e in builder.entities if e['type'] == 'DATE')}
+
+Now define relationships between these entities on the right →
+"""
     
-    # Return summary and update all dropdowns (5 relationships × 2 dropdowns each = 10 updates)
+    # Return summary and update all dropdowns
     dropdown_updates = [gr.update(choices=entity_names, value=None)] * 10
     return [summary] + dropdown_updates
 
@@ -266,7 +193,7 @@ def generate_network_graph(*args):
     try:
         builder = NetworkGraphBuilder()
         
-        # Collect entities (first 30 args: 6 records × 5 fields)
+        # Collect entities
         num_records = 6
         fields_per_record = 5
         
@@ -291,7 +218,7 @@ def generate_network_graph(*args):
             if date:
                 builder.add_entity(date, 'DATE', record_id)
         
-        # Collect relationships (next args: 5 relationships × 3 fields)
+        # Collect relationships
         relationship_start = 30
         num_relationships = 5
         
@@ -304,33 +231,28 @@ def generate_network_graph(*args):
             if source and target:
                 builder.add_relationship(source, target, rel_type)
         
-        # Get layout type (last arg)
-        layout_type = args[-1] if len(args) > relationship_start else 'spring'
-        
         # Build graph
         G = builder.build_graph()
         
         if len(G.nodes) == 0:
             return None, "❌ **No entities to display.** Please enter entities in Step 1 and click 'Identify Entities' first."
         
-        # Create visualization (even if no relationships, show isolated nodes)
-        fig = builder.create_plotly_graph(G, layout_type)
+        # Create visualization
+        html_graph = builder.create_pyvis_graph(G)
         
         # Create statistics
         stats = f"""
-        ### 📈 Network Statistics
-        - **Nodes (Entities):** {G.number_of_nodes()}
-        - **Edges (Relationships):** {G.number_of_edges()}
-        """
+### 📈 Network Statistics
+- **Nodes (Entities):** {G.number_of_nodes()}
+- **Edges (Relationships):** {G.number_of_edges()}
+"""
         
         if len(G.edges) == 0:
             stats += "\n⚠️ **No relationships defined** - showing isolated nodes only.\n"
-            stats += "\n*Define relationships in Step 2 to see connections between entities.*\n"
         else:
             stats += f"- **Network Density:** {nx.density(G):.3f}\n"
             stats += f"- **Average Connections per Node:** {sum(dict(G.degree()).values()) / G.number_of_nodes():.2f}\n"
-        
-        if G.number_of_edges() > 0:
+            
             # Find most connected nodes
             degrees = dict(G.degree())
             top_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:3]
@@ -338,19 +260,21 @@ def generate_network_graph(*args):
             for node, degree in top_nodes:
                 stats += f"- {node}: {degree} connections\n"
         
-        return fig, stats
+        return html_graph, stats
     
     except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
         error_msg = f"""
-        ### ❌ Error Generating Graph
-        
-        An error occurred: {str(e)}
-        
-        **Troubleshooting:**
-        1. Make sure you've clicked "Identify Entities" first
-        2. Check that you have at least one entity entered
-        3. If problem persists, try refreshing the page
-        """
+### ❌ Error Generating Graph
+
+{str(e)}
+
+**Technical details:**
+```
+{error_trace}
+```
+"""
         return None, error_msg
 
 def create_interface():
@@ -359,232 +283,107 @@ def create_interface():
         # Basic Network Explorer
         
         Build interactive social network graphs by entering entities extracted through Named Entity Recognition (NER).
-        This tool demonstrates how NER can be used to visualize relationships and connections in text data.
         
         ### How to use this tool:
-        1. **📝 Enter entities** in the records on the left (people, locations, events, organizations, dates)
-        2. **🔗 Click "Identify Entities"** to populate the relationship dropdowns
-        3. **🤝 Define relationships** on the right by selecting entities and connection types
-        4. **🎨 Click "Generate Network Graph"** to visualize your network
-        5. **👁️ Explore** the interactive graph - hover over nodes and edges for details
-        6. **🔄 Refresh the page** to start over with new data
+        1. **📝 Enter entities** in the records on the left
+        2. **🔗 Click "Identify Entities"** to populate the dropdowns
+        3. **🤝 Define relationships** on the right
+        4. **🎨 Click "Generate Network Graph"** to visualize
+        5. **👁️ Explore** - drag nodes, zoom, hover for details
         """)
         
-        # Add tip box
         gr.HTML("""
         <div style="background-color: #fff3cd; border: 1px solid #ffeaa7; border-radius: 8px; padding: 12px; margin: 15px 0;">
-            <strong style="color: #856404;">💡 Top tip:</strong> This tool works best when you have already identified entities from text using NER. Try the NER Explorer Tool first to extract entities automatically!
+            <strong style="color: #856404;">💡 Top tip:</strong> Start with just a few entities and relationships to see how it works!
         </div>
         """)
         
-        # Entity input section
         entity_inputs = []
         
-        # Two-column layout: Entities on left, Relationships on right
         with gr.Row():
-            # LEFT COLUMN: Entity Inputs
+            # LEFT: Entity Inputs
             with gr.Column(scale=1):
-                with gr.Accordion("📚 Step 1: Enter Entities from Your Records", open=True):
-                    # First 4 records (always visible)
+                with gr.Accordion("📚 Step 1: Enter Entities", open=True):
                     for i in range(4):
                         with gr.Group():
-                            gr.Markdown(f"### Record {i+1}")
-                            with gr.Row():
-                                person = gr.Textbox(label="👤 Person", placeholder="e.g., Winston Churchill")
-                                location = gr.Textbox(label="📍 Location", placeholder="e.g., London")
-                            with gr.Row():
-                                event = gr.Textbox(label="📅 Event", placeholder="e.g., Battle of Britain")
-                                org = gr.Textbox(label="🏢 Organization", placeholder="e.g., Royal Air Force")
+                            gr.Markdown(f"**Record {i+1}**")
+                            person = gr.Textbox(label="👤 Person", placeholder="e.g., Winston Churchill")
+                            location = gr.Textbox(label="📍 Location", placeholder="e.g., London")
+                            event = gr.Textbox(label="📅 Event", placeholder="e.g., Battle of Britain")
+                            org = gr.Textbox(label="🏢 Organization", placeholder="e.g., War Cabinet")
                             date = gr.Textbox(label="🗓️ Date", placeholder="e.g., 1940")
-                            
                             entity_inputs.extend([person, location, event, org, date])
                     
-                    # Additional records (collapsible)
-                    with gr.Accordion("➕ Additional Records (5-6)", open=False):
+                    with gr.Accordion("➕ Records 5-6", open=False):
                         for i in range(4, 6):
                             with gr.Group():
-                                gr.Markdown(f"### Record {i+1}")
-                                with gr.Row():
-                                    person = gr.Textbox(label="👤 Person", placeholder="e.g., Winston Churchill")
-                                    location = gr.Textbox(label="📍 Location", placeholder="e.g., London")
-                                with gr.Row():
-                                    event = gr.Textbox(label="📅 Event", placeholder="e.g., Battle of Britain")
-                                    org = gr.Textbox(label="🏢 Organization", placeholder="e.g., Royal Air Force")
-                                date = gr.Textbox(label="🗓️ Date", placeholder="e.g., 1940")
-                                
+                                gr.Markdown(f"**Record {i+1}**")
+                                person = gr.Textbox(label="👤 Person")
+                                location = gr.Textbox(label="📍 Location")
+                                event = gr.Textbox(label="📅 Event")
+                                org = gr.Textbox(label="🏢 Organization")
+                                date = gr.Textbox(label="🗓️ Date")
                                 entity_inputs.extend([person, location, event, org, date])
                 
                 collect_btn = gr.Button("🔍 Identify Entities", variant="primary", size="lg")
                 entity_summary = gr.Markdown()
             
-            # RIGHT COLUMN: Relationship Builder (ALWAYS VISIBLE)
+            # RIGHT: Relationships & Graph
             with gr.Column(scale=1):
-                with gr.Accordion("🤝 Step 2: Define Relationships Between Entities", open=True):
-                    gr.Markdown("*First identify entities, then define relationships below:*")
-                
+                with gr.Accordion("🤝 Step 2: Define Relationships", open=True):
                     relationship_inputs = []
-                    
                     for i in range(5):
                         with gr.Row():
-                            source = gr.Dropdown(label=f"From", choices=[], interactive=True, scale=2)
-                            rel_type = gr.Dropdown(
-                                label="Type",
-                                choices=RELATIONSHIP_TYPES,
-                                value="related_to",
-                                interactive=True,
-                                scale=2
-                            )
-                            target = gr.Dropdown(label=f"To", choices=[], interactive=True, scale=2)
-                        
+                            source = gr.Dropdown(label="From", choices=[], scale=2)
+                            rel_type = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="related_to", scale=2)
+                            target = gr.Dropdown(label="To", choices=[], scale=2)
                         relationship_inputs.extend([source, rel_type, target])
                 
-                with gr.Accordion("🎨 Step 3: Customize and Generate", open=True):
-                    layout_type = gr.Dropdown(
-                        label="Graph Layout",
-                        choices=['spring', 'circular', 'kamada_kawai', 'shell'],
-                        value='spring',
-                        info="Choose how nodes are arranged"
-                    )
-                    
-                    generate_btn = gr.Button("🔍 Generate Network Graph", variant="primary", size="lg")
-        
-        # Output section
-        gr.HTML("<hr style='margin: 30px 0;'>")
-        
-        with gr.Row():
-            network_stats = gr.Markdown()
-        
-        with gr.Row():
-            network_plot = gr.Plot(label="Interactive Network Graph")
+                generate_btn = gr.Button("🔍 Generate Network Graph", variant="primary", size="lg")
+                
+                gr.HTML("<hr style='margin: 20px 0;'>")
+                network_stats = gr.Markdown()
+                network_plot = gr.HTML(label="Interactive Network Graph")
         
         # Examples
-        with gr.Column():
-            gr.Markdown("""
-            ### 💡 No example entities to test? No problem!
-            Simply click on one of the examples provided below, and the fields will be populated for you.
-            """, elem_id="examples-heading")
+        gr.Markdown("### 💡 Try an example:")
         gr.Examples(
-            examples=[
-                [
-                    # === ENTITY RECORDS ===
-                    # Record 1
-                    "Winston Churchill", "London", "Battle of Britain", "War Cabinet", "1940",
-                    # Record 2
-                    "Clement Attlee", "London", "Potsdam Conference", "Labour Party", "1945",
-                    # Record 3
-                    "Field Marshal Montgomery", "North Africa", "Battle of El Alamein", "Eighth Army", "1942",
-                    # Record 4
-                    "Winston Churchill", "Yalta", "Yalta Conference", "War Cabinet", "February 1945",
-                    # Record 5
-                    "King George VI", "London", "Victory in Europe Day", "British Monarchy", "May 1945",
-                    # Record 6
-                    "Field Marshal Montgomery", "Lüneburg Heath", "German Surrender", "British Army", "May 1945",
-                    # === RELATIONSHIPS ===
-                    # Relationship 1
-                    "Winston Churchill", "works_with", "Clement Attlee",
-                    # Relationship 2
-                    "Winston Churchill", "participated_in", "Battle of Britain",
-                    # Relationship 3
-                    "Field Marshal Montgomery", "participated_in", "Battle of El Alamein",
-                    # Relationship 4
-                    "Winston Churchill", "participated_in", "Yalta Conference",
-                    # Relationship 5
-                    "Clement Attlee", "participated_in", "Potsdam Conference",
-                    # Layout type
-                    "spring"
-                ],
-                [
-                    # === ENTITY RECORDS ===
-                    # Record 1 - Pride and Prejudice
-                    "Elizabeth Bennet", "Longbourn", "Meryton Assembly", "", "Autumn 1811",
-                    # Record 2
-                    "Mr Darcy", "Pemberley", "Meryton Assembly", "", "Autumn 1811",
-                    # Record 3
-                    "Jane Bennet", "Longbourn", "Netherfield Ball", "", "November 1811",
-                    # Record 4
-                    "Mr Bingley", "Netherfield", "Netherfield Ball", "", "November 1811",
-                    # Record 5
-                    "Elizabeth Bennet", "Rosings", "Easter Visit", "", "Spring 1812",
-                    # Record 6
-                    "Mr Darcy", "Rosings", "First Proposal", "", "Spring 1812",
-                    # === RELATIONSHIPS ===
-                    # Relationship 1
-                    "Elizabeth Bennet", "knows", "Mr Darcy",
-                    # Relationship 2
-                    "Jane Bennet", "knows", "Mr Bingley",
-                    # Relationship 3
-                    "Elizabeth Bennet", "located_in", "Longbourn",
-                    # Relationship 4
-                    "Mr Darcy", "located_in", "Pemberley",
-                    # Relationship 5
-                    "Elizabeth Bennet", "participated_in", "Meryton Assembly",
-                    # Layout type
-                    "spring"
-                ]
-            ],
-            inputs=entity_inputs + relationship_inputs + [layout_type],
-            label="Examples"
+            examples=[[
+                "Winston Churchill", "London", "Battle of Britain", "War Cabinet", "1940",
+                "Clement Attlee", "London", "Potsdam Conference", "Labour Party", "1945",
+                "Field Marshal Montgomery", "North Africa", "Battle of El Alamein", "Eighth Army", "1942",
+                "Winston Churchill", "Yalta", "Yalta Conference", "War Cabinet", "February 1945",
+                "", "", "", "", "",
+                "", "", "", "", "",
+                "Winston Churchill", "works_with", "Clement Attlee",
+                "Winston Churchill", "participated_in", "Battle of Britain",
+                "Field Marshal Montgomery", "participated_in", "Battle of El Alamein",
+                "", "", "",
+                "", "", "",
+            ]],
+            inputs=entity_inputs + relationship_inputs,
+            label="WWII Example"
         )
         
-        # Add custom CSS to match NER tool styling
-        gr.HTML("""
-        <style>
-            /* Make the Examples label text black */
-            .gradio-examples-label {
-                color: black !important;
-            }
-            h4.examples-label, .examples-label {
-                color: black !important;
-            }
-            #examples-heading + div label,
-            #examples-heading + div .label-text {
-                color: black !important;
-            }
-        </style>
-        """)
-        
-        # Wire up the interface
-        # Collect entities button - updates the relationship dropdowns
+        # Wire up
         collect_btn.click(
             fn=collect_entities_from_records,
             inputs=entity_inputs,
-            outputs=[entity_summary] + relationship_inputs[::3] + relationship_inputs[2::3]  # Update source and target dropdowns
+            outputs=[entity_summary] + relationship_inputs[::3] + relationship_inputs[2::3]
         )
         
-        # Generate graph button
-        all_inputs = entity_inputs + relationship_inputs + [layout_type]
+        all_inputs = entity_inputs + relationship_inputs
         generate_btn.click(
             fn=generate_network_graph,
             inputs=all_inputs,
             outputs=[network_plot, network_stats]
         )
         
-        # Information footer
         gr.HTML("""
-        <hr style="margin-top: 40px; margin-bottom: 20px;">
-        <div style="background-color: #f8f9fa; padding: 20px; border-radius: 8px; margin-top: 20px;">
-            <h4 style="margin-top: 0;">ℹ️ About This Tool</h4>
-            <p style="font-size: 14px; line-height: 1.8;">
-                This tool demonstrates how <strong>Named Entity Recognition (NER)</strong> can be combined with 
-                <strong>network analysis</strong> to visualize relationships in text data. In real-world applications,
-                entities would be automatically extracted from text using NER models, and relationships could be 
-                identified through co-occurrence analysis, dependency parsing, or machine learning.
-            </p>
-            <p style="font-size: 14px; line-height: 1.8; margin-bottom: 0;">
-                <strong>Built with:</strong> Gradio, NetworkX, and Plotly | 
-                <strong>Graph Layouts:</strong> Spring (force-directed), Circular, Kamada-Kawai, Shell
-            </p>
-        </div>
-        
-        <br> 
-        <hr style="margin-top: 40px; margin-bottom: 20px;">
-        <div style="background-color: #f8f9fa; padding: 20px; border-radius: 8px; margin-top: 20px; text-align: center;">
-            <p style="font-size: 14px; line-height: 1.8; margin: 0;">
-                This <strong>Basic Network Explorer</strong> tool was created as part of a Bodleian Libraries (University of Oxford) Sassoon Research Fellowship.
-            </p><br><br> 
-            <p style="font-size: 14px; line-height: 1.8; margin: 0;">
-                The code for this tool was built with the aid of Claude Sonnet 4.5.
-            </p>
+        <hr style="margin: 40px 0;">
+        <div style="text-align: center; color: #666; font-size: 14px;">
+            <p>Basic Network Explorer | Bodleian Libraries (University of Oxford) Sassoon Research Fellowship</p>
+            <p>Built with the aid of Claude Sonnet 4.5</p>
         </div>
         """)