Update handler.py

handler.py

@@ -1,39 +1,28 @@
 # handler.py
-
-import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-
-
-_generator = None
-
-def init():
-
-    global _generator
-    model_dir = "."  
-    device = 0 if torch.cuda.is_available() else -1
-
-    tokenizer = AutoTokenizer.from_pretrained(model_dir)
-    model     = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to(device if device>=0 else "cpu")
-
-    _generator = pipeline(
-        "text2text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=device
-    )
-
-def run(payload: dict) -> list:
-
-    text   = payload.get("inputs", "")
-    params = payload.get("parameters", {})
-    return _generator(text, **params)
+import os
 
 class EndpointHandler:
-
     def __init__(self, model_dir: str):
-        # simply delegate to init()
-        init()
+        # 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
+        self.generator = pipeline(
+            "text2text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device=0  # set to -1 if you want CPU only
+        )
 
     def __call__(self, payload: dict) -> list:
-        # delegate to run()
-        return run(payload)
+        """
+        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", "")
+        params = payload.get("parameters", {})
+        # run generation
+        outputs = self.generator(text, **params)
+        return outputs