Fix compatibility with transformers >= 5

ImageReward/models/BLIP/blip.py

@@ -17,7 +17,7 @@ def init_tokenizer():
     tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
     tokenizer.add_special_tokens({'bos_token':'[DEC]'})
     tokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})       
-    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]  
+    tokenizer.enc_token_id = tokenizer.convert_tokens_to_ids('[ENC]')
     return tokenizer
 
 

ImageReward/models/BLIP/med.py

@@ -40,9 +40,23 @@
     # Compatibility for later versions of transformers
     from transformers.modeling_utils import PreTrainedModel
     from transformers.pytorch_utils import apply_chunking_to_forward
-    from transformers.pytorch_utils import find_pruneable_heads_and_indices
     from transformers.pytorch_utils import prune_linear_layer
 
+    try:
+        # transformers < 5 exposed this in pytorch_utils.
+        from transformers.pytorch_utils import find_pruneable_heads_and_indices
+    except ImportError:
+        # transformers >= 5 dropped it — vendor the original implementation.
+        def find_pruneable_heads_and_indices(heads, n_heads, head_size, already_pruned_heads):
+            mask = torch.ones(n_heads, head_size)
+            heads = set(heads) - already_pruned_heads
+            for head in heads:
+                head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
+                mask[head] = 0
+            mask = mask.view(-1).contiguous().eq(1)
+            index = torch.arange(len(mask))[mask].long()
+            return heads, index
+
 from transformers.utils import logging
 from transformers.models.bert.configuration_bert import BertConfig
 
@@ -556,6 +570,32 @@ class BertPreTrainedModel(PreTrainedModel):
     base_model_prefix = "bert"
     _keys_to_ignore_on_load_missing = [r"position_ids"]
 
+    def tie_weights(self, *args, **kwargs):
+        # No output head is tied to the input embeddings in this BLIP BERT, and transformers >= 5
+        # reworked tie_weights (all_tied_weights_keys) in a way this vendored model doesn't populate.
+        # Tying is a no-op for inference here, so skip it to stay compatible across versions.
+        return
+
+    def get_head_mask(self, head_mask, num_hidden_layers, is_attention_chunked=False):
+        # transformers >= 5 dropped get_head_mask from ModuleUtilsMixin; vendor the original.
+        if head_mask is not None:
+            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
+            if is_attention_chunked is True:
+                head_mask = head_mask.unsqueeze(-1)
+        else:
+            head_mask = [None] * num_hidden_layers
+        return head_mask
+
+    def _convert_head_mask_to_5d(self, head_mask, 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)
+        assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
+        head_mask = head_mask.to(dtype=self.dtype)
+        return head_mask
+
     def _init_weights(self, module):
         """ Initialize the weights """
         if isinstance(module, (nn.Linear, nn.Embedding)):