Multiplot support, bokeh and plotly, multiple graph layout support.

app.py

@@ -1,52 +1,66 @@
-
 import random
 
 import gradio as gr
-import spaces
+import networkx as nx
 
 from lib.graph_extract import triplextract, parse_triples
-from lib.visualize import create_bokeh_plot #, create_plotly_plot
+from lib.visualize import create_graph, create_bokeh_plot, create_plotly_plot
 from lib.samples import snippets
 
 WORD_LIMIT = 300
 
-def process_text(text, entity_types, predicates):
+def process_text(text, entity_types, predicates, layout_type, visualization_type):
     if not text:
-        return None, "Please enter some text."
+        return None, None, "Please enter some text."
 
     words = text.split()
     if len(words) > WORD_LIMIT:
-        return None, f"Please limit your input to {WORD_LIMIT} words. Current word count: {len(words)}"
+        return None, None, f"Please limit your input to {WORD_LIMIT} words. Current word count: {len(words)}"
 
     entity_types = [et.strip() for et in entity_types.split(",") if et.strip()]
     predicates = [p.strip() for p in predicates.split(",") if p.strip()]
 
     if not entity_types:
-        return None, "Please enter at least one entity type."
+        return None, None, "Please enter at least one entity type."
     if not predicates:
-        return None, "Please enter at least one predicate."
+        return None, None, "Please enter at least one predicate."
 
     try:
         prediction = triplextract(text, entity_types, predicates)
         if prediction.startswith("Error"):
-            return None, prediction
+            return None, None, prediction
 
         entities, relationships = parse_triples(prediction)
 
         if not entities and not relationships:
-            return (
-                None,
-                "No entities or relationships found. Try different text or check your input.",
-            )
-
-        fig = create_bokeh_plot(entities, relationships)
-        return (
-            fig,
-            f"Entities: {entities}\nRelationships: {relationships}\n\nRaw output:\n{prediction}",
-        )
+            return None, None, "No entities or relationships found. Try different text or check your input."
+
+        G = create_graph(entities, relationships)
+
+        if visualization_type == 'Bokeh':
+            fig = create_bokeh_plot(G, layout_type)
+        else:
+            fig = create_plotly_plot(G, layout_type)
+
+        output_text = f"Entities: {entities}\nRelationships: {relationships}\n\nRaw output:\n{prediction}"
+        return G, fig, output_text
     except Exception as e:
-        print(f"Error in process_text: {e}")
-        return None, f"An error occurred: {str(e)}"
+        print(f"Error in process_text: {str(e)}")
+        return None, None, f"An error occurred: {str(e)}"
+
+def update_graph(G, layout_type, visualization_type):
+    if G is None:
+        return None, "Please process text first."
+    
+    try:
+        if visualization_type == 'Bokeh':
+            fig = create_bokeh_plot(G, layout_type)
+        else:
+            fig = create_plotly_plot(G, layout_type)
+        return fig, ""
+    except Exception as e:
+        print(f"Error in update_graph: {e}")
+        return None, f"An error occurred while updating the graph: {str(e)}"
 
 def update_inputs(sample_name):
     sample = snippets[sample_name]
@@ -60,34 +74,42 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
     
     with gr.Row():
         with gr.Column(scale=1):
-            sample_dropdown = gr.Dropdown(
-                choices=list(snippets.keys()),
-                label="Select Sample",
-                value=default_sample_name
-            )
-            input_text = gr.Textbox(
-                label="Input Text",
-                lines=5,
-                value=default_sample.text_input
-            )
+            sample_dropdown = gr.Dropdown(choices=list(snippets.keys()), label="Select Sample", value=default_sample_name)
+            input_text = gr.Textbox(label="Input Text", lines=5, value=default_sample.text_input)
             entity_types = gr.Textbox(label="Entity Types", value=default_sample.entity_types)
             predicates = gr.Textbox(label="Predicates", value=default_sample.predicates)
-            submit_btn = gr.Button("Extract Knowledge Graph")
+            layout_type = gr.Dropdown(choices=['spring', 'fruchterman_reingold', 'circular', 'random', 'spectral', 'shell'], 
+                                      label="Layout Type", value='spring')
+            visualization_type = gr.Radio(choices=['Bokeh', 'Plotly'], label="Visualization Type", value='Bokeh')
+            process_btn = gr.Button("Process Text")
         with gr.Column(scale=2):
             output_graph = gr.Plot(label="Knowledge Graph")
             error_message = gr.Textbox(label="Textual Output")
 
-    sample_dropdown.change(
-        update_inputs,
-        inputs=[sample_dropdown],
-        outputs=[input_text, entity_types, predicates]
-    )
+    graph_state = gr.State(None)
+
+    def process_and_update(text, entity_types, predicates, layout_type, visualization_type):
+        G, fig, output = process_text(text, entity_types, predicates, layout_type, visualization_type)
+        return G, fig, output
+
+    def update_graph_wrapper(G, layout_type, visualization_type):
+        if G is not None:
+            fig, _ = update_graph(G, layout_type, visualization_type)
+            return fig
 
-    submit_btn.click(
-        process_text,
-        inputs=[input_text, entity_types, predicates],
-        outputs=[output_graph, error_message],
-    )
+    sample_dropdown.change(update_inputs, inputs=[sample_dropdown], outputs=[input_text, entity_types, predicates])
+
+    process_btn.click(process_and_update,
+                      inputs=[input_text, entity_types, predicates, layout_type, visualization_type],
+                      outputs=[graph_state, output_graph, error_message])
+    
+    layout_type.change(update_graph_wrapper,
+                       inputs=[graph_state, layout_type, visualization_type],
+                       outputs=[output_graph])
+    
+    visualization_type.change(update_graph_wrapper,
+                              inputs=[graph_state, layout_type, visualization_type],
+                              outputs=[output_graph])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)

lib/visualize.py

@@ -1,30 +1,55 @@
 import plotly.graph_objects as go
 import networkx as nx
-
-import networkx as nx
+import numpy as np
 from bokeh.models import (BoxSelectTool, HoverTool, MultiLine, NodesAndLinkedEdges, 
                           Plot, Range1d, Scatter, TapTool, LabelSet, ColumnDataSource)
 from bokeh.palettes import Spectral4
 from bokeh.plotting import from_networkx
 
-def create_bokeh_plot(entities, relationships):
-    # Create a NetworkX graph
+def create_graph(entities, relationships):
     G = nx.Graph()
     for entity_id, entity_data in entities.items():
-        G.add_node(entity_id, label=f"{entity_data['value']} ({entity_data['type']})")
+        G.add_node(entity_id, label=f"{entity_data.get('value', 'Unknown')} ({entity_data.get('type', 'Unknown')})")
+    
     for source, relation, target in relationships:
         G.add_edge(source, target, label=relation)
-
-    plot = Plot(width=600, height=600,  # Increased size for better visibility
+    
+    return G
+
+def improved_spectral_layout(G, scale=1):
+    pos = nx.spectral_layout(G)
+    # Add some random noise to prevent overlapping
+    pos = {node: (x + np.random.normal(0, 0.1), y + np.random.normal(0, 0.1)) for node, (x, y) in pos.items()}
+    # Scale the layout
+    pos = {node: (x * scale, y * scale) for node, (x, y) in pos.items()}
+    return pos
+
+def create_bokeh_plot(G, layout_type='spring'):
+    plot = Plot(width=600, height=600,
                 x_range=Range1d(-1.2, 1.2), y_range=Range1d(-1.2, 1.2))
     plot.title.text = "Knowledge Graph Interaction"
 
-    # Use tooltips to show node and edge labels on hover
     node_hover = HoverTool(tooltips=[("Entity", "@label")])
     edge_hover = HoverTool(tooltips=[("Relation", "@label")])
     plot.add_tools(node_hover, edge_hover, TapTool(), BoxSelectTool())
 
-    graph_renderer = from_networkx(G, nx.spring_layout, scale=1,k=0.5, iterations=50, center=(0, 0))
+    # Create layout based on layout_type
+    if layout_type == 'spring':
+        pos = nx.spring_layout(G, k=0.5, iterations=50)
+    elif layout_type == 'fruchterman_reingold':
+        pos = nx.fruchterman_reingold_layout(G, k=0.5, iterations=50)
+    elif layout_type == 'circular':
+        pos = nx.circular_layout(G)
+    elif layout_type == 'random':
+        pos = nx.random_layout(G)
+    elif layout_type == 'spectral':
+        pos = improved_spectral_layout(G)
+    elif layout_type == 'shell':
+        pos = nx.shell_layout(G)
+    else:
+        pos = nx.spring_layout(G, k=0.5, iterations=50)
+
+    graph_renderer = from_networkx(G, pos, scale=1, center=(0, 0))
 
     graph_renderer.node_renderer.glyph = Scatter(size=15, fill_color=Spectral4[0])
     graph_renderer.node_renderer.selection_glyph = Scatter(size=15, fill_color=Spectral4[2])
@@ -48,9 +73,7 @@ def create_bokeh_plot(entities, relationships):
     plot.renderers.append(labels)
 
     # Add edge labels
-    edge_x = []
-    edge_y = []
-    edge_labels = []
+    edge_x, edge_y, edge_labels = [], [], []
     for (start_node, end_node, label) in G.edges(data='label'):
         start_x, start_y = graph_renderer.layout_provider.graph_layout[start_node]
         end_x, end_y = graph_renderer.layout_provider.graph_layout[end_node]
@@ -65,69 +88,30 @@ def create_bokeh_plot(entities, relationships):
     plot.renderers.append(edge_labels)
 
     return plot
-    
-# def create_bokeh_plot(entities, relationships):
-#     # Create a NetworkX graph
-#     G = nx.Graph()
-#     for entity_id, entity_data in entities.items():
-#         G.add_node(entity_id, **entity_data)
-#     for source, relation, target in relationships:
-#         G.add_edge(source, target)
-
-#     # Create a Bokeh plot
-#     plot = figure(title="Knowledge Graph", x_range=(-1.1,1.1), y_range=(-1.1,1.1),
-#                   width=400, height=400, tools="pan,wheel_zoom,box_zoom,reset")
 
-#     # Create graph renderer
-#     graph_renderer = from_networkx(G, nx.spring_layout, scale=1, center=(0,0))
-
-#     # Add graph renderer to plot
-#     plot.renderers.append(graph_renderer)
-
-#     return plot
-
-def create_plotly_plot(entities, relationships):
-    G = nx.DiGraph()  # Use DiGraph for directed edges
-
-    for entity_id, entity_data in entities.items():
-        G.add_node(entity_id, **entity_data)
-
-    for source, relation, target in relationships:
-        G.add_edge(source, target, relation=relation)
-
-    pos = nx.spring_layout(G, k=0.5, iterations=50)  # Adjust layout parameters
-
-    edge_trace = go.Scatter(
-        x=[],
-        y=[],
-        line=dict(width=1, color="#888"),
-        hoverinfo="text",
-        mode="lines",
-        text=[],
-    )
-
-    node_trace = go.Scatter(
-        x=[],
-        y=[],
-        mode="markers+text",
-        hoverinfo="text",
-        marker=dict(
-            showscale=True,
-            colorscale="Viridis",
-            reversescale=True,
-            color=[],
-            size=15,
-            colorbar=dict(
-                thickness=15,
-                title="Node Connections",
-                xanchor="left",
-                titleside="right",
-            ),
-            line_width=2,
-        ),
-        text=[],
-        textposition="top center",
-    )
+def create_plotly_plot(G, layout_type='spring'):
+    # Create layout based on layout_type
+    if layout_type == 'spring':
+        pos = nx.spring_layout(G, k=0.5, iterations=50)
+    elif layout_type == 'fruchterman_reingold':
+        pos = nx.fruchterman_reingold_layout(G, k=0.5, iterations=50)
+    elif layout_type == 'circular':
+        pos = nx.circular_layout(G)
+    elif layout_type == 'random':
+        pos = nx.random_layout(G)
+    elif layout_type == 'spectral':
+        pos = improved_spectral_layout(G)
+    elif layout_type == 'shell':
+        pos = nx.shell_layout(G)
+    else:
+        pos = nx.spring_layout(G, k=0.5, iterations=50)
+
+    edge_trace = go.Scatter(x=[], y=[], line=dict(width=1, color="#888"), hoverinfo="text", mode="lines", text=[])
+    node_trace = go.Scatter(x=[], y=[], mode="markers+text", hoverinfo="text",
+                            marker=dict(showscale=True, colorscale="Viridis", reversescale=True, color=[], size=15,
+                                        colorbar=dict(thickness=15, title="Node Connections", xanchor="left", titleside="right"),
+                                        line_width=2),
+                            text=[], textposition="top center")
 
     edge_labels = []
 
@@ -137,57 +121,26 @@ def create_plotly_plot(entities, relationships):
         edge_trace["x"] += (x0, x1, None)
         edge_trace["y"] += (y0, y1, None)
         
-        # Calculate midpoint for edge label
         mid_x, mid_y = (x0 + x1) / 2, (y0 + y1) / 2
-        edge_labels.append(
-            go.Scatter(
-                x=[mid_x],
-                y=[mid_y],
-                mode="text",
-                text=[G.edges[edge]["relation"]],
-                textposition="middle center",
-                hoverinfo="none",
-                showlegend=False,
-                textfont=dict(size=8),
-            )
-        )
+        edge_labels.append(go.Scatter(x=[mid_x], y=[mid_y], mode="text", text=[G.edges[edge]["label"]],
+                                      textposition="middle center", hoverinfo="none", showlegend=False, textfont=dict(size=8)))
 
     for node in G.nodes():
         x, y = pos[node]
         node_trace["x"] += (x,)
         node_trace["y"] += (y,)
-        node_info = f"{entities[node]['value']} ({entities[node]['type']})"
-        node_trace["text"] += (node_info,)
+        node_trace["text"] += (G.nodes[node]["label"],)
         node_trace["marker"]["color"] += (len(list(G.neighbors(node))),)
 
-    fig = go.Figure(
-        data=[edge_trace, node_trace] + edge_labels,
-        layout=go.Layout(
-            title="Knowledge Graph",
-            titlefont_size=16,
-            showlegend=False,
-            hovermode="closest",
-            margin=dict(b=20, l=5, r=5, t=40),
-            annotations=[],
-            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-            width=800,
-            height=600,
-        ),
-    )
-
-    # Enable dragging of nodes
-    fig.update_layout(
-        newshape=dict(line_color="#009900"),
-        # Enable zoom
-        xaxis=dict(
-            scaleanchor="y",
-            scaleratio=1,
-        ),
-        yaxis=dict(
-            scaleanchor="x",
-            scaleratio=1,
-        ),
-    )
+    fig = go.Figure(data=[edge_trace, node_trace] + edge_labels,
+                    layout=go.Layout(title="Knowledge Graph", titlefont_size=16, showlegend=False, hovermode="closest",
+                                     margin=dict(b=20, l=5, r=5, t=40), annotations=[],
+                                     xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                                     yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                                     width=800, height=600))
+
+    fig.update_layout(newshape=dict(line_color="#009900"),
+                      xaxis=dict(scaleanchor="y", scaleratio=1),
+                      yaxis=dict(scaleanchor="x", scaleratio=1))
 
     return fig