Update main.py

main.py

@@ -12,8 +12,34 @@ os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
 os.environ["TOKENIZERS_PARALLELISM"]       = "false"
 os.environ["HF_HUB_VERBOSITY"]            = "error"
 
+import torch
 import numpy as np
 import faiss
+
+# ── Compatibility patch ────────────────────────────────────────────────────────
+# `find_pruneable_heads_and_indices` was removed from transformers.pytorch_utils
+# in newer releases, but the jina-bert-v2 custom modeling code still imports it
+# from there.  We re-inject the original implementation before anything loads.
+import transformers.pytorch_utils as _pt_utils
+if not hasattr(_pt_utils, "find_pruneable_heads_and_indices"):
+    from typing import Set, List, Tuple
+    def _find_pruneable_heads_and_indices(
+        heads: List[int],
+        n_heads: int,
+        head_size: int,
+        already_pruned_heads: Set[int],
+    ) -> Tuple[Set[int], torch.LongTensor]:
+        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.LongTensor = torch.arange(len(mask))[mask].long()
+        return heads, index
+    _pt_utils.find_pruneable_heads_and_indices = _find_pruneable_heads_and_indices
+# ──────────────────────────────────────────────────────────────────────────────
+
 from fastapi import FastAPI, HTTPException, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field