Create handler.py

handler.py

@@ -0,0 +1,47 @@
+from unsloth import FastLanguageModel
+from transformers import AutoTokenizer
+import torch
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True)
+        self.model, _ = FastLanguageModel.from_pretrained(
+            model_name=path,
+            max_seq_length=2048,
+            dtype=torch.float16,
+            load_in_4bit=True,
+        )
+        self.model.to(self.device)
+        self.model.eval()
+
+    def __call__(self, data):
+        prompt = data.get("inputs", "")
+        if not prompt:
+            return {"error": "Missing 'inputs' in request payload."}
+
+        generation_params = {
+            "max_new_tokens": data.get("max_new_tokens", 128),
+            "temperature": data.get("temperature", 0.7),
+            "top_p": data.get("top_p", 0.9),
+            "top_k": data.get("top_k", 50),
+            "do_sample": data.get("do_sample", True),
+            "repetition_penalty": data.get("repetition_penalty", 1.1),
+        }
+
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                **generation_params
+            )
+
+        generated_text = self.tokenizer.decode(
+            outputs[0],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True
+        )
+
+        return {"generated_text": generated_text}