Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import ast
-from graphviz import Digraph
+import networkx as nx
+import matplotlib.pyplot as plt
 
 client = InferenceClient("Qwen/Qwen2.5-72B-Instruct")
 
@@ -81,15 +82,29 @@ def sampling(num_samples, num_associations):
                                        top_p=0.1
                                        ).choices[0].get('message')['content']))
 
-  dot = Digraph(comment=f'SynthNet, {num_samples} samples, {num_associations} associations', graph_attr={'rankdir': 'LR'})
+  G = nx.DiGraph()
 
   for entry in triplets:
     source, label, target = entry['properties']
-    dot.node(source, source)
-    dot.node(target, target)
-    dot.edge(source, target, label=label)
+    G.add_node(source, label=source)
+    G.add_node(target, label=target)
+    G.add_edge(source, target, label=label)
 
-  dot.render('synthnet', format='png')
+  pos = nx.spring_layout(G)
+  nx.draw_networkx_nodes(G, pos, node_size=500, node_color='lightblue')
+
+  edge_labels = nx.get_edge_attributes(G, 'label')  # Get edge labels
+  nx.draw_networkx_edges(G, pos, arrowstyle='->', arrowsize=25)
+  nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
+
+  node_labels = nx.get_node_attributes(G, 'label')  # Get node labels
+  nx.draw_networkx_labels(G, pos, labels=node_labels, font_size=8, font_family="sans-serif")
+
+  plt.axis('off')
+  plt.tight_layout()
+
+  plt.savefig('synthnet.png')
+  plt.close()
   return 'synthnet.png'
 
 demo = gr.Interface(
@@ -104,4 +119,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)