Update app.py

app.py

@@ -6,6 +6,8 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from pyvis.network import Network
+import networkx as nx
 
 DESCRIPTION = """
 # GWQ PREV
@@ -27,6 +29,19 @@ model = AutoModelForCausalLM.from_pretrained(
 model.config.sliding_window = 4096
 model.eval()
 
+def create_knowledge_graph(text):
+    # Simple example: Create a 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"
 
 @spaces.GPU(duration=120)
 def generate(
@@ -37,6 +52,7 @@ 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})
@@ -67,6 +83,10 @@ 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}"
 
 demo = gr.ChatInterface(
     fn=generate,
@@ -106,13 +126,13 @@ demo = gr.ChatInterface(
             step=0.05,
             value=1.2,
         ),
+        gr.Checkbox(label="Visualize Knowledge Graph", value=False),
     ],
     stop_btn=None,
     examples=[
         ["Write a Python function to reverses a string if it's length is a multiple of 4. def reverse_string(str1): if len(str1) % 4 == 0: return ''.join(reversed(str1)) return str1 print(reverse_string('abcd')) print(reverse_string('python')) "],
         ["Rectangle $ABCD$ is the base of pyramid $PABCD$. If $AB = 10$, $BC = 5$, $\overline{PA}\perp \text{plane } ABCD$, and $PA = 8$, then what is the volume of $PABCD$?"],
         ["Difference between List comprehension and Lambda in Python lst  =  [x ** 2  for x in range (1, 11)   if  x % 2 == 1] print(lst)"],
-        ["How many hours does it take a man to eat a Helicopter?"],
         ["How Many R's in the Word 'STRAWBERRY' ?"],
     ],
     cache_examples=False,
@@ -122,6 +142,5 @@ demo = gr.ChatInterface(
     fill_height=True,
 )
 
-
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()