Create handler.py

handler.py

@@ -0,0 +1,79 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from peft import PeftModel
+import torch
+import os
+import traceback
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        base_model_id = "microsoft/Phi-4-mini-instruct"
+        adapter_path = path # Diretório local no container onde o repo foi baixado
+
+        try:
+            print(f"Iniciando Handler: Carregando modelo base {base_model_id}")
+            # Carregar em bfloat16 ou float16 se disponível
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            self.base_model = AutoModelForCausalLM.from_pretrained(
+                base_model_id,
+                torch_dtype=dtype,
+                trust_remote_code=True
+                # device_map é gerenciado pelo endpoint
+            )
+
+            print(f"Carregando tokenizer de {base_model_id}")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                base_model_id,
+                trust_remote_code=True
+            )
+            if self.tokenizer.pad_token is None or self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
+                self.tokenizer.pad_token = self.tokenizer.unk_token
+                print("Definido tokenizer.pad_token = tokenizer.unk_token")
+
+            print(f"Carregando adaptador LoRA de {adapter_path}")
+            self.model = PeftModel.from_pretrained(self.base_model, adapter_path)
+            self.model.eval()
+            print("Adaptador LoRA carregado.")
+
+            self.pipeline = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                # device=0 # Geralmente não necessário, endpoint gerencia
+            )
+            print("Pipeline de text-generation criado. Handler pronto.")
+
+        except Exception as e:
+            print(f"ERRO FATAL durante __init__ do Handler: {e}")
+            print(traceback.format_exc())
+            raise e # Levanta o erro para falhar a inicialização do endpoint
+
+
+    def __call__(self, data):
+        try:
+            inputs = data.pop("inputs", data)
+            parameters = data.pop("parameters", None) or {}
+
+            print(f"Handler __call__ recebeu inputs: {inputs}")
+            print(f"Handler __call__ recebeu parâmetros: {parameters}")
+
+            # Preparar o prompt - Detecta se input é lista de dicts (chat) ou string
+            prompt_text = inputs
+            if isinstance(inputs, list) and len(inputs) > 0 and isinstance(inputs[0], dict) and 'role' in inputs[0]:
+                 print("Aplicando chat template...")
+                 # Cuidado: add_generation_prompt=True é para gerar RESPOSTA do assistant
+                 prompt_text = self.tokenizer.apply_chat_template(inputs, tokenize=False, add_generation_prompt=True)
+
+            print(f"Texto do prompt para o pipeline: {prompt_text}")
+
+            # Gerar texto usando o pipeline
+            outputs = self.pipeline(prompt_text, **parameters)
+
+            print(f"Handler __call__ gerou outputs: {outputs}")
+            # Retorna a saída (geralmente uma lista de dicionários)
+            return outputs
+
+        except Exception as e:
+            print(f"ERRO durante __call__ do Handler: {e}")
+            print(traceback.format_exc())
+            # Retornar erro de forma estruturada ajuda na depuração
+            return [{"error": str(e), "traceback": traceback.format_exc()}]