Update handler.py

handler.py

@@ -7,6 +7,13 @@ class EndpointHandler:
         # We point directly to the original weights
         self.model_id = "zhihan1996/DNABERT-2-117M"
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
+
+        # --- FIX: Disable Flash Attention to avoid the Triton error ---
+        from transformers import AutoConfig
+        config = AutoConfig.from_pretrained(self.model_id, trust_remote_code=True)
+        config.use_flash_attn = False  # This bypasses the broken line 114 code
+        # --------------------------------------------------------------
+        
         self.model = AutoModel.from_pretrained(self.model_id, trust_remote_code=True)
         if torch.cuda.is_available():
             self.model = self.model.to("cuda")