Usar solo una función para ambos modelos

app.py

@@ -7,28 +7,13 @@ from model import model_params, AudioDataset
 token = os.getenv("HF_TOKEN")
 client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct", token=token)
 
-def predict_class(audio_path):
-    dataset_path = f"data/mixed_data" # PARA CLASIFICADOR
+def predict(audio_path, dataset_path):
     model, _, _, id2label = model_params(dataset_path)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# Usar a GPU o CPU
     model.to(device)# Usar a GPU o CPU
     audio_dataset = AudioDataset(dataset_path, {})
     inputs = audio_dataset.preprocess_audio(audio_path)
-    inputs = {"input_values": inputs.to(device).unsqueeze(0)}
-    with torch.no_grad():
-        outputs = model(**inputs)
-        predicted_class_ids = outputs.logits.argmax(-1)
-        label = id2label[predicted_class_ids.item()]
-    return label
-
-def predict_mon(audio_path):
-    dataset_path = f"data/baby_cry_detection" # PARA CLASIFICADOR
-    model, _, _, id2label = model_params(dataset_path)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# Usar a GPU o CPU
-    model.to(device)# Usar a GPU o CPU
-    audio_dataset = AudioDataset(dataset_path, {})
-    inputs = audio_dataset.preprocess_audio(audio_path)
-    inputs = {"input_values": inputs.to(device).unsqueeze(0)}
+    inputs = {"input_values": inputs.to(device).unsqueeze(0)}  # Add batch dimension
     with torch.no_grad():
         outputs = model(**inputs)
         predicted_class_ids = outputs.logits.argmax(-1)
@@ -161,7 +146,11 @@ with gr.Blocks(theme=my_theme) as demo:
                     )
                 classify_btn = gr.Button("¿Por qué llora?")
                 classification_output = gr.Textbox(label="Tu bebé llora por:")
-                classify_btn.click(predict_class, inputs=audio_input, outputs=classification_output)
+                classify_btn.click(
+                    lambda audio: predict(audio, dataset_path="data/mixed_data"),
+                    inputs=audio_input,
+                    outputs=classification_output
+                    )
             with gr.Column():
                 gr.Markdown("<h2>Assistant</h2>")
                 system_message = "You are a Chatbot specialized in baby health and care."
@@ -190,8 +179,12 @@ with gr.Blocks(theme=my_theme) as demo:
                     type="filepath", # Para no usar numpy y preprocesar siempre igual
                     )
                 classify_btn = gr.Button("¿Por qué llora?")
-                classification_output = gr.Textbox(label="Tu bebé llora por:")
-                classify_btn.click(predict_mon, inputs=audio_input, outputs=classification_output)
+                classification_output = gr.Textbox(label="Tu bebé está:")
+                classify_btn.click(
+                    lambda audio: predict(audio, dataset_path="data/baby_cry_detection"),
+                    inputs=audio_input,
+                    outputs=classification_output
+                    )
             with gr.Column():
                 gr.Markdown("<h2>Assistant</h2>")
                 system_message = "You are a Chatbot specialized in baby health and care."
@@ -208,8 +201,7 @@ with gr.Blocks(theme=my_theme) as demo:
                     ],
                 )
                 gr.Markdown("Este chatbot no sustituye a un profesional de la salud. Ante cualquier preocupación o duda, consulta con tu pediatra.")
-        gr.Markdown("Contenido de la Página 2")
-        boton_volver_inicio_2 = gr.Button("Volver a la pantalla inicial").click(cambiar_pestaña, outputs=[pagina_2, pantalla_inicial])
+                boton_volver_inicio_2 = gr.Button("Volver a la pantalla inicial").click(cambiar_pestaña, outputs=[pagina_2, pantalla_inicial])
     boton_pagina_1.click(cambiar_pestaña, outputs=[pantalla_inicial, pagina_1])
     boton_pagina_2.click(cambiar_pestaña, outputs=[pantalla_inicial, pagina_2])
 demo.launch()