Add an assert about the attention_mask dtype.

The attention mask in AMPLIFY should be additive, as a consequence, masked positions should have -inf values and unmasked positions should have 0 (or any constant value, but 0 is standard).

amplify.py

@@ -246,6 +246,10 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
 
         # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
         if attention_mask is not None and not torch.all(attention_mask == 0):
+            assert attention_mask.dtype != torch.bool and 1.0 not in attention_mask, (
+                "AMPLIFY expects an additive attention_mask.\n"
+                "Modify the output of the tokenizer with attention_mask = torch.where(attention_mask, float(0.0), float("-inf"))"
+            )
             attention_mask = (
                 attention_mask.unsqueeze(1)
                 .unsqueeze(1)