Spaces:

AI4PD
/

hexviz

Sleeping

aksell commited on May 9, 2023

Commit

b65afcb

1 Parent(s): 6d4dc1d

Fix T5 bug with tokens returned

Files changed (1) hide show

hexviz/attention.py CHANGED Viewed

@@ -191,7 +191,7 @@ def get_attention(
         if remove_special_tokens:
             # Remove attention to </s> (last) token
             attentions = [attention[:, :, :-1, :-1] for attention in attentions]
-            tokenized_sequence = inputs[:-1]
         attentions = torch.stack([attention.squeeze(0) for attention in attentions])
     else:

         if remove_special_tokens:
             # Remove attention to </s> (last) token
             attentions = [attention[:, :, :-1, :-1] for attention in attentions]
+            tokenized_sequence = tokenized_sequence[:-1]
         attentions = torch.stack([attention.squeeze(0) for attention in attentions])
     else: