Update app.py

app.py

@@ -2,39 +2,39 @@ from GPT4KG import KnowledgeGraph
 import gradio as gr
 from PIL import Image
 
-def generate_graph(input_text,api_key):
-  try:
-    kg = KnowledgeGraph(api_key,"temp.kg")
-  except:
+def generate_graph(input_text,api_key,graph):
+  if graph[0] =="":
     kg = KnowledgeGraph(api_key)
+    graph[0] = kg
+  else:
+    kg = graph[0]
   kg.learn(str(input_text))
-  kg.save_graph("temp.kg")
-  kg.display_graph("temp.png")
-  return Image.open("temp.png")
+  img = kg.display_graph()
+  return img,graph
 
-def answer_question(question,api_key):
-  try:
-    kg = KnowledgeGraph(api_key,"temp.kg")
-  except:
+def answer_question(question,api_key,graph):
+  if graph[0] =="":
     kg = KnowledgeGraph(api_key)
+    graph[0] = kg
+  else:
+    kg = graph[0]
   return kg.chat_qa(question)
 
 title = "GPT-4 Knowledge Graph Generator"
 description = "Enter text to generate a knowledge graph using GPT4KG:"
 
 with gr.Blocks() as demo:
-  with open("temp.kg","w") as f:
-    f.write("")
   gr.Markdown("""<h1><center>GPT-4 Knowledge Graph Generator</center></h1>""")
   output_image = gr.Image(label="Knowledge Graph", type="pil")
   api_key = gr.Textbox(lines=1, label="OpenAI API Key")
+  graph = gr.State([""])
   input_text = gr.Textbox(lines=5, label="Information to be added to graph")
   submit_btn = gr.Button("Add info to graph")
-  submit_btn.click(fn=generate_graph, inputs=[input_text,api_key], outputs=[output_image])
+  submit_btn.click(fn=generate_graph, inputs=[input_text,api_key,graph], outputs=[output_image,graph])
 
   question = gr.Textbox(lines=1, label="Question about the info in this graph")
   answer = gr.Textbox(lines=1, label="Answer")
   qa_btn = gr.Button("Ask question")
-  qa_btn.click(fn=answer_question, inputs=[question,api_key], outputs=[answer])
+  qa_btn.click(fn=answer_question, inputs=[question,api_key,graph], outputs=[answer])
 
 demo.launch()