Update app.py

app.py

@@ -9,8 +9,8 @@ dataset = load_dataset('dar-tau/grammar-attention-maps-opt-350m')['train']
 def analyze_sentence(index):
     row = dataset[index]
     text = row['text']
-    _, seq_len, _ = row['attention_maps_shape']
-    
+    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(1).sum(0))
     plt.xticks(np.arange(seq_len - 1) + 0.5,