Swapped to bokeh, changed up samples

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -5,12 +5,11 @@ import gradio as gr
 import spaces
 
 from lib.graph_extract import triplextract, parse_triples
-from lib.visualize import create_cytoscape_plot
+from lib.visualize import create_bokeh_plot #, create_plotly_plot
 from lib.samples import snippets
 
 WORD_LIMIT = 300
 
-@spaces.GPU
 def process_text(text, entity_types, predicates):
     if not text:
         return None, "Please enter some text."
@@ -40,7 +39,7 @@ def process_text(text, entity_types, predicates):
                 "No entities or relationships found. Try different text or check your input.",
             )
 
-        fig = create_cytoscape_plot(entities, relationships)
+        fig = create_bokeh_plot(entities, relationships)
         return (
             fig,
             f"Entities: {entities}\nRelationships: {relationships}\n\nRaw output:\n{prediction}",

lib/graph_extract.py

@@ -54,6 +54,8 @@ print("Model and tokenizer loaded successfully.")
 generation_config = GenerationConfig.from_pretrained("sciphi/triplex")
 generation_config.max_length = 2048
 generation_config.pad_token_id = tokenizer.eos_token_id
+
+
 @spaces.GPU
 def triplextract(text, entity_types, predicates):
     input_format = """Perform Named Entity Recognition (NER) and extract knowledge graph triplets from the text. NER identifies named entities of given entity types, and triple extraction identifies relationships between entities using specified predicates. Return the result as a JSON object with an "entities_and_triples" key containing an array of entities and triples.

lib/samples.py

@@ -19,8 +19,8 @@ snippets = {
         we were all going direct to Heaven, we were all going direct the other way – in short, 
         the period was so far like the present period, that some of its noisiest authorities 
         insisted on its being received, for good or for evil, in the superlative degree of comparison only.""",
-        entity_types="TIME, EMOTION, LOCATION, EVENT, OUTCOME, PLACE",
-        predicates="WAS, HAD, WERE"
+        entity_types="EMOTION, EVENT, OUTCOME, PLACE",
+        predicates="WAS, HAD, WERE, IS"
     ),
 
     'tech_company': Snippet(
@@ -29,8 +29,8 @@ snippets = {
         software, and online services. The company's flagship products include the iPhone smartphone, 
         iPad tablet, and Mac personal computer. As of 2023, Apple has over 150,000 employees worldwide 
         and generates annual revenue exceeding $350 billion.""",
-        entity_types="COMPANY, PERSON, PRODUCT, LOCATION, DATE, NUMBER, EVENT, SUBJECT",
-        predicates="FOUNDED, HEADQUARTERED_IN, PRODUCES, HAS, EMPLOYEES, "
+        entity_types="COMPANY, PERSON, PRODUCT, LOCATION, DATE, EVENT",
+        predicates="FOUNDED, PRODUCES, HAS, IN, EMPLOYEES, "
     ),
 
     'climate_change': Snippet(

lib/visualize.py

@@ -1,10 +1,92 @@
 import plotly.graph_objects as go
 import networkx as nx
 
-import plotly.graph_objects as go
 import networkx as nx
+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
+    G = nx.Graph()
+    for entity_id, entity_data in entities.items():
+        G.add_node(entity_id, label=f"{entity_data['value']} ({entity_data['type']})")
+    for source, relation, target in relationships:
+        G.add_edge(source, target, label=relation)
+
+    plot = Plot(width=600, height=600,  # Increased size for better visibility
+                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))
+
+    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])
+    graph_renderer.node_renderer.hover_glyph = Scatter(size=15, fill_color=Spectral4[1])
+
+    graph_renderer.edge_renderer.glyph = MultiLine(line_color="#000", line_alpha=0.9, line_width=3)
+    graph_renderer.edge_renderer.selection_glyph = MultiLine(line_color=Spectral4[2], line_width=4)
+    graph_renderer.edge_renderer.hover_glyph = MultiLine(line_color=Spectral4[1], line_width=3)
+
+    graph_renderer.selection_policy = NodesAndLinkedEdges()
+    graph_renderer.inspection_policy = NodesAndLinkedEdges()
+
+    plot.renderers.append(graph_renderer)
+
+    # Add node labels
+    x, y = zip(*graph_renderer.layout_provider.graph_layout.values())
+    node_labels = nx.get_node_attributes(G, 'label')
+    source = ColumnDataSource({'x': x, 'y': y, 'label': [node_labels[node] for node in G.nodes()]})
+    labels = LabelSet(x='x', y='y', text='label', source=source, background_fill_color='white',
+                      text_font_size='8pt', background_fill_alpha=0.7)
+    plot.renderers.append(labels)
+
+    # Add 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]
+        edge_x.append((start_x + end_x) / 2)
+        edge_y.append((start_y + end_y) / 2)
+        edge_labels.append(label)
+
+    edge_label_source = ColumnDataSource({'x': edge_x, 'y': edge_y, 'label': edge_labels})
+    edge_labels = LabelSet(x='x', y='y', text='label', source=edge_label_source,
+                           background_fill_color='white', text_font_size='8pt',
+                           background_fill_alpha=0.7)
+    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_cytoscape_plot(entities, relationships):
+def create_plotly_plot(entities, relationships):
     G = nx.DiGraph()  # Use DiGraph for directed edges
 
     for entity_id, entity_data in entities.items():

requirements.txt

@@ -4,4 +4,5 @@ matplotlib==3.7.2
 torch==2.0.1
 transformers==4.43.3
 accelerate==0.33.0
-networkx
+networkx
+bokeh==3.5.1