Update app.py

app.py

@@ -437,45 +437,75 @@ def mask(sentence):
 
 
 
-
 #plotly tree
+import plotly.graph_objs as go
+import textwrap
+import re
+from collections import defaultdict
+
 def generate_plot(original_sentence):
     paraphrased_sentences = generate_paraphrase(original_sentence)
     first_paraphrased_sentence = paraphrased_sentences[0]
     masked_sentence = mask_non_stopword(first_paraphrased_sentence)
     masked_versions = mask(masked_sentence)
+    
     nodes = []
     nodes.append(original_sentence)
     nodes.extend(paraphrased_sentences)
     nodes.extend(masked_versions)
-    edges = [
-        (0, 1),
-        (0, 2),
-        (0, 3),
-        (0, 4),
-        (0, 5),
-        (1, 6),
-        (1, 7),
-        (1, 8),
-        (1, 9),
-        (1, 10)
-    ]
-    wrapped_nodes = ['<br>'.join(textwrap.wrap(node, width=30)) for node in nodes]
+    nodes[0] += ' L0'
+    para_len = len(paraphrased_sentences)
+    for i in range(1, para_len+1):
+        nodes[i] += ' L1'
+    for i in range(para_len+1, len(nodes)):
+        nodes[i] += ' L2'
+    
+    cleaned_nodes = [re.sub(r'\sL[0-9]$', '', node) for node in nodes]
+    wrapped_nodes = ['<br>'.join(textwrap.wrap(node, width=30)) for node in cleaned_nodes]
+    
+    def get_levels_and_edges(nodes):
+        levels = {}
+        edges = []
+        for i, node in enumerate(nodes):
+            level = int(node.split()[-1][1])
+            levels[i] = level
+
+        # Add edges from L0 to all L1 nodes
+        root_node = next(i for i, level in levels.items() if level == 0)
+        for i, level in levels.items():
+            if level == 1:
+                edges.append((root_node, i))
+
+        # Identify the first L1 node
+        first_l1_node = next(i for i, level in levels.items() if level == 1)
+        # Add edges from the first L1 node to all L2 nodes
+        for i, level in levels.items():
+            if level == 2:
+                edges.append((first_l1_node, i))
+
+        return levels, edges
+
+    # Get levels and dynamic edges
+    levels, edges = get_levels_and_edges(nodes)
+    max_level = max(levels.values())
+
+    # Calculate positions
+    positions = {}
+    level_widths = defaultdict(int)
+    for node, level in levels.items():
+        level_widths[level] += 1
+
+    x_offsets = {level: - (width - 1) / 2 for level, width in level_widths.items()}
+    y_gap = 4
+
+    for node, level in levels.items():
+        positions[node] = (x_offsets[level], -level * y_gap)
+        x_offsets[level] += 1
+
+    # Create figure
     fig = go.Figure()
-    positions = {
-        0: (0, 0),
-        1: (-4, -4),
-        2: (-2, -4),
-        3: (0, -4),
-        4: (2, -4),
-        5: (4, -4),
-        6: (-4.5, -8),
-        7: (-3, -8),
-        8: (-1.5, -8),
-        9: (0, -8),
-        10: (2, -8)  # Example addition for index 10
-    }
 
+    # Add nodes to the figure
     for i, node in enumerate(wrapped_nodes):
         x, y = positions[i]
         fig.add_trace(go.Scatter(
@@ -500,6 +530,7 @@ def generate_plot(original_sentence):
             width=200
         )
 
+    # Add edges to the figure
     for edge in edges:
         x0, y0 = positions[edge[0]]
         x1, y1 = positions[edge[1]]