Update app.py

app.py

@@ -6,10 +6,12 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 from enum import Enum
 
-class VisType(Enum):
-    ALL = 'ALL'
 
+class VisType(Enum):
+    SUM = 'Sum over Layers'
+    
 
+num_layers = 24
 dataset = load_dataset('dar-tau/grammar-attention-maps-opt-350m')['train']
 tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m', add_prefix_space=True)
 
@@ -21,10 +23,13 @@ def analyze_sentence(index, vis_type):
     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).clip(0, 1)
-    fig = plt.figure(figsize=(0.5 + 0.5 * len(tokenized), 0.4 * len(tokenized)))
+    fig = plt.figure(figsize=(0.5 + 0.4 * len(tokenized), 0.35 * len(tokenized)))
     attn_maps = attn_maps[:, 1:, 1:]
-    if vis_type == VisType.ALL.value:
+    if vis_type == VisType.SUM.value:
         plot_data = attn_maps.sum(0)
+    elif vis_type.startswith('Layer #'):
+        layer_to_inspect = int(vis_type.split('#')[1])
+        plot_data = attn_maps[layer_to_inspect]
     else:
         print(vis_type)
         0/0
@@ -34,7 +39,8 @@ def analyze_sentence(index, vis_type):
     plt.ylabel('TARGET')
     plt.xlabel('SOURCE')    
     plt.grid()
-    metrics = {k: v for k, v in row.items() if k not in ['text', 'attention_maps', 'attention_maps_shape']}
+    metrics = {'Metrics': 0}
+    metrics.update({k: v for k, v in row.items() if k not in ['text', 'attention_maps', 'attention_maps_shape']})    
     return fig, metrics
 
 demo = gr.Blocks()
@@ -43,8 +49,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=[x.value for x in VisType], 
-                                   min_width=150, value=VisType.ALL, type='value')
+        vis_dropdown = gr.Dropdown(label="Visualization", choices=[x.value for x in VisType] + ['Layer #' + i for i in range(num_layers)], 
+                                   min_width=150, value=VisType.SUM, type='value')
         btn = gr.Button("Run", min_width=30)
     output = gr.Plot(label="Plot", container=True)
     metrics = gr.Label("Metrics")