Update app.py

app.py

@@ -13,11 +13,11 @@ 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)
+    tokenized = tokenizer.batch_decode(tokenizer.encode(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)
-    fig = plt.figure()
+    seq_len = attn_map_shape[1]
+    attn_maps = np.array(row['attention_maps']).reshape(*attn_map_shape).clip(0, 1)
+    fig = plt.figure(figsize=(8, 8))
     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);