Update app.py

app.py

@@ -26,142 +26,131 @@ from huggingface_hub import login
 HF_TOKEN = os.getenv('HF_TOKEN')
 login(token=HF_TOKEN)
 
-# Available models
+# Modelos disponibles
 AVAILABLE_MODELS = {
-    "bloomz-560m": "bigscience/bloomz-560m",
-    "bloomz-7B1": "bigscience/bloomz-7b1",
-    "GPT-J-6B": "EleutherAI/gpt-j-6b",
-    "mT5-XL": "google/mt5-xl", 
+    "BLOOMZ-560M": "bigscience/bloomz-560m"
 }
 
-# Initialize model and tokenizer
+# Inicializar modelo y tokenizer
 current_model = None
 current_tokenizer = None
 current_model_name = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def load_model(model_name):
-    """Load the selected model and tokenizer."""
+def cargar_modelo(nombre_modelo):
+    """Carga el modelo y el tokenizer seleccionado."""
     global current_model, current_tokenizer, current_model_name
-    if current_model_name != model_name:
-        current_model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name]).to(device)
-        current_tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
-        current_model_name = model_name
+    if current_model_name != nombre_modelo:
+        current_model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[nombre_modelo]).to(device)
+        current_tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[nombre_modelo])
+        current_model_name = nombre_modelo
 
-# Load the default model at startup
-load_model("bloomz-560m")
+# Cargar el modelo por defecto
+cargar_modelo("BLOOMZ-560M")
 
 @spaces.GPU()
-def get_next_token_predictions(text, model_name, top_k=10):
-    """Generate the next token predictions with their probabilities."""
+def obtener_predicciones(texto, nombre_modelo, top_k=10):
+    """Genera las predicciones de las siguientes palabras con sus probabilidades."""
     global current_model, current_tokenizer
     
-    # Load the model if it has changed
-    if current_model_name != model_name:
-        load_model(model_name)
+    # Cargar modelo si ha cambiado
+    if current_model_name != nombre_modelo:
+        cargar_modelo(nombre_modelo)
     
-    inputs = current_tokenizer(text, return_tensors="pt").to(device)
+    entradas = current_tokenizer(texto, return_tensors="pt").to(device)
 
     with torch.no_grad():
-        outputs = current_model(**inputs)
-        logits = outputs.logits[0, -1, :]
-        probs = torch.nn.functional.softmax(logits, dim=-1)
+        salidas = current_model(**entradas)
+        logits = salidas.logits[0, -1, :]
+        probabilidades = torch.nn.functional.softmax(logits, dim=-1)
     
-    top_k_probs, top_k_indices = torch.topk(probs, k=top_k)
+    top_k_prob, top_k_indices = torch.topk(probabilidades, k=top_k)
     top_k_tokens = [current_tokenizer.decode([idx.item()]) for idx in top_k_indices]
     
-    return top_k_tokens, top_k_probs.cpu().tolist()
-
-def plot_probabilities(tokens, probs):
-    """Generate a horizontal bar chart for token probabilities."""
-    fig, ax = plt.subplots(figsize=(8, 5))
-    ax.barh(tokens[::-1], probs[::-1], color="skyblue")
-    ax.set_xlabel("Probability")
-    ax.set_title("Next Token Predictions")
-    plt.tight_layout()
-
-    # Save plot as an image and return the file path
-    plot_path = "token_probabilities.png"
-    plt.savefig(plot_path)
-    plt.close(fig)
-    
-    return plot_path
+    return top_k_tokens, top_k_prob.cpu().tolist()
+
+def generar_barplot(tokens, probabilidades):
+    """Genera una gráfica de barras de Gradio con las palabras más probables."""
+    datos = {"Palabra": tokens, "Probabilidad": probabilidades}
+    return datos
 
-def predict_next_token(model_name, text, top_k, custom_token=""):
-    """Get predictions and update the UI with text and a chart."""
-    if custom_token:
-        text += custom_token
+def predecir_siguiente_palabra(nombre_modelo, texto, top_k, token_custom=""):
+    """Obtiene predicciones y actualiza la UI."""
+    if token_custom:
+        texto += token_custom
 
-    tokens, probs = get_next_token_predictions(text, model_name, top_k)
+    tokens, probabilidades = obtener_predicciones(texto, nombre_modelo, int(top_k))
 
-    # Generate bar chart
-    plot_path = plot_probabilities(tokens, probs)
+    # Generar gráfico con Gradio BarPlot
+    barplot_data = generar_barplot(tokens, probabilidades)
 
-    return gr.update(choices=[f"'{t}'" for t in tokens]), plot_path
+    return gr.update(choices=[f"'{t}'" for t in tokens]), barplot_data
 
-def append_selected_token(text, selected_token):
-    """Append selected token from dropdown to the text input."""
-    if selected_token:
-        clean_token = selected_token.strip("'")
-        text += f" {clean_token}"
-    return text
+def agregar_token_seleccionado(texto, token_seleccionado):
+    """Agrega el token seleccionado al texto de entrada."""
+    if token_seleccionado:
+        token_limpio = token_seleccionado.strip("'")
+        texto += f" {token_limpio}"
+    return texto
 
-# Create the UI
+# Crear la interfaz en español
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔥 Interactive Text Prediction with Transformers")
+    gr.Markdown("# 🔥 Predicción de Texto con Modelos Transformadores")
     gr.Markdown(
-        "This application lets you interactively generate text using multiple transformer models. "
-        "Choose a model, type your text, and explore token predictions."
+        "Esta aplicación permite generar palabras utilizando un modelo de lenguaje. "
+        "Selecciona un modelo, introduce un texto y explora las palabras más probables a continuación."
     )
     
     with gr.Row():
-        model_dropdown = gr.Dropdown(
+        dropdown_modelo = gr.Dropdown(
             choices=list(AVAILABLE_MODELS.keys()),
-            value="distilgpt2",
-            label="Select Model"
+            value="BLOOMZ-560M",
+            label="📌 Modelo de lenguaje"
+        )
+
+        dropdown_top_k = gr.Dropdown(
+            choices=["5", "10", "15", "20"],
+            value="10",
+            label="🔢 Número de palabras a mostrar"
         )
     
     with gr.Row():
-        text_input = gr.Textbox(
+        texto_entrada = gr.Textbox(
             lines=5,
-            label="Input Text",
-            placeholder="Type your text here...",
-            value="The quick brown fox"
+            label="📝 Texto de entrada",
+            placeholder="Escribe aquí...",
+            value="Mi abuela me dejó una gran"
         )
     
     with gr.Row():
-        top_k_slider = gr.Slider(
-            minimum=1,
-            maximum=20,
-            value=10,
-            step=1,
-            label="Top-k Predictions"
-        )
+        boton_predecir = gr.Button("🔮 Predecir")
 
     with gr.Row():
-        predict_button = gr.Button("Predict")
-    
-    with gr.Row():
-        token_dropdown = gr.Dropdown(
-            label="Predicted Tokens",
+        dropdown_tokens = gr.Dropdown(
+            label="🔠 Palabras predichas",
             choices=[]
         )
-        append_button = gr.Button("Append Token")
+        boton_agregar = gr.Button("➕ Agregar palabra")
 
     with gr.Row():
-        chart_output = gr.Image(label="Token Probability Chart")
+        barplot_resultados = gr.BarPlot(
+            label="📊 Palabras más probables",
+            x="Palabra",
+            y="Probabilidad",
+            title="Predicciones del modelo"
+        )
 
-    # Button click events
-    predict_button.click(
-        predict_next_token,
-        inputs=[model_dropdown, text_input, top_k_slider],
-        outputs=[token_dropdown, chart_output]
+    # Acciones de botones
+    boton_predecir.click(
+        predecir_siguiente_palabra,
+        inputs=[dropdown_modelo, texto_entrada, dropdown_top_k],
+        outputs=[dropdown_tokens, barplot_resultados]
     )
 
-    append_button.click(
-        append_selected_token,
-        inputs=[text_input, token_dropdown],
-        outputs=text_input
+    boton_agregar.click(
+        agregar_token_seleccionado,
+        inputs=[texto_entrada, dropdown_tokens],
+        outputs=texto_entrada
     )
 
 demo.queue().launch()