Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import numpy as np
+import torch
 from datasets import load_dataset
 from transformers import AutoTokenizer
 import matplotlib.pyplot as plt
@@ -42,7 +43,7 @@ def analyze_sentence(index, vis_type, vis_format):
         plt.xlabel('SOURCE')    
         plt.grid()
         fig_output, graph_output = fig, ""
-    else:
+    elif vis_format == 'Graph':
         ex = [{
             "words": [{"text": x, "tag": ""} for x in tokenized[1:]],
             "arcs": [{"start": j, "end": i, "label": "", "dir": "right"} 
@@ -58,6 +59,14 @@ def analyze_sentence(index, vis_type, vis_format):
                         "</div>"
                        )
         fig_output = None
+    else:
+        fig_output = None
+        top_values, top_indices = torch.tensor(plot_data).flatten().topk(30)
+        topk_data = []
+        for val, ind in zip(top_values, top_indices):
+            ind = np.unravel_index(ind, plot_data.shape)
+            topk_data += [(tokenized[1+ind[0]], tokenized[1+ind[1]])]
+        graph_output = '<div>' + str(topk_data) + '</div>'
     metrics = {'Metrics': 1}
     
     metrics.update({k: v for k, v in row.items() if k not in ['text', 'attention_maps', 'attention_maps_shape']})    
@@ -75,7 +84,7 @@ with demo:
                                    min_width=70, value=VisType.SUM, type='value')
         btn = gr.Button("Run", min_width=30)
     
-    vis_format_checkbox = gr.Radio(['Plot', 'Graph'])
+    vis_format_checkbox = gr.Radio(['Plot', 'Graph', 'Text'])
         
     output = gr.Plot(label="Plot", container=True)
     graph_output = gr.HTML()