File size: 1,463 Bytes
0ef5b19 |
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 |
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}
|