update tokenizer

handler.py

@@ -1,16 +1,30 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
 from typing import Dict
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 class EndpointHandler:
-    def __init__(self, path=""):
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(path)
+    """
+    Minimal custom handler for InternLM2 / NuExtract-2-8B
+    """
+
+    def __init__(self, path: str = "./model"):
+        # allow execution of custom model code
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            path, trust_remote_code=True
+        )
+        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
 
     def __call__(self, data: Dict[str, str]) -> Dict[str, str]:
-        inputs = data.get("inputs", "")
-        if not inputs:
+        prompt = data.get("inputs", "")
+        if not prompt:
             return {"error": "No input provided."}
-        inputs = self.tokenizer(inputs, return_tensors="pt")
-        outputs = self.model.generate(**inputs, max_new_tokens=100)
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return {"generated_text": response}
+
+        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)
+        return {"generated_text": answer}