from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio

# Especificar el dispositivo en el que se cargará el modelo (en este caso, "cuda" para GPU)
device = "cuda"

# Cargar el modelo preentrenado y el tokenizador asociado
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")


# Cargar el tokenizador asociado
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

def speak(prompt):
    # Tokenizar el prompt y convertirlo a tensores de PyTorch, luego enviarlos al dispositivo especificado
    model_inputs = tokenizer([prompt], return_tensors="pt").to(device)
    model.to(device)

    # Generar texto condicionalmente a partir del prompt utilizando el modelo
    generated_ids = model.generate(**model_inputs, max_new_tokens=100, do_sample=True)

    # Decodificar los identificadores generados en texto y imprimir el resultado
    resulting_text = tokenizer.batch_decode(generated_ids)[0]
    return resulting_text


iface = gr.Interface(fn=speak, inputs="text", outputs="text")
iface.launch()