Create handler.py

handler.py

@@ -0,0 +1,41 @@
+from unsloth import FastLanguageModel
+import torch
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
+            model_name=path,
+            max_seq_length=2048,
+            dtype=torch.float16,
+            load_in_4bit=True,
+        )
+        FastLanguageModel.for_inference(self.model)
+
+    def __call__(self, data: dict):
+        inputs_text = data.pop("inputs", "")
+        parameters = data.pop("parameters", {})
+
+        # Formata no template do LLaMA 3
+        messages = [{"role": "user", "content": inputs_text}]
+        formatted = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+        inputs = self.tokenizer(
+            formatted,
+            return_tensors="pt"
+        ).to("cuda")
+
+        outputs = self.model.generate(
+            **inputs,
+            max_new_tokens=parameters.get("max_new_tokens", 512),
+            temperature=parameters.get("temperature", 0.7),
+            do_sample=True,
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
+
+        # Retorna só a resposta, sem o prompt
+        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return {"generated_text": decoded}