Update handler.py

handler.py

@@ -6,19 +6,27 @@ class EndpointHandler:
     def __init__(self, path=""):
         self.model_id = "zhihan1996/DNABERT-2-117M"
         
-        # 1. Use a standard BERT config instead of the custom DNABERT one
-        # This prevents the 'flash_attn_triton.py' from ever being triggered
-        config = BertConfig.from_pretrained(self.model_id)
+        # 1. MANUALLY build the configuration. 
+        # By NOT using .from_pretrained(), we avoid the custom modeling code entirely.
+        config = BertConfig(
+            vocab_size=32000, 
+            hidden_size=768, 
+            num_hidden_layers=12, 
+            num_attention_heads=12, 
+            intermediate_size=3072,
+            max_position_embeddings=512,
+            type_vocab_size=2,
+            is_decoder=False
+        )
         
-        # 2. Load the tokenizer normally
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
         
-        # 3. Load as a standard BertModel
-        # We use 'trust_remote_code=False' here to force standard layers
+        # 2. Load the weights into a standard, safe BertModel
+        # We use the original weights but our 'clean' config
         self.model = BertModel.from_pretrained(
             self.model_id, 
             config=config, 
-            trust_remote_code=False, 
+            trust_remote_code=False,
             ignore_mismatched_sizes=True
         )
         
@@ -27,12 +35,9 @@ class EndpointHandler:
         self.model.eval()
 
     def __call__(self, data: Dict[str, Any]) -> List[float]:
-        # Handle input strings or dictionaries
         inputs = data.pop("inputs", data)
-        if isinstance(inputs, list):
-            inputs = inputs[0]
-
-        # Standard tokenization
+        
+        # Tokenization
         encoded_input = self.tokenizer(
             inputs, 
             return_tensors='pt', 
@@ -45,10 +50,10 @@ class EndpointHandler:
             encoded_input = {k: v.to("cuda") for k, v in encoded_input.items()}
             
         with torch.no_grad():
+            # This now runs standard PyTorch Attention layers
             outputs = self.model(**encoded_input)
             
-        # Get the hidden states and perform mean pooling
-        # index 0 is the last_hidden_state
+        # Sequence embedding via mean pooling
         embeddings = outputs[0][0].mean(dim=0).cpu().numpy().tolist()
         
         return embeddings