Update app.py

app.py

@@ -11,14 +11,13 @@ tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m', add_prefix_space=
 def analyze_sentence(index):
     row = dataset[index]
     text = row['text']
+    tokenized = tokenizer.tokenize(text, add_special_tokens=False)
     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:])
-    plt.xticks(np.arange(seq_len - 1) + 0.5,
-           tokenizer.tokenize(text, add_special_tokens=False), rotation=90);
-    plt.yticks(np.arange(seq_len - 1) + 0.5,
-               tokenizer.tokenize(text, add_special_tokens=False), rotation=0);
+    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()