Update main.py

main.py

@@ -17,23 +17,20 @@ import numpy as np
 import faiss
 
 # ── Compatibility patches ──────────────────────────────────────────────────────
-# The jina-bert-v2 custom modeling code was written against an older transformers
-# API. Two things were removed / tightened in newer releases:
+# jina-bert-v2 (trust_remote_code) was written against transformers 4.x.
+# Transformers 5.x removed / broke three things the model relies on.
+# All patches are no-ops when the symbol already exists.
 #
-# 1. `find_pruneable_heads_and_indices` was removed from transformers.pytorch_utils.
-# 2. `PretrainedConfig` no longer sets is_decoder / add_cross_attention as instance
-#    defaults in __init__. A tightened __getattribute__ now raises AttributeError
-#    instead of the old silent fallback, breaking JinaBertConfig access patterns.
-#
-# Both patches are guarded with hasattr/flag checks so they are no-ops if a future
-# transformers version re-adds these symbols.
+# 1. find_pruneable_heads_and_indices  — removed from pytorch_utils
+# 2. PretrainedConfig.is_decoder etc  — no longer set as instance defaults
+# 3. PreTrainedModel.get_head_mask     — removed from modeling_utils in T5
 
+# ── patch 1: pytorch_utils ────────────────────────────────────────────────────
 import transformers.pytorch_utils as _pt_utils
 if not hasattr(_pt_utils, "find_pruneable_heads_and_indices"):
     def _find_pruneable_heads_and_indices(
         heads, n_heads: int, head_size: int, already_pruned_heads
     ):
-        import torch
         mask = torch.ones(n_heads, head_size)
         heads = set(heads) - already_pruned_heads
         for head in heads:
@@ -44,25 +41,51 @@ if not hasattr(_pt_utils, "find_pruneable_heads_and_indices"):
         return heads, index
     _pt_utils.find_pruneable_heads_and_indices = _find_pruneable_heads_and_indices
 
+# ── patch 2: PretrainedConfig legacy defaults ─────────────────────────────────
 import transformers.configuration_utils as _cfg_utils
 _PC = _cfg_utils.PretrainedConfig
 if not hasattr(_PC, "_jina_compat_patched"):
-    # Attributes that used to be set in PretrainedConfig.__init__ with defaults
-    # but were removed from the base class in newer transformers versions.
-    _LEGACY_DEFAULTS = {
-        "is_decoder":               False,
-        "add_cross_attention":      False,
+    _LEGACY_CFG_DEFAULTS = {
+        "is_decoder":                  False,
+        "add_cross_attention":         False,
         "cross_attention_hidden_size": None,
-        "use_cache":                True,
+        "use_cache":                   True,
     }
     def _pc_getattr(self, key: str):
-        if key in _LEGACY_DEFAULTS:
-            return _LEGACY_DEFAULTS[key]
+        if key in _LEGACY_CFG_DEFAULTS:
+            return _LEGACY_CFG_DEFAULTS[key]
         raise AttributeError(
             f"'{type(self).__name__}' object has no attribute '{key}'"
         )
     _PC.__getattr__ = _pc_getattr
     _PC._jina_compat_patched = True
+
+# ── patch 3: PreTrainedModel.get_head_mask ────────────────────────────────────
+import transformers.modeling_utils as _mod_utils
+_PTM = _mod_utils.PreTrainedModel
+if not hasattr(_PTM, "get_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 _get_head_mask(self, head_mask, num_hidden_layers, is_attention_chunked=False):
+        if head_mask is not None:
+            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
+            if is_attention_chunked:
+                head_mask = head_mask.unsqueeze(-1)
+        else:
+            head_mask = [None] * num_hidden_layers
+        return head_mask
+
+    if not hasattr(_PTM, "_convert_head_mask_to_5d"):
+        _PTM._convert_head_mask_to_5d = _convert_head_mask_to_5d
+    _PTM.get_head_mask = _get_head_mask
 # ──────────────────────────────────────────────────────────────────────────────
 
 from fastapi import FastAPI, HTTPException, UploadFile, File, Form