Upload handler.py

handler.py

@@ -0,0 +1,13 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+class EndpointHandler:
+    def __init__(self, model_path: str, task="text-generation"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
+        self.pipe = pipeline(task=task, model=self.model, tokenizer=self.tokenizer)
+
+    def __call__(self, inputs: dict) -> dict:
+        prompt = inputs.get("inputs", "")
+        params = inputs.get("parameters", {})
+        outputs = self.pipe(prompt, **params)
+        return {"generated_text": outputs[0]["generated_text"]}