Rename models/med.py to models/bert.py

models/{med.py → bert.py}

@@ -224,6 +224,12 @@ class BertSelfAttention(nn.Module):
 
         past_key_value = (key_layer, value_layer)
 
+        # compatible with higher versions of transformers 
+        if key_layer.shape[0] > query_layer.shape[0]:
+            key_layer = key_layer[:query_layer.shape[0], :, :, :]
+            attention_mask = attention_mask[:query_layer.shape[0], :, :]
+            value_layer = value_layer[:query_layer.shape[0], :, :, :]
+
         # Take the dot product between "query" and "key" to get the raw attention scores.
         attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
 
@@ -392,12 +398,10 @@ class BertLayer(nn.Module):
         mode=None,
     ):
         
-        if mode == 'mlr':
-
+        if mode == 'tagging':
+            
             assert encoder_hidden_states is not None, "encoder_hidden_states must be given for cross-attention layers"
 
-            # print('attention_output.shape',attention_output.shape)
-            # print('encoder_hidden_states.shape',encoder_hidden_states.shape)
             cross_attention_outputs = self.crossattention(
                 hidden_states,
                 attention_mask,