File size: 3,449 Bytes
0d6bfde |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 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()}] |