trust_remote_code=True, on tokenizer

handler.py

@@ -1,23 +1,22 @@
 from typing import Dict
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class EndpointHandler:
-    """
-    Minimal custom handler for InternLM2 / NuExtract-2-8B
-    """
+    """Custom handler for NuExtract-2-8B (InternLM2 based)."""
 
-    def __init__(self, path: str = "./model"):
-        # allow execution of custom model code
+    def __init__(self, path: str = "") -> None:
+        # ↓↓↓ allow the repo’s custom configuration & modelling code
         self.tokenizer = AutoTokenizer.from_pretrained(
-            path, trust_remote_code=True
+            path,
+            trust_remote_code=True            # ← mandatory
         )
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
-            trust_remote_code=True,           # ← key line
-            torch_dtype=torch.float16,        # load in fp16 to fit on one A10/T4
-            device_map="auto"                 # send to GPU if available
-        ).eval()                              # put in inference mode
+            trust_remote_code=True,            # ← mandatory
+            torch_dtype=torch.float16,         # fits on a 16 GB GPU
+            device_map="auto"                  # put tensors on the GPU
+        ).eval()
 
     def __call__(self, data: Dict[str, str]) -> Dict[str, str]:
         prompt = data.get("inputs", "")
@@ -25,6 +24,6 @@ class EndpointHandler:
             return {"error": "No input provided."}
 
         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        outputs = self.model.generate(**inputs, max_new_tokens=128)
-        answer  = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        output_ids = self.model.generate(**inputs, max_new_tokens=128)
+        answer = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
         return {"generated_text": answer}