Update app.py

app.py

@@ -5,9 +5,12 @@ import seaborn as sns
 
 dataset = load_dataset('dar-tau/grammar-attention-maps-opt-350m')['train']
 
-def choose_sentence(name):
-    return "Hello " + name + "!!"
+def analyze_sentence(index):
+    row = dataset[index]
+    attn_maps = np.array(row['attention_maps']).reshape(*row['attention_maps_shape'])
+    return row['text'], sns.heatmap(attn_maps.sum(1).sum(0))
 
 
-iface = gr.Interface(fn=choose_sentence, inputs=[gr.Dropdown(choices=dataset['text'])], outputs="text")
+iface = gr.Interface(fn=analyze_sentence, inputs=[gr.Dropdown(choices=dataset['text'], type='index')], 
+                     outputs=[gr.Label(), gr.Plot(label="Plot")])
 iface.launch()