Create handler.py

handler.py

@@ -0,0 +1,43 @@
+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}