Update handler.py

handler.py

@@ -1,28 +1,23 @@
 # handler.py
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 class EndpointHandler:
     def __init__(self, model_dir: str):
-        # load tokenizer & model from the same folder where handler.py lives
-        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
-        # build a HF pipeline; device_map=“auto” will pick GPU if available
+        # model_dir is already the repo root
+        tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        model     = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to("cuda")
+
+        # build a text2text pipeline on GPU (device=0)
         self.generator = pipeline(
             "text2text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
-            device=0  # set to -1 if you want CPU only
+            model=model,
+            tokenizer=tokenizer,
+            device=0
         )
 
     def __call__(self, payload: dict) -> list:
-        """
-        Expects a JSON payload like:
-          {"inputs": "<your question here>", "parameters": {"max_new_tokens": 200}}
-        Returns the raw list of dicts that HF pipeline emits.
-        """
-        text = payload.get("inputs", "")
+        # receive {"inputs": "...", "parameters": {...}}
+        text   = payload.get("inputs", "")
         params = payload.get("parameters", {})
-        # run generation
-        outputs = self.generator(text, **params)
-        return outputs
+        return self.generator(text, **params)