vendor get_head_mask, removed from transformers 5

minigpt4/models/Qformer.py

@@ -56,6 +56,20 @@ from transformers.models.bert.configuration_bert import BertConfig
 logger = logging.get_logger(__name__)
 
 
+def _get_head_mask(head_mask, num_hidden_layers):
+    """ModuleUtilsMixin.get_head_mask was removed in transformers 5.x. Head masking
+    is a training/analysis feature; at inference head_mask is always None, so the
+    contract is simply one None per layer."""
+    if head_mask is None:
+        return [None] * num_hidden_layers
+    if head_mask.dim() == 1:
+        head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+        head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
+    elif head_mask.dim() == 2:
+        head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
+    return head_mask
+
+
 class BertEmbeddings(nn.Module):
     """Construct the embeddings from word and position embeddings."""
 
@@ -940,7 +954,7 @@ class BertModel(BertPreTrainedModel):
         # attention_probs has shape bsz x n_heads x N x N
         # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
         # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
-        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+        head_mask = _get_head_mask(head_mask, self.config.num_hidden_layers)
 
         encoder_outputs = self.encoder(
             embedding_output,