Update app.py

app.py

@@ -22,7 +22,7 @@ def analyze_sentence(index, vis_type):
     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.5 + 0.5 * len(tokenized), 0.4 * len(tokenized)))
-    plot_data = attn_maps[:, 1:, 1:]
+    attn_maps = attn_maps[:, 1:, 1:]
     if vis_type == VisType.ALL:
         plot_data = attn_maps.sum(0)
     sns.heatmap(plot_data)
@@ -40,7 +40,8 @@ with demo:
         sentence_dropdown = gr.Dropdown(label="Sentence", 
                                         choices=[x.split('</s> ')[1] for x in dataset['text']], 
                                value=0, min_width=500, type='index')
-        vis_dropdown = gr.Dropdown(label="Visualization", choices=list(VisType), value=0, type='index')
+        vis_dropdown = gr.Dropdown(label="Visualization", choices=list(VisType), 
+                                   min_width=80, value=0, type='index')
         btn = gr.Button("Run")
     output = gr.Plot(label="Plot", container=True)
     metrics = gr.Label("Metrics")