Update app.py

app.py

@@ -6,12 +6,11 @@ Automatically generated by Colab.
 Original file is located at
     https://colab.research.google.com/drive/1pFGR4uvXMMWVJFQeFmn--arumSxqa5Yy
 """
-import os
-os.system('apt install graphviz')
 
 from transformers import AutoTokenizer
 from transformers import AutoModelForSeq2SeqLM
-import plotly.graph_objects as go
+import plotly.graph_objs as go
+import textwrap
 from transformers import pipeline
 import re
 import time
@@ -436,32 +435,119 @@ def mask(sentence):
     return masked_sentences
 
 # Function to generate the tree and return the Graphviz source
-def generate_tree(original_sentence: str) -> str:
+# def generate_tree(original_sentence: str) -> str:
+#     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)
+#     dot = graphviz.Digraph()
+#     dot.attr(rankdir='LR', size='10,10', dpi=' 2743')
+
+#     existing_nodes = set()
+
+#     def add_paraphrases(parent, paraphrases):
+#         if parent not in existing_nodes:
+#             dot.node(parent, parent, shape='box')
+#             existing_nodes.add(parent)
+
+#         for paraphrase in paraphrases:
+#             if paraphrase not in existing_nodes:
+#                 dot.node(paraphrase, paraphrase, shape='box')
+#                 existing_nodes.add(paraphrase)
+#             dot.edge(parent, paraphrase)
+
+#     add_paraphrases(original_sentence, paraphrased_sentences) #whenever a new branch is to be created call this function along with the original sentence and the list of its paraphrases
+#     add_paraphrases(paraphrased_sentences[0], masked_versions)
+
+#     graph_path = dot.render(filename='paraphrase_tree_dynamic', format='png')
+#     return masked_sentence, masked_versions, dot.source
+
+
+
+#plotly tree
+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)
-    dot = graphviz.Digraph()
-    dot.attr(rankdir='LR', size='10,10', dpi=' 2743')
-
-    existing_nodes = set()
+    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)
+    ]
+    wrapped_nodes = ['<br>'.join(textwrap.wrap(node, width=30)) for node in nodes]
+    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
+    }
+
+    for i, node in enumerate(wrapped_nodes):
+        x, y = positions[i]
+        fig.add_trace(go.Scatter(
+            x=[x],
+            y=[y],
+            mode='markers',
+            marker=dict(size=10, color='blue'),
+            hoverinfo='none'
+        ))
+        fig.add_annotation(
+            x=x,
+            y=y,
+            text=node,
+            showarrow=False,
+            yshift=20,  # Adjust the y-shift value to avoid overlap
+            align="center",
+            font=dict(size=10),
+            bordercolor='black',
+            borderwidth=1,
+            borderpad=4,
+            bgcolor='white',
+            width=200
+        )
+
+    for edge in edges:
+        x0, y0 = positions[edge[0]]
+        x1, y1 = positions[edge[1]]
+        fig.add_trace(go.Scatter(
+            x=[x0, x1],
+            y=[y0, y1],
+            mode='lines',
+            line=dict(color='black', width=2)
+        ))
+
+    fig.update_layout(
+        showlegend=False,
+        margin=dict(t=50, b=50, l=50, r=50),
+        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+        width=1470,
+        height=800  # Increase height to provide more space
+    )
 
-    def add_paraphrases(parent, paraphrases):
-        if parent not in existing_nodes:
-            dot.node(parent, parent, shape='box')
-            existing_nodes.add(parent)
+    return masked_sentence, masked_versions, fig
 
-        for paraphrase in paraphrases:
-            if paraphrase not in existing_nodes:
-                dot.node(paraphrase, paraphrase, shape='box')
-                existing_nodes.add(paraphrase)
-            dot.edge(parent, paraphrase)
 
-    add_paraphrases(original_sentence, paraphrased_sentences) #whenever a new branch is to be created call this function along with the original sentence and the list of its paraphrases
-    add_paraphrases(paraphrased_sentences[0], masked_versions)
 
-    graph_path = dot.render(filename='paraphrase_tree_dynamic', format='png')
-    return masked_sentence, masked_versions, dot.source
 
 # Function for the Gradio interface
 def model(prompt):
@@ -472,11 +558,8 @@ def model(prompt):
     for i in range(len(common_subs)):
         common_subs[i]["Paraphrased Sentence"] = res[i]
     result = highlight_phrases_with_colors(res, common_grams)
-    masked_sentence, masked_versions, tree_source = generate_tree(sentence)
-    graph = graphviz.Source(tree_source)
-    png_content = graph.render(filename='paraphrase_tree_dynamic', format='png')
-    # tree = f'<div style="width: 100%; overflow-x: auto;">{svg_content}</div>'
-    return generated, generated, result, masked_sentence, masked_versions, png_content
+    masked_sentence, masked_versions, tree = generate_plot(sentence)
+    return generated, generated, result, masked_sentence, masked_versions, tree
 
 with gr.Blocks(theme = gr.themes.Monochrome()) as demo:
     gr.Markdown("# Paraphrases the Text and Highlights the Non-melting Points")
@@ -504,7 +587,7 @@ with gr.Blocks(theme = gr.themes.Monochrome()) as demo:
         masked_versions = gr.Textbox(label="Sentence Generated by Masking Model")
 
     with gr.Row():
-        tree = gr.Image(label="Paraphrase Tree")
+        tree = gr.Plot()
 
     submit_button.click(model, inputs=user_input, outputs=[ai_output, selected_sentence, html_output, masked_sentence, masked_versions, tree])
     clear_button.click(lambda: "", inputs=None, outputs=user_input)