Updated knowledge graph

app.py

@@ -471,6 +471,15 @@ for i in range(len(t)//3):
   r.append(t[3*i:3*i+3])
 r
 
+def get_edge_labels(t:list):
+  dct = {}
+  length_of_t = len(t)
+  for i in range(length_of_t):
+    t[i][0] = t[i][0].replace('"',"").replace("'","").strip()
+    t[i][2] = t[i][2].replace('"',"").replace("'","").strip()
+    t[i][1] = t[i][1].replace('"',"").replace("'","")
+    dct[(t[i][0],t[i][2] )] =  t[i][1]
+  return dct
 def knowledge_graph(prompt): 
   
   response = openai.Completion.create(
@@ -493,27 +502,42 @@ def knowledge_graph(prompt):
     t.append(r[3*i:3*i+3])
   # t = [['"Ankit"', '"went_to"', '"market"'], ['"Ankit"', '"called"', '"Raj"']]
   import networkx as nx
+  import random
   print(t)
   G = nx.Graph()
+  new_nodes = []
+  print('Edge labels')
+  edge_labels = get_edge_labels(t)
+  print(edge_labels)
+  print(f't after edge labesl = {t}')
   for i in t:
-    G.add_node(i[0])
-    G.add_node(i[2])
+    if not i[0] in new_nodes:
+      new_nodes.append(i[0])
+      G.add_node(i[0])
+    if not i[2] in new_nodes:
+      new_nodes.append(i[2])
+      G.add_node(i[2])
+    # G.add_node(i[0])
+    # G.add_node(i[2])
     G.add_edge(i[0],i[2])
   pos = nx.spring_layout(G)
   nx.draw(G,pos,labels={node: node for node in G.nodes()})
-  print('Drew')
+  
   x = nx.draw_networkx_edge_labels(
       G, pos,
-      edge_labels={(t[0][0], t[0][2]): t[0][1], 
-                  (t[1][0], t[1][2]): t[1][1]
-                  },
+      edge_labels=edge_labels,
       font_color='red'
   )
   # print(x)
-  plt.savefig("/tmp/generate_graph.png")
-  img = Image.open("/tmp/generate_graph.png")
-  os.remove("/tmp/generate_graph.png")
-  return img
+  random_name = f'generated_img_{random.randint(1,100000)}.png'
+  plt.savefig(f"/tmp/{random_name}")
+  plt.clf()
+  img = Image.open(f"/tmp/{random_name}")
+  os.remove(f"/tmp/{random_name}")
+  
+  return np.asarray(img)
+  
+c =knowledge_graph("Alice went to office. Called bob. Went to grocery shopping. Then went home")
   
 
 
@@ -653,7 +677,7 @@ home</p>""")
     hbtn.click(video, inputs=fil1, outputs=vid2)
       
   with gr.Accordion("Generate Knowledge Graph",open=False):
-    gr.Markdown("**Genrate KNowledge Graph**")
+    gr.Markdown("**Genrate Knowledge Graph**")
     with gr.Row():
       with gr.Column():
         prompt_to_generate_graph = gr.Text()