Update handler.py

handler.py

@@ -1,18 +1,34 @@
 import torch
+import os
 from typing import Any, Dict, List
+
+# Force the environment variable inside the script as well
+os.environ["HF_HUB_TRUST_REMOTE_CODE"] = "True"
+
 from transformers import AutoConfig, AutoTokenizer, AutoModelForMaskedLM
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # We explicitly load the config first to satisfy the toolkit's hunger for info
-        self.config = AutoConfig.from_pretrained(path, trust_remote_code=True)
+        # We point to the specific InstaDeep model directly to avoid
+        # any local repository naming conflicts during the 'path' resolution
+        self.model_id = "InstaDeepAI/nucleotide-transformer-v2-50m-multi-species"
         
-        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+        # 1. Load Config first with explicit trust
+        self.config = AutoConfig.from_pretrained(
+            self.model_id, 
+            trust_remote_code=True
+        )
+        
+        # 2. Load Tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.model_id, 
+            trust_remote_code=True
+        )
         
-        # We load the model using the config we just initialized
+        # 3. Load Model
         self.model = AutoModelForMaskedLM.from_pretrained(
-            path, 
-            config=self.config, 
+            self.model_id,
+            config=self.config,
             trust_remote_code=True
         )
         
@@ -21,12 +37,12 @@ class EndpointHandler:
         self.model.eval()
 
     def __call__(self, data: Dict[str, Any]) -> List[float]:
-        # Handle input safely
-        inputs = data.get("inputs", data)
+        # Handle inputs from the toolkit JSON
+        inputs = data.pop("inputs", data)
         if isinstance(inputs, list):
             inputs = inputs[0]
 
-        # Use 1000bp chunks for the 12.2kb APRIL promoter
+        # APRIL Promoter Chunking (12.2kb)
         chunk_size = 1000
         stride = 500
         chunks = [inputs[i:i + chunk_size] for i in range(0, len(inputs), stride)]
@@ -34,15 +50,21 @@ class EndpointHandler:
         all_embeddings = []
         with torch.no_grad():
             for chunk in chunks:
-                tokens = self.tokenizer(chunk, return_tensors='pt', padding=True, truncation=True, max_length=chunk_size)
+                tokens = self.tokenizer(
+                    chunk, 
+                    return_tensors='pt', 
+                    padding=True, 
+                    truncation=True, 
+                    max_length=chunk_size
+                )
                 if torch.cuda.is_available():
                     tokens = {k: v.to("cuda") for k, v in tokens.items()}
                 
                 outputs = self.model(**tokens, output_hidden_states=True)
-                # Mean pool last hidden state
+                # Last hidden state mean pooling
                 chunk_emb = torch.mean(outputs.hidden_states[-1], dim=1).squeeze()
                 all_embeddings.append(chunk_emb)
 
-        # Average the chunks into one vector for LangGraph
+        # Average the chunks for one representative vector
         final_embedding = torch.stack(all_embeddings).mean(dim=0).cpu().numpy().tolist()
         return final_embedding