Update app.py

app.py

@@ -5,9 +5,11 @@ from transformers import AutoTokenizer
 import matplotlib.pyplot as plt
 import seaborn as sns
 
+
 dataset = load_dataset('dar-tau/grammar-attention-maps-opt-350m')['train']
 tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m', add_prefix_space=True)
 
+
 def analyze_sentence(index):
     row = dataset[index]
     text = row['text']
@@ -15,13 +17,14 @@ def analyze_sentence(index):
     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)
-    plot = sns.heatmap(attn_maps.sum(0)[1:, 1:])
+    fig = plt.figure()
+    sns.heatmap(attn_maps.sum(0)[1:, 1:])
     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()
-    return text, plot
+    return text, fig
 
 
 iface = gr.Interface(fn=analyze_sentence, inputs=[gr.Dropdown(choices=dataset['text'], value=0, type='index')],