Prueba_Mistral / app.py
Ernestasio's picture
Create app.py
e919c00
raw
history blame
1.06 kB
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()