Merge pull request #6 from nitinsurya/main

GLiNER#263 Fix get_attention_mask function that allows for batched in…

Author: Ingvarstep

src/flashdeberta/model.py

@@ -301,6 +301,12 @@ def __init__(self, config):
         self.gradient_checkpointing = False
 
     def get_attention_mask(self, attention_mask):
+        if attention_mask.dim() <= 2:
+            extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
+        elif attention_mask.dim() == 3:
+            attention_mask = attention_mask.unsqueeze(1)
+
         return attention_mask
 
 class FlashDebertaV2PreTrainedModel(PreTrainedModel):