Update level_classifier_tool_2.py

level_classifier_tool_2.py

@@ -21,20 +21,41 @@ class HFEmbeddingBackend:
     MODEL: Any = field(init=False, repr=False)
 
     def __post_init__(self):
-        os.environ.setdefault("SPACES_ZERO_DISABLED", "1")
-        try:
-            torch.backends.cuda.sdp_kernel(enable_math=True, enable_flash=False, enable_mem_efficient=False)
-        except Exception:
-            pass
-
-        self.TOK = AutoTokenizer.from_pretrained(self.model_name)
-        self.MODEL = AutoModel.from_pretrained(self.model_name, attn_implementation="eager")
-        try:
-            self.MODEL.config.attn_implementation = "eager"
-        except Exception:
-            pass
-
-        self.MODEL.to(self.device).eval()
+    # 1) Try to disable Spaces ZeroGPU monkey-patch proactively
+    os.environ.setdefault("SPACES_ZERO_DISABLED", "1")
+    try:
+        # If Spaces was already imported somewhere, explicitly disable its patch.
+        from spaces import zero as _spaces_zero  # safe import; no-op if not installed
+        if hasattr(_spaces_zero, "disable"):
+            _spaces_zero.disable()
+    except Exception:
+        pass
+
+    # 2) Keep attention off Flash/MemEfficient (avoid vectorized mask paths)
+    try:
+        torch.backends.cuda.sdp_kernel(enable_math=True, enable_flash=False, enable_mem_efficient=False)
+    except Exception:
+        pass
+
+    # 3) Load tokenizer/model and force eager attention (non-vmap route)
+    self.TOK = AutoTokenizer.from_pretrained(self.model_name)
+    self.MODEL = AutoModel.from_pretrained(self.model_name, attn_implementation="eager")
+
+    # (extra safety) disable any sliding/windowed attention that can trigger the vmap mask path
+    try:
+        if hasattr(self.MODEL.config, "sliding_window"):
+            self.MODEL.config.sliding_window = None
+        if hasattr(self.MODEL, "generation_config") and hasattr(self.MODEL.generation_config, "sliding_window"):
+            self.MODEL.generation_config.sliding_window = None
+    except Exception:
+        pass
+
+    try:
+        self.MODEL.config.attn_implementation = "eager"
+    except Exception:
+        pass
+
+    self.MODEL.to(self.device).eval()
 
     def encode(self, texts: Iterable[str], batch_size: int = 32) -> "Tuple[torch.Tensor, List[str]]":
         """