Update app.py

app.py

@@ -24,7 +24,6 @@ def analyze_sentence(index, vis_type, vis_format):
     attn_map_shape = row['attention_maps_shape'][1:]
     seq_len = attn_map_shape[1]
     attn_maps = np.array(row['attention_maps']).reshape(*attn_map_shape).clip(0, 1)
-    fig = plt.figure(figsize=(0.1 + 0.3 * len(tokenized), 0.25 * len(tokenized)))
     attn_maps = attn_maps[:, 1:, 1:]
     if vis_type == VisType.SUM.value:
         plot_data = attn_maps.sum(0)
@@ -34,15 +33,26 @@ def analyze_sentence(index, vis_type, vis_format):
     else:
         print(vis_type)
         0/0
-    sns.heatmap(plot_data)
-    plt.xticks(np.arange(seq_len - 1) + 0.5, tokenized[1:], rotation=90);
-    plt.yticks(np.arange(seq_len - 1) + 0.5, tokenized[1:], rotation=0);
-    plt.ylabel('TARGET')
-    plt.xlabel('SOURCE')    
-    plt.grid()
+    if vis_format == 'Plot':
+        fig = plt.figure(figsize=(0.1 + 0.3 * len(tokenized), 0.25 * len(tokenized)))
+        sns.heatmap(plot_data)
+        plt.xticks(np.arange(seq_len - 1) + 0.5, tokenized[1:], rotation=90);
+        plt.yticks(np.arange(seq_len - 1) + 0.5, tokenized[1:], rotation=0);
+        plt.ylabel('TARGET')
+        plt.xlabel('SOURCE')    
+        plt.grid()
+        fig_output, graph_output = fig, ""
+    else:
+        ex = [{
+            "words": [{"text": x} for x in tokenized],
+            "arcs": [{"start": i, "end": j} for j in range(i) else range(len(tokenized))]
+               ]
+        graph_output = displacy.render(ex, style="dep", manual=True)
+        fig_output = None
     metrics = {'Metrics': 1}
+    
     metrics.update({k: v for k, v in row.items() if k not in ['text', 'attention_maps', 'attention_maps_shape']})    
-    return fig, metrics
+    return fig_output, graph_output, metrics
 
 demo = gr.Blocks()
 with demo:
@@ -57,8 +67,9 @@ with demo:
     vis_format_checkbox = gr.Radio(['Plot', 'Graph'])
         
     output = gr.Plot(label="Plot", container=True)
+    graph_output = gr.HTML(label="Graph")
     metrics = gr.Label("Metrics")
-    btn.click(analyze_sentence, [sentence_dropdown, vis_dropdown, vis_format_checkbox], [output, metrics])
+    btn.click(analyze_sentence, [sentence_dropdown, vis_dropdown, vis_format_checkbox], [output, graph_output, metrics])
 
 
 if __name__ == "__main__":