switching from graphviz to networkx

app.py

@@ -15,7 +15,6 @@ from dotenv import load_dotenv
 from huggingface_hub import hf_hub_download
 import spaces
 import traceback
-from graphviz import Digraph
 from PIL import Image, ImageDraw, ImageFont
 from io import BytesIO
 import functools
@@ -320,39 +319,39 @@ def add_nodes_edges(dot, node, config, max_weight, min_weight, parent=None, is_r
         for child in node.get('children', []):
             add_nodes_edges(dot, child, config, max_weight, min_weight, parent=node_id, is_root=False, depth=depth+1, trim_cutoff=trim_cutoff)
 
+import networkx as nx
+import matplotlib.pyplot as plt
+
+def add_nodes_edges_nx(G, node, parent=None, is_root=True):
+    node_id = str(id(node))
+    token = node.get('token', '').strip()
+
+    if is_root or token:
+        G.add_node(node_id, label=token if not is_root else "*")
+        if parent:
+            G.add_edge(parent, node_id)
+
+        for child in node.get('children', []):
+            add_nodes_edges_nx(G, child, parent=node_id, is_root=False)
+
 def create_tree_diagram(data, config, max_weight, min_weight, trim_cutoff=0):
-    import os
-    from graphviz import Digraph
+    G = nx.DiGraph()
+    add_nodes_edges_nx(G, data)
 
-    # Ensure the system can find the dot command
-    os.environ["PATH"] += os.pathsep + "/usr/bin"
+    # Draw the tree using matplotlib
+    plt.figure(figsize=(12, 12))
+    pos = nx.spring_layout(G, k=0.5, iterations=50)
+    labels = nx.get_node_attributes(G, 'label')
+    nx.draw(G, pos, labels=labels, with_labels=True, node_size=5000, node_color="lightblue", font_size=10, font_weight="bold", edge_color="gray", arrows=False)
 
-    # Initialize the Digraph object and explicitly set the dot command path
-    dot = Digraph(comment='Definition Tree', format='png', engine='dot')
-    dot.command = '/usr/bin/dot'
-    
-    dot.attr(rankdir='LR', size='5040,5000', margin='0.06', nodesep='0.06', ranksep='1', dpi='120', bgcolor='white')
+    # Save the image to a BytesIO object
+    output = BytesIO()
+    plt.savefig(output, format='png')
+    plt.close()
+    output.seek(0)
 
-    add_nodes_edges(dot, data, config, max_weight, min_weight, trim_cutoff=trim_cutoff)
+    return output
 
-    # Save to a temporary file first
-    temp_filename = "temp_tree_diagram"
-    dot.render(temp_filename, format='png', cleanup=True)
-    
-    # Read the file back into a BytesIO object
-    with open(f"{temp_filename}.png", "rb") as f:
-        output = BytesIO(f.read())
-
-    # Add white background
-    with Image.open(output) as img:
-        bg = Image.new("RGB", (img.width, 5000), (255, 255, 255))
-        y_offset = (5000 - img.height) // 2
-        bg.paste(img, (0, y_offset))
-        final_output = BytesIO()
-        bg.save(final_output, 'PNG')
-        final_output.seek(0)
-
-    return final_output
 
 
 

requirements.txt

@@ -12,6 +12,5 @@ spaces
 graphviz
 Pillow
 accelerate
+networkx
 #fnord
-#fnord
-#fnord