Update app.py

app.py

@@ -1,12 +1,13 @@
+import os
+from collections.abc import Iterator
+from threading import Thread
+import networkx as nx
+import matplotlib.pyplot as plt
+
 import gradio as gr
 import spaces
 import torch
-from collections.abc import Iterator
-from threading import Thread
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from pyvis.network import Network
-import networkx as nx
-import os
 
 DESCRIPTION = """
 # GWQ PREV
@@ -28,19 +29,20 @@ model = AutoModelForCausalLM.from_pretrained(
 model.config.sliding_window = 4096
 model.eval()
 
-def create_knowledge_graph(text):
-    # Simple example: Create a graph from the text
+def create_knowledge_graph_image(text):
+    # Example: Create a simple knowledge graph from the text
     G = nx.Graph()
     words = text.split()
     for i in range(len(words) - 1):
         G.add_edge(words[i], words[i + 1])
-    return G
-
-def visualize_knowledge_graph(graph):
-    net = Network(notebook=True, cdn_resources='in_line')
-    net.from_nx(graph)
-    net.show("knowledge_graph.html")
-    return "knowledge_graph.html"
+    
+    # Draw the graph
+    plt.figure(figsize=(8, 6))
+    pos = nx.spring_layout(G)
+    nx.draw(G, pos, with_labels=True, node_color='lightblue', edge_color='gray', node_size=2000, font_size=10)
+    plt.savefig("knowledge_graph.png")
+    plt.close()
+    return "knowledge_graph.png"
 
 @spaces.GPU(duration=120)
 def generate(
@@ -51,7 +53,6 @@ def generate(
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
-    visualize_graph: bool = False,
 ) -> Iterator[str]:
     conversation = chat_history.copy()
     conversation.append({"role": "user", "content": message})
@@ -82,10 +83,9 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
-    if visualize_graph:
-        graph = create_knowledge_graph("".join(outputs))
-        graph_file = visualize_knowledge_graph(graph)
-        yield f"Knowledge graph saved to {graph_file}"
+    # After generating the text, create the knowledge graph image
+    knowledge_graph_image = create_knowledge_graph_image("".join(outputs))
+    yield knowledge_graph_image
 
 demo = gr.ChatInterface(
     fn=generate,
@@ -125,7 +125,6 @@ demo = gr.ChatInterface(
             step=0.05,
             value=1.2,
         ),
-        gr.Checkbox(label="Visualize Knowledge Graph", value=False),
     ],
     stop_btn=None,
     examples=[
@@ -139,6 +138,7 @@ demo = gr.ChatInterface(
     description=DESCRIPTION,
     css_paths="style.css",
     fill_height=True,
+    additional_outputs=[gr.Image(label="Knowledge Graph")]
 )
 
 if __name__ == "__main__":