Create handler.py

handler.py

@@ -0,0 +1,29 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16, device_map="auto")
+        self.pipe = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.95,
+        )
+
+    def __call__(self, data):
+        prompt = data.get("inputs", "")
+        if not prompt:
+            return {"error": "No input provided"}
+
+        # Generate response
+        response = self.pipe(prompt)[0]['generated_text']
+
+        # Remove the original prompt from the response
+        response = response[len(prompt):].strip()
+
+        return {"generated_text": response}