Update handler.py

handler.py

@@ -1,31 +1,51 @@
 import sys
-from typing import Any, Dict, List
+from unittest.mock import MagicMock
+
+# 1. GLOBAL BLACKOUT: Must be at the very top, before any other imports
+# This makes Triton invisible to every script the model downloads.
+sys.modules["triton"] = MagicMock()
+sys.modules["triton.language"] = MagicMock()
+
 import torch
+from typing import Any, Dict, List
 from transformers import AutoTokenizer, AutoModel, AutoConfig
 
 class EndpointHandler:
     def __init__(self, path=""):
         self.model_id = "zhihan1996/DNABERT-2-117M"
         
-        # 1. Load tokenizer and config
+        # 2. Config level: Explicitly set flash_attn to False in the config object
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
         config = AutoConfig.from_pretrained(self.model_id, trust_remote_code=True)
         
-        # 2. Force the config to disable flash attention
+        # Some custom implementations check for 'use_flash_attn' or 'flash_attn'
         config.use_flash_attn = False
+        if hasattr(config, "auto_map"):
+            # Force it to use the standard modeling rather than the Triton-based one
+            config.auto_map["AutoModel"] = "modeling_bert.BertModel"
+
+        # 3. Load Model
+        self.model = AutoModel.from_pretrained(
+            self.model_id, 
+            trust_remote_code=True, 
+            config=config
+        )
         
-        # 3. Load the model
-        self.model = AutoModel.from_pretrained(self.model_id, trust_remote_code=True, config=config)
-        
+        # 4. Layer Level: Double-check the individual attention layers
+        # This is our last-resort safety net
+        for module in self.model.modules():
+            if hasattr(module, "use_flash_attn"):
+                module.use_flash_attn = False
+
         if torch.cuda.is_available():
             self.model = self.model.to("cuda")
         self.model.eval()
 
     def __call__(self, data: Dict[str, Any]) -> List[float]:
-        # Extract inputs
         inputs = data.pop("inputs", data)
-        
-        # Tokenize
+        if isinstance(inputs, list):
+            inputs = inputs[0]
+
         encoded_input = self.tokenizer(
             inputs, 
             return_tensors='pt', 
@@ -36,20 +56,10 @@ class EndpointHandler:
         
         if torch.cuda.is_available():
             encoded_input = {k: v.to("cuda") for k, v in encoded_input.items()}
-
-        # THE "TOTAL ECLIPSE": 
-        # We temporarily move triton out of sys.modules so the model 
-        # thinks it's not installed and falls back to standard PyTorch math.
-        real_triton = sys.modules.pop("triton", None)
-        
-        try:
-            with torch.no_grad():
-                outputs = self.model(**encoded_input)
-                # Mean pooling for embedding
-                embeddings = outputs[0][0].mean(dim=0).cpu().numpy().tolist()
-        finally:
-            # Restore triton so we don't break the rest of the HF environment
-            if real_triton:
-                sys.modules["triton"] = real_triton
-        
+            
+        with torch.no_grad():
+            outputs = self.model(**encoded_input)
+            
+        # Mean pooling
+        embeddings = outputs[0][0].mean(dim=0).cpu().numpy().tolist()
         return embeddings