handler.py

import json
from transformers import AutoModelForCausalLM, AutoTokenizer

# Inicialização global do modelo e tokenizador
MODEL_NAME = "unsloth/mistral-7b-v0.3-bnb-4bit"
model = None
tokenizer = None

def initialize():
    """
    Função para carregar o modelo e o tokenizador durante a inicialização do endpoint.
    """
    global model, tokenizer
    print("Carregando modelo e tokenizador...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
    print("Modelo e tokenizador carregados com sucesso.")

def handle(data, context):
    """
    Função principal para processar requisições de entrada e gerar respostas.
    """
    global model, tokenizer
    if model is None or tokenizer is None:
        initialize()

    # Parse da entrada
    try:
        inputs = json.loads(data)
        text = inputs.get("text", "")
        max_length = inputs.get("max_length", 128)
    except Exception as e:
        return {"error": f"Erro ao processar entrada: {str(e)}"}

    # Geração de texto com o modelo
    try:
        tokens = tokenizer.encode(text, return_tensors="pt")
        output_tokens = model.generate(tokens, max_length=max_length, eos_token_id=2)
        generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
    except Exception as e:
        return {"error": f"Erro ao gerar resposta: {str(e)}"}

    return {"generated_text": generated_text}