Update app.py

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
 import networkx as nx
-import plotly.graph_objects as go
-import numpy as np
+from pyvis.network import Network
+import tempfile
+import os
+import base64
 
-# Entity type colors
+# Entity type colors (matching your NER tool)
 ENTITY_COLORS = {
     'PERSON': '#00B894',      # Green
     'LOCATION': '#A0E7E5',    # Light Cyan
@@ -14,6 +16,8 @@ ENTITY_COLORS = {
 
 # Relationship types for dropdown
 RELATIONSHIP_TYPES = [
+    'related_to',
+    'knows',
     'works_with',
     'located_in',
     'participated_in',
@@ -22,17 +26,19 @@ RELATIONSHIP_TYPES = [
     'employed_by',
     'founded',
     'attended',
-    'knows',
-    'related_to',
     'collaborates_with',
     'married_to',
     'sibling_of',
     'parent_of',
     'lives_at',
     'wrote',
+    'visited',
+    'born_in',
+    'died_in',
     'other'
 ]
 
+
 class NetworkGraphBuilder:
     def __init__(self):
         self.entities = []
@@ -82,129 +88,133 @@ class NetworkGraphBuilder:
         
         return G
     
-    def create_plotly_graph(self, G):
-        """Create interactive Plotly visualization"""
+    def create_pyvis_graph(self, G):
+        """Create interactive PyVis visualization"""
         if len(G.nodes) == 0:
             return None
         
-        # Use spring layout for positioning
-        pos = nx.spring_layout(G, k=2, iterations=50, seed=42)
-        
-        # Create edge traces
-        edge_x = []
-        edge_y = []
-        edge_labels = []
-        edge_mid_x = []
-        edge_mid_y = []
-        
-        for edge in G.edges(data=True):
-            x0, y0 = pos[edge[0]]
-            x1, y1 = pos[edge[1]]
-            edge_x.extend([x0, x1, None])
-            edge_y.extend([y0, y1, None])
-            
-            # Midpoint for edge label
-            mid_x = (x0 + x1) / 2
-            mid_y = (y0 + y1) / 2
-            edge_mid_x.append(mid_x)
-            edge_mid_y.append(mid_y)
-            edge_labels.append(edge[2].get('relationship', ''))
-        
-        edge_trace = go.Scatter(
-            x=edge_x, y=edge_y,
-            line=dict(width=2, color='#888'),
-            hoverinfo='none',
-            mode='lines'
-        )
-        
-        # Edge labels trace
-        edge_label_trace = go.Scatter(
-            x=edge_mid_x, y=edge_mid_y,
-            mode='text',
-            text=edge_labels,
-            textposition='middle center',
-            textfont=dict(size=10, color='#666'),
-            hoverinfo='none'
+        # Create PyVis network with dark theme like the example
+        net = Network(
+            height="600px",
+            width="100%",
+            bgcolor="#1a1a2e",
+            font_color="white",
+            directed=False
         )
         
-        # Create node traces - one per entity type for legend
-        node_traces = []
+        # Configure physics for interactive dragging
+        net.set_options("""
+        {
+            "nodes": {
+                "borderWidth": 2,
+                "borderWidthSelected": 4,
+                "font": {
+                    "size": 14,
+                    "face": "arial",
+                    "color": "white"
+                },
+                "shadow": {
+                    "enabled": true,
+                    "color": "rgba(0,0,0,0.5)",
+                    "size": 10
+                }
+            },
+            "edges": {
+                "color": {
+                    "inherit": false
+                },
+                "smooth": {
+                    "enabled": true,
+                    "type": "continuous"
+                },
+                "shadow": {
+                    "enabled": true,
+                    "color": "rgba(0,0,0,0.3)"
+                }
+            },
+            "physics": {
+                "enabled": true,
+                "barnesHut": {
+                    "gravitationalConstant": -8000,
+                    "centralGravity": 0.3,
+                    "springLength": 150,
+                    "springConstant": 0.04,
+                    "damping": 0.09
+                },
+                "stabilization": {
+                    "enabled": true,
+                    "iterations": 100
+                }
+            },
+            "interaction": {
+                "hover": true,
+                "tooltipDelay": 100,
+                "dragNodes": true,
+                "dragView": true,
+                "zoomView": true
+            }
+        }
+        """)
         
-        for entity_type, color in ENTITY_COLORS.items():
-            nodes_of_type = [n for n, d in G.nodes(data=True) if d.get('entity_type') == entity_type]
-            if not nodes_of_type:
-                continue
-                
-            node_x = []
-            node_y = []
-            node_text = []
-            node_hover = []
-            node_sizes = []
+        # Add nodes with styling based on entity type and degree
+        for node in G.nodes():
+            data = G.nodes[node]
+            entity_type = data.get('entity_type', 'UNKNOWN')
+            color = ENTITY_COLORS.get(entity_type, '#CCCCCC')
             
-            for node in nodes_of_type:
-                x, y = pos[node]
-                node_x.append(x)
-                node_y.append(y)
-                node_text.append(node)
-                
-                # Calculate connections
-                connections = list(G.neighbors(node))
-                degree = len(connections)
-                node_sizes.append(30 + (degree * 10))
-                
-                # Hover text
-                hover = f"<b>{node}</b><br>Type: {entity_type}<br>Connections: {degree}"
-                if connections:
-                    hover += f"<br>Connected to: {', '.join(connections[:5])}"
-                    if len(connections) > 5:
-                        hover += f"<br>... +{len(connections) - 5} more"
-                node_hover.append(hover)
+            # Size based on connections (degree)
+            degree = G.degree(node)
+            size = 25 + (degree * 8)
             
-            node_trace = go.Scatter(
-                x=node_x, y=node_y,
-                mode='markers+text',
-                name=entity_type,
-                text=node_text,
-                textposition='top center',
-                textfont=dict(size=12, color='#333'),
-                hoverinfo='text',
-                hovertext=node_hover,
-                marker=dict(
-                    size=node_sizes,
-                    color=color,
-                    line=dict(width=2, color='white'),
-                    symbol='circle'
-                )
+            # Create tooltip
+            connections = list(G.neighbors(node))
+            title = f"<b>{node}</b><br>Type: {entity_type}<br>Connections: {len(connections)}"
+            if connections:
+                title += f"<br>Connected to: {', '.join(connections[:5])}"
+                if len(connections) > 5:
+                    title += f"<br>... +{len(connections) - 5} more"
+            
+            net.add_node(
+                node,
+                label=node,
+                color=color,
+                size=size,
+                title=title,
+                font={'size': 12, 'color': 'white'}
             )
-            node_traces.append(node_trace)
-        
-        # Create figure
-        fig = go.Figure(
-            data=[edge_trace, edge_label_trace] + node_traces,
-            layout=go.Layout(
-                title=dict(
-                    text='Interactive Network Graph',
-                    font=dict(size=20)
-                ),
-                showlegend=True,
-                legend=dict(
-                    yanchor="top",
-                    y=0.99,
-                    xanchor="left",
-                    x=0.01,
-                    bgcolor="rgba(255,255,255,0.8)"
-                ),
-                hovermode='closest',
-                margin=dict(b=20, l=5, r=5, t=50),
-                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-                plot_bgcolor='#fafafa',
-                paper_bgcolor='#fafafa',
-                height=600
+        
+        # Add edges with relationship labels
+        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,
+                color='#888888',
+                width=2,
+                font={'size': 10, 'color': '#cccccc', 'align': 'middle'}
             )
-        )
         
-        return fig
+        # Generate HTML
+        html = net.generate_html()
+        
+        # Wrap in iframe-compatible format
+        # Encode as base64 data URI for iframe src
+        html_bytes = html.encode('utf-8')
+        b64_html = base64.b64encode(html_bytes).decode('utf-8')
+        
+        iframe_html = f'''
+        <iframe 
+            src="data:text/html;base64,{b64_html}" 
+            width="100%" 
+            height="620px" 
+            style="border: 2px solid #333; border-radius: 10px; background-color: #1a1a2e;"
+            sandbox="allow-scripts allow-same-origin"
+        ></iframe>
+        '''
+        
+        return iframe_html
 
 
 def collect_entities_from_records(
@@ -252,26 +262,43 @@ def collect_entities_from_records(
         'DATE': sum(1 for e in builder.entities if e['type'] == 'DATE'),
     }
     
-    # Create summary
-    summary = f"""
-### 📊 Identified Entities ({len(builder.entities)} total)
-| Type | Count |
-|------|-------|
-| 👤 People | {counts['PERSON']} |
-| 📍 Locations | {counts['LOCATION']} |
-| 📅 Events | {counts['EVENT']} |
-| 🏢 Organizations | {counts['ORGANIZATION']} |
-| 🗓️ Dates | {counts['DATE']} |
-
-**Entities found:** {', '.join(entity_names) if entity_names else 'None'}
-
-➡️ Now define relationships between these entities below
-"""
+    # Create HTML summary that spans full width
+    summary_html = f'''
+    <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 20px; border-radius: 12px; margin: 10px 0;">
+        <h3 style="color: white; margin: 0 0 15px 0; font-size: 18px;">📊 Identified Entities ({len(builder.entities)} total)</h3>
+        <div style="display: flex; flex-wrap: wrap; gap: 15px; margin-bottom: 15px;">
+            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
+                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['PERSON']}</div>
+                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">👤 People</div>
+            </div>
+            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
+                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['LOCATION']}</div>
+                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">📍 Locations</div>
+            </div>
+            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
+                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['EVENT']}</div>
+                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">📅 Events</div>
+            </div>
+            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
+                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['ORGANIZATION']}</div>
+                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">🏢 Organizations</div>
+            </div>
+            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
+                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['DATE']}</div>
+                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">🗓️ Dates</div>
+            </div>
+        </div>
+        <div style="background: rgba(255,255,255,0.15); padding: 12px; border-radius: 8px;">
+            <div style="color: rgba(255,255,255,0.9); font-size: 13px; word-wrap: break-word;">
+                <strong>Entities:</strong> {', '.join(entity_names) if entity_names else 'None found'}
+            </div>
+        </div>
+    </div>
+    '''
     
     # Return summary and update all 10 dropdowns (5 source + 5 target)
-    # Each dropdown gets updated with the entity names
     return (
-        summary,
+        summary_html,
         gr.update(choices=entity_names, value=None),  # source 1
         gr.update(choices=entity_names, value=None),  # target 1
         gr.update(choices=entity_names, value=None),  # source 2
@@ -341,76 +368,87 @@ def generate_network_graph(
         G = builder.build_graph()
         
         if len(G.nodes) == 0:
-            empty_fig = go.Figure()
-            empty_fig.add_annotation(
-                text="No entities to display.<br>Please enter entities in Step 1 and click 'Identify Entities' first.",
-                xref="paper", yref="paper",
-                x=0.5, y=0.5, showarrow=False,
-                font=dict(size=16, color="#666")
-            )
-            empty_fig.update_layout(
-                height=400,
-                plot_bgcolor='#fafafa',
-                paper_bgcolor='#fafafa'
-            )
-            return empty_fig, "❌ **No entities to display.** Please enter entities in Step 1 first."
+            empty_html = '''
+            <div style="background-color: #1a1a2e; height: 400px; display: flex; align-items: center; justify-content: center; border-radius: 10px; border: 2px solid #333;">
+                <div style="text-align: center; color: white;">
+                    <div style="font-size: 48px; margin-bottom: 20px;">📊</div>
+                    <h3>No entities to display</h3>
+                    <p style="color: #888;">Enter entities above and click "Identify Entities" first</p>
+                </div>
+            </div>
+            '''
+            return empty_html, "❌ **No entities to display.** Please enter entities in Step 1 first."
         
         # Create visualization
-        fig = builder.create_plotly_graph(G)
+        graph_html = builder.create_pyvis_graph(G)
         
         # Create statistics
-        stats = f"""
-### 📈 Network Statistics
-| Metric | Value |
-|--------|-------|
-| **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. Add relationships in Step 2!\n"
-        else:
+        stats_html = f'''
+        <div style="background: #f8f9fa; padding: 15px; border-radius: 10px; border: 1px solid #ddd;">
+            <h4 style="margin: 0 0 10px 0;">📈 Network Statistics</h4>
+            <table style="width: 100%; border-collapse: collapse;">
+                <tr>
+                    <td style="padding: 8px; border-bottom: 1px solid #eee;"><strong>Nodes</strong></td>
+                    <td style="padding: 8px; border-bottom: 1px solid #eee; text-align: right;">{G.number_of_nodes()}</td>
+                </tr>
+                <tr>
+                    <td style="padding: 8px; border-bottom: 1px solid #eee;"><strong>Edges</strong></td>
+                    <td style="padding: 8px; border-bottom: 1px solid #eee; text-align: right;">{G.number_of_edges()}</td>
+                </tr>
+        '''
+        
+        if len(G.edges) > 0:
             density = nx.density(G)
             avg_degree = sum(dict(G.degree()).values()) / G.number_of_nodes()
-            stats += f"| **Network Density** | {density:.3f} |\n"
-            stats += f"| **Avg. Connections** | {avg_degree:.2f} |\n"
+            stats_html += f'''
+                <tr>
+                    <td style="padding: 8px; border-bottom: 1px solid #eee;"><strong>Density</strong></td>
+                    <td style="padding: 8px; border-bottom: 1px solid #eee; text-align: right;">{density:.3f}</td>
+                </tr>
+                <tr>
+                    <td style="padding: 8px;"><strong>Avg. Connections</strong></td>
+                    <td style="padding: 8px; text-align: right;">{avg_degree:.2f}</td>
+                </tr>
+            </table>
+            '''
             
-            # Find most connected nodes
+            # Most connected
             degrees = dict(G.degree())
             top_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:3]
-            stats += "\n**Most Connected Entities:**\n"
+            stats_html += '''
+            <div style="margin-top: 15px;">
+                <strong>Most Connected:</strong>
+                <ul style="margin: 5px 0; padding-left: 20px;">
+            '''
             for node, degree in top_nodes:
-                stats += f"- **{node}**: {degree} connections\n"
+                stats_html += f'<li>{node}: {degree}</li>'
+            stats_html += '</ul></div>'
+        else:
+            stats_html += '''
+                </table>
+                <div style="margin-top: 10px; padding: 10px; background: #fff3cd; border-radius: 5px; color: #856404;">
+                    ⚠️ No relationships defined - nodes are isolated
+                </div>
+            '''
+        
+        stats_html += '</div>'
         
-        return fig, stats
+        return graph_html, stats_html
     
     except Exception as e:
         import traceback
         error_trace = traceback.format_exc()
         
-        error_fig = go.Figure()
-        error_fig.add_annotation(
-            text=f"Error: {str(e)}",
-            xref="paper", yref="paper",
-            x=0.5, y=0.5, showarrow=False,
-            font=dict(size=14, color="red")
-        )
-        error_fig.update_layout(height=300, plot_bgcolor='#fafafa', paper_bgcolor='#fafafa')
-        
-        error_msg = f"""
-### ❌ Error Generating Graph
-
-{str(e)}
-
-<details>
-<summary>Technical details</summary>
-
-```
-{error_trace}
-```
-</details>
-"""
-        return error_fig, error_msg
+        error_html = f'''
+        <div style="background-color: #f8d7da; height: 200px; display: flex; align-items: center; justify-content: center; border-radius: 10px; border: 2px solid #f5c6cb;">
+            <div style="text-align: center; color: #721c24; padding: 20px;">
+                <div style="font-size: 48px; margin-bottom: 10px;">⚠️</div>
+                <h3>Error generating graph</h3>
+                <p>{str(e)}</p>
+            </div>
+        </div>
+        '''
+        return error_html, f"❌ Error: {str(e)}"
 
 
 def load_austen_example():
@@ -419,11 +457,11 @@ def load_austen_example():
         # Record 1
         "Elizabeth Bennet", "Longbourn", "Meryton Ball", "Bennet Family", "1811",
         # Record 2
-        "Mr. Darcy", "Pemberley", "Netherfield Ball", "Darcy Estate", "1811",
+        "Mr. Darcy", "Pemberley", "Netherfield Ball", "Darcy Estate", "",
         # Record 3
-        "Jane Bennet", "Netherfield", "", "Bennet Family", "",
+        "Jane Bennet", "Netherfield", "", "", "",
         # Record 4
-        "Mr. Bingley", "Netherfield", "", "", "",
+        "Mr. Bingley", "London", "", "", "",
         # Record 5
         "Mr. Wickham", "Meryton", "", "Militia", "",
         # Record 6
@@ -437,11 +475,11 @@ def load_wwii_example():
         # Record 1
         "Winston Churchill", "London", "Battle of Britain", "War Cabinet", "1940",
         # Record 2
-        "Clement Attlee", "London", "Potsdam Conference", "Labour Party", "1945",
+        "Franklin D. Roosevelt", "Washington D.C.", "D-Day", "Allied Forces", "1944",
         # Record 3
         "Field Marshal Montgomery", "North Africa", "Battle of El Alamein", "Eighth Army", "1942",
         # Record 4
-        "Franklin D. Roosevelt", "Washington D.C.", "D-Day", "Allied Forces", "1944",
+        "Clement Attlee", "Potsdam", "Potsdam Conference", "Labour Party", "1945",
         # Record 5
         "", "", "", "", "",
         # Record 6
@@ -450,136 +488,213 @@ def load_wwii_example():
 
 
 def create_interface():
-    with gr.Blocks(title="Basic Network Explorer", theme=gr.themes.Soft()) as demo:
+    with gr.Blocks(title="Network Explorer", theme=gr.themes.Soft(), css="""
+        .record-box {
+            border: 1px solid #e0e0e0;
+            border-radius: 8px;
+            padding: 12px;
+            background: #fafafa;
+        }
+        .full-width {
+            width: 100%;
+        }
+    """) as demo:
         gr.Markdown("""
-        # 🕸️ Basic Network Explorer
-        
-        Build interactive network graphs by entering entities extracted through Named Entity Recognition (NER).
-        Explore relationships between people, places, events, organizations and dates.
+        # 🕸️ Network Explorer
         
-        ### How to use this tool:
-        1. **📝 Enter entities** in the records below (or load an example)
-        2. **🔍 Click "Identify Entities"** to collect and list all entities
-        3. **🤝 Define relationships** between entities using the dropdowns
-        4. **🎨 Click "Generate Network Graph"** to visualize
-        5. **👁️ Explore** - hover over nodes for details, zoom and pan the graph
+        Build interactive network graphs from named entities. Enter people, places, events, 
+        organizations and dates, then define relationships to visualize connections.
         """)
         
-        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> Start with just a few entities and relationships to see how it works!
-        </div>
-        """)
+        # Quick start buttons
+        gr.Markdown("### 💡 Quick Start - Load an Example:")
+        with gr.Row():
+            austen_btn = gr.Button("📚 Jane Austen (Pride & Prejudice)", variant="secondary", size="sm")
+            wwii_btn = gr.Button("⚔️ WWII History", variant="secondary", size="sm")
+        
+        gr.HTML("<hr style='margin: 15px 0;'>")
+        
+        # ==================== STEP 1: ENTITY INPUT (2x2 Grid) ====================
+        gr.Markdown("## 📝 Step 1: Enter Entities")
         
-        # Store all entity input components
         entity_inputs = []
         
-        # Example buttons
+        # First row: Records 1 & 2
         with gr.Row():
-            gr.Markdown("### 💡 Quick Start - Load an Example:")
+            with gr.Column():
+                gr.Markdown("**Record 1**")
+                with gr.Group():
+                    p1 = gr.Textbox(label="👤 Person", placeholder="e.g., Elizabeth Bennet")
+                    l1 = gr.Textbox(label="📍 Location", placeholder="e.g., Longbourn")
+                    e1 = gr.Textbox(label="📅 Event", placeholder="e.g., Meryton Ball")
+                    o1 = gr.Textbox(label="🏢 Organization", placeholder="e.g., Bennet Family")
+                    d1 = gr.Textbox(label="🗓️ Date", placeholder="e.g., 1811")
+                entity_inputs.extend([p1, l1, e1, o1, d1])
+            
+            with gr.Column():
+                gr.Markdown("**Record 2**")
+                with gr.Group():
+                    p2 = gr.Textbox(label="👤 Person", placeholder="e.g., Mr. Darcy")
+                    l2 = gr.Textbox(label="📍 Location", placeholder="e.g., Pemberley")
+                    e2 = gr.Textbox(label="📅 Event", placeholder="e.g., Netherfield Ball")
+                    o2 = gr.Textbox(label="🏢 Organization", placeholder="e.g., Darcy Estate")
+                    d2 = gr.Textbox(label="🗓️ Date", placeholder="")
+                entity_inputs.extend([p2, l2, e2, o2, d2])
+        
+        # Second row: Records 3 & 4
         with gr.Row():
-            austen_btn = gr.Button("📚 Jane Austen (Pride & Prejudice)", variant="secondary")
-            wwii_btn = gr.Button("⚔️ WWII History", variant="secondary")
+            with gr.Column():
+                gr.Markdown("**Record 3**")
+                with gr.Group():
+                    p3 = gr.Textbox(label="👤 Person", placeholder="e.g., Jane Bennet")
+                    l3 = gr.Textbox(label="📍 Location", placeholder="e.g., Netherfield")
+                    e3 = gr.Textbox(label="📅 Event", placeholder="")
+                    o3 = gr.Textbox(label="🏢 Organization", placeholder="")
+                    d3 = gr.Textbox(label="🗓️ Date", placeholder="")
+                entity_inputs.extend([p3, l3, e3, o3, d3])
+            
+            with gr.Column():
+                gr.Markdown("**Record 4**")
+                with gr.Group():
+                    p4 = gr.Textbox(label="👤 Person", placeholder="e.g., Mr. Bingley")
+                    l4 = gr.Textbox(label="📍 Location", placeholder="e.g., London")
+                    e4 = gr.Textbox(label="📅 Event", placeholder="")
+                    o4 = gr.Textbox(label="🏢 Organization", placeholder="")
+                    d4 = gr.Textbox(label="🗓️ Date", placeholder="")
+                entity_inputs.extend([p4, l4, e4, o4, d4])
+        
+        # Optional records 5-6
+        with gr.Accordion("➕ Additional Records (5-6)", open=False):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("**Record 5**")
+                    with gr.Group():
+                        p5 = gr.Textbox(label="👤 Person")
+                        l5 = gr.Textbox(label="📍 Location")
+                        e5 = gr.Textbox(label="📅 Event")
+                        o5 = gr.Textbox(label="🏢 Organization")
+                        d5 = gr.Textbox(label="🗓️ Date")
+                    entity_inputs.extend([p5, l5, e5, o5, d5])
+                
+                with gr.Column():
+                    gr.Markdown("**Record 6**")
+                    with gr.Group():
+                        p6 = gr.Textbox(label="👤 Person")
+                        l6 = gr.Textbox(label="📍 Location")
+                        e6 = gr.Textbox(label="📅 Event")
+                        o6 = gr.Textbox(label="🏢 Organization")
+                        d6 = gr.Textbox(label="🗓️ Date")
+                    entity_inputs.extend([p6, l6, e6, o6, d6])
+        
+        # Identify button
+        collect_btn = gr.Button("🔍 Identify Entities", variant="primary", size="lg")
+        
+        # Full-width entity summary
+        entity_summary = gr.HTML()
         
         gr.HTML("<hr style='margin: 20px 0;'>")
         
+        # ==================== STEP 2: RELATIONSHIPS (Grid) ====================
+        gr.Markdown("## 🤝 Step 2: Define Relationships")
+        gr.Markdown("*Select entities from the dropdowns to create connections*")
+        
+        # Relationship inputs in a grid
+        relationship_inputs = []
+        
         with gr.Row():
-            # LEFT COLUMN: Entity Inputs
-            with gr.Column(scale=1):
-                gr.Markdown("## 📚 Step 1: Enter Entities")
-                
-                with gr.Accordion("Records 1-4", 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., Elizabeth Bennet", scale=1)
-                                location = gr.Textbox(label="📍 Location", placeholder="e.g., Longbourn", scale=1)
-                            with gr.Row():
-                                event = gr.Textbox(label="📅 Event", placeholder="e.g., Meryton Ball", scale=1)
-                                org = gr.Textbox(label="🏢 Organization", placeholder="e.g., Bennet Family", scale=1)
-                            date = gr.Textbox(label="🗓️ Date", placeholder="e.g., 1811")
-                            entity_inputs.extend([person, location, event, org, date])
-                
-                with gr.Accordion("Records 5-6 (Optional)", 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", scale=1)
-                                location = gr.Textbox(label="📍 Location", scale=1)
-                            with gr.Row():
-                                event = gr.Textbox(label="📅 Event", scale=1)
-                                org = gr.Textbox(label="🏢 Organization", scale=1)
-                            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()
+            with gr.Column():
+                gr.Markdown("**Relationship 1**")
+                src1 = gr.Dropdown(label="From", choices=[])
+                rel1 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="related_to")
+                tgt1 = gr.Dropdown(label="To", choices=[])
+                relationship_inputs.extend([src1, rel1, tgt1])
             
-            # RIGHT COLUMN: Relationships
-            with gr.Column(scale=1):
-                gr.Markdown("## 🤝 Step 2: Define Relationships")
-                gr.Markdown("*First click 'Identify Entities' to populate the dropdowns*")
-                
-                # Create relationship inputs with explicit variable names
-                source1 = gr.Dropdown(label="From", choices=[], scale=2)
-                with gr.Row():
-                    rel_type1 = gr.Dropdown(label="Relationship Type", choices=RELATIONSHIP_TYPES, value="related_to", scale=2)
-                    target1 = gr.Dropdown(label="To", choices=[], scale=2)
-                
-                gr.HTML("<hr style='margin: 10px 0; border-top: 1px dashed #ccc;'>")
-                
-                source2 = gr.Dropdown(label="From", choices=[], scale=2)
-                with gr.Row():
-                    rel_type2 = gr.Dropdown(label="Relationship Type", choices=RELATIONSHIP_TYPES, value="related_to", scale=2)
-                    target2 = gr.Dropdown(label="To", choices=[], scale=2)
-                
-                gr.HTML("<hr style='margin: 10px 0; border-top: 1px dashed #ccc;'>")
-                
-                source3 = gr.Dropdown(label="From", choices=[], scale=2)
-                with gr.Row():
-                    rel_type3 = gr.Dropdown(label="Relationship Type", choices=RELATIONSHIP_TYPES, value="related_to", scale=2)
-                    target3 = gr.Dropdown(label="To", choices=[], scale=2)
-                
-                gr.HTML("<hr style='margin: 10px 0; border-top: 1px dashed #ccc;'>")
-                
-                source4 = gr.Dropdown(label="From", choices=[], scale=2)
-                with gr.Row():
-                    rel_type4 = gr.Dropdown(label="Relationship Type", choices=RELATIONSHIP_TYPES, value="related_to", scale=2)
-                    target4 = gr.Dropdown(label="To", choices=[], scale=2)
-                
-                gr.HTML("<hr style='margin: 10px 0; border-top: 1px dashed #ccc;'>")
-                
-                source5 = gr.Dropdown(label="From", choices=[], scale=2)
-                with gr.Row():
-                    rel_type5 = gr.Dropdown(label="Relationship Type", choices=RELATIONSHIP_TYPES, value="related_to", scale=2)
-                    target5 = gr.Dropdown(label="To", choices=[], scale=2)
-        
-        # Collect relationship inputs
-        relationship_inputs = [
-            source1, rel_type1, target1,
-            source2, rel_type2, target2,
-            source3, rel_type3, target3,
-            source4, rel_type4, target4,
-            source5, rel_type5, target5
-        ]
+            with gr.Column():
+                gr.Markdown("**Relationship 2**")
+                src2 = gr.Dropdown(label="From", choices=[])
+                rel2 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="related_to")
+                tgt2 = gr.Dropdown(label="To", choices=[])
+                relationship_inputs.extend([src2, rel2, tgt2])
+            
+            with gr.Column():
+                gr.Markdown("**Relationship 3**")
+                src3 = gr.Dropdown(label="From", choices=[])
+                rel3 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="related_to")
+                tgt3 = gr.Dropdown(label="To", choices=[])
+                relationship_inputs.extend([src3, rel3, tgt3])
         
-        gr.HTML("<hr style='margin: 30px 0;'>")
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("**Relationship 4**")
+                src4 = gr.Dropdown(label="From", choices=[])
+                rel4 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="related_to")
+                tgt4 = gr.Dropdown(label="To", choices=[])
+                relationship_inputs.extend([src4, rel4, tgt4])
+            
+            with gr.Column():
+                gr.Markdown("**Relationship 5**")
+                src5 = gr.Dropdown(label="From", choices=[])
+                rel5 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="related_to")
+                tgt5 = gr.Dropdown(label="To", choices=[])
+                relationship_inputs.extend([src5, rel5, tgt5])
+            
+            with gr.Column():
+                # Empty column for balance, or tips
+                gr.HTML("""
+                <div style="background: #e7f3ff; padding: 15px; border-radius: 8px; margin-top: 25px;">
+                    <strong>💡 Tip:</strong> Click "Identify Entities" first to populate these dropdowns!
+                </div>
+                """)
         
-        # Generate button
-        generate_btn = gr.Button("🎨 Generate Network Graph", variant="primary", size="lg")
+        gr.HTML("<hr style='margin: 20px 0;'>")
         
-        # Output section
-        gr.Markdown("## 📊 Step 3: View Results")
+        # ==================== STEP 3: GENERATE & VIEW ====================
+        gr.Markdown("## 🎨 Step 3: Generate Network Graph")
         
+        generate_btn = gr.Button("🎨 Generate Network Graph", variant="primary", size="lg")
+        
+        # Full-width network graph
         with gr.Row():
-            with gr.Column(scale=2):
-                network_plot = gr.Plot(label="Interactive Network Graph")
+            with gr.Column(scale=3):
+                network_plot = gr.HTML(label="Interactive Network Graph")
             with gr.Column(scale=1):
-                network_stats = gr.Markdown()
+                network_stats = gr.HTML()
+        
+        # Color legend
+        gr.HTML("""
+        <div style="background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); padding: 20px; border-radius: 10px; margin-top: 20px;">
+            <h4 style="color: white; margin: 0 0 15px 0;">🎨 Entity Color Legend</h4>
+            <div style="display: flex; flex-wrap: wrap; gap: 20px;">
+                <span style="display: flex; align-items: center; gap: 8px; color: white;">
+                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #00B894; display: inline-block; border: 2px solid white;"></span>
+                    Person
+                </span>
+                <span style="display: flex; align-items: center; gap: 8px; color: white;">
+                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #A0E7E5; display: inline-block; border: 2px solid white;"></span>
+                    Location
+                </span>
+                <span style="display: flex; align-items: center; gap: 8px; color: white;">
+                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #4ECDC4; display: inline-block; border: 2px solid white;"></span>
+                    Event
+                </span>
+                <span style="display: flex; align-items: center; gap: 8px; color: white;">
+                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #55A3FF; display: inline-block; border: 2px solid white;"></span>
+                    Organization
+                </span>
+                <span style="display: flex; align-items: center; gap: 8px; color: white;">
+                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #FF6B6B; display: inline-block; border: 2px solid white;"></span>
+                    Date
+                </span>
+            </div>
+            <p style="color: #888; margin: 15px 0 0 0; font-size: 13px;">
+                🖱️ <strong>Interaction:</strong> Drag nodes to rearrange • Scroll to zoom • Hover for details
+            </p>
+        </div>
+        """)
         
-        # Wire up the example buttons
+        # ==================== WIRE UP EVENTS ====================
+        
+        # Example buttons
         austen_btn.click(
             fn=load_austen_example,
             inputs=[],
@@ -592,21 +707,21 @@ def create_interface():
             outputs=entity_inputs
         )
         
-        # Wire up collect entities button
+        # Collect entities
         collect_btn.click(
             fn=collect_entities_from_records,
             inputs=entity_inputs,
             outputs=[
                 entity_summary,
-                source1, target1,
-                source2, target2,
-                source3, target3,
-                source4, target4,
-                source5, target5
+                src1, tgt1,
+                src2, tgt2,
+                src3, tgt3,
+                src4, tgt4,
+                src5, tgt5
             ]
         )
         
-        # Wire up generate graph button
+        # Generate graph
         all_inputs = entity_inputs + relationship_inputs
         generate_btn.click(
             fn=generate_network_graph,
@@ -614,41 +729,12 @@ def create_interface():
             outputs=[network_plot, network_stats]
         )
         
-        # Color legend
-        gr.HTML("""
-        <div style="background-color: #f8f9fa; padding: 15px; border-radius: 8px; margin-top: 20px;">
-            <h4 style="margin-top: 0;">🎨 Entity Color Legend</h4>
-            <div style="display: flex; flex-wrap: wrap; gap: 15px; align-items: center;">
-                <span style="display: flex; align-items: center; gap: 5px;">
-                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #00B894; display: inline-block;"></span>
-                    <strong>Person</strong>
-                </span>
-                <span style="display: flex; align-items: center; gap: 5px;">
-                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #A0E7E5; display: inline-block;"></span>
-                    <strong>Location</strong>
-                </span>
-                <span style="display: flex; align-items: center; gap: 5px;">
-                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #4ECDC4; display: inline-block;"></span>
-                    <strong>Event</strong>
-                </span>
-                <span style="display: flex; align-items: center; gap: 5px;">
-                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #55A3FF; display: inline-block;"></span>
-                    <strong>Organization</strong>
-                </span>
-                <span style="display: flex; align-items: center; gap: 5px;">
-                    <span style="width: 20px; height: 20px; border-radius: 50%; background-color: #FF6B6B; display: inline-block;"></span>
-                    <strong>Date</strong>
-                </span>
-            </div>
-        </div>
-        """)
-        
         # Footer
         gr.HTML("""
         <hr style="margin: 40px 0 20px 0;">
         <div style="text-align: center; color: #666; font-size: 14px; padding: 20px;">
-            <p><strong>Basic Network Explorer</strong> | Bodleian Libraries (University of Oxford) Sassoon Research Fellowship</p>
-            <p>Built with the aid of Claude</p>
+            <p><strong>Network Explorer</strong> | Bodleian Libraries, University of Oxford</p>
+            <p style="color: #888; font-size: 12px;">Built with the aid of Claude</p>
         </div>
         """)