jinaai
/

jina-bert-implementation

Model card Files Files and versions

Positional Interpolation

#14

by Markus28 - opened Mar 6, 2024

base: refs/heads/main

←

from: refs/pr/14

Discussion Files changed

This PR is in draft mode

Files changed (1) hide show

modeling_bert.py +2 -1

modeling_bert.py CHANGED Viewed

@@ -787,7 +787,8 @@ class JinaBertEncoder(nn.Module):
             # Device catch-up
             self.alibi = self.alibi.to(hidden_states.device)
-        alibi_bias = self.alibi[:, :, :seqlen, :seqlen]
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 logger.warning_once(

             # Device catch-up
             self.alibi = self.alibi.to(hidden_states.device)
+        unpadded_seqlens = torch.sum(attention_mask, dim=1).unsqueeze(1).unsqueeze(1).unsqueeze(1)
+        alibi_bias = self.alibi[:, :, :seqlen, :seqlen] * 512 / unpadded_seqlens
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 logger.warning_once(