Todo bien esta noche muack

app.py

@@ -2,25 +2,82 @@ import os
 import torch
 import gradio as gr
 from huggingface_hub import InferenceClient
-from model import model_params, AudioDataset
+from model import predict_params, AudioDataset
+from interfaz import estilo, my_theme
 
 token = os.getenv("HF_TOKEN")
 client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct", token=token)
+model_cache = {}
 
-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)}  # Add batch dimension
+def load_model_and_dataset(model_path, dataset_path, filter_white_noise):
+    if (model_path, dataset_path, filter_white_noise) not in model_cache:
+        model, _, _, id2label = predict_params(dataset_path, model_path, filter_white_noise)
+        model_cache[(model_path, dataset_path, filter_white_noise)] = (model, id2label)
+    return model_cache[(model_path, dataset_path, filter_white_noise)]
+
+def predict(audio_path, model_path, dataset_path, filter_white_noise):
+    model, id2label = load_model_and_dataset(model_path, dataset_path, filter_white_noise)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+    audios = AudioDataset(dataset_path, {}, filter_white_noise).preprocess_audio(audio_path)
+    inputs = {"input_values": audios.to(device).unsqueeze(0)}
     with torch.no_grad():
         outputs = model(**inputs)
-        predicted_class_ids = outputs.logits.argmax(-1)
-        label = id2label[predicted_class_ids.item()]
+        logits = outputs.logits
+        predicted_class_ids = torch.argmax(logits, dim=-1).item()
+        label = id2label[predicted_class_ids]
+        if dataset_path == "data/mixed_data":
+            label_mapping = {0: 'Hambre', 1: 'Problemas para respirar', 2: 'Dolor', 3: 'Cansancio/Incomodidad'}
+            label = label_mapping.get(predicted_class_ids, label)
     return label
 
-def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
+def predict_stream(audio_path):
+    model_mon, _ = load_model_and_dataset(
+        model_path="distilhubert-finetuned-cry-detector",
+        dataset_path="data/baby_cry_detection",
+        filter_white_noise=False
+        )
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model_mon.to(device)
+    model_mon.eval()
+    audio_dataset = AudioDataset(dataset_path="data/baby_cry_detection", label2id={}, filter_white_noise=False)
+    processed_audio = audio_dataset.preprocess_audio(audio_path)
+    inputs = {"input_values": processed_audio.to(device).unsqueeze(0)}
+    with torch.no_grad():
+        outputs = model_mon(**inputs)
+        logits = outputs.logits
+        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        crying_probabilities = probabilities[:, 1]
+        avg_crying_probability = crying_probabilities.mean().item()*100
+    if avg_crying_probability < 25:
+        model_class, id2label = load_model_and_dataset(
+            model_path="distilhubert-finetuned-mixed-data",
+            dataset_path="data/mixed_data",
+            filter_white_noise=True
+            )
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model_class.to(device)
+        model_class.eval()
+        audio_dataset_class = AudioDataset(dataset_path="data/mixed_data", label2id={}, filter_white_noise=True)
+        processed_audio_class = audio_dataset_class.preprocess_audio(audio_path)
+        inputs_class = {"input_values": processed_audio_class.to(device).unsqueeze(0)}
+        with torch.no_grad():
+            outputs_class = model_class(**inputs_class)
+            logits_class = outputs_class.logits
+            predicted_class_ids_class = torch.argmax(logits_class, dim=-1).item()
+            label_class = id2label[predicted_class_ids_class]
+            label_mapping = {0: 'Hambre', 1: 'Problemas para respirar', 2: 'Dolor', 3: 'Cansancio/Incomodidad'}
+            label_class = label_mapping.get(predicted_class_ids_class, label_class)
+        return f"Bebé llorando por {label_class}. Probabilidad: {avg_crying_probability:.1f})"
+    else:
+        return f"No está llorando. Proabilidad: {avg_crying_probability:.1f})"
+
+def chatbot_config(message, history: list[tuple[str, str]]):
+    system_message = "You are a Chatbot specialized in baby health and care."
+    max_tokens = 512
+    temperature = 0.7
+    top_p = 0.95
     messages = [{"role": "system", "content": system_message}]
     for val in history:
         if val[0]:
@@ -29,102 +86,22 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
             messages.append({"role": "assistant", "content": val[1]})
     messages.append({"role": "user", "content": message})
     response = ""
-    for message in client.chat_completion(messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p): # Creo que lo importante para el modelo
-        token = message.choices[0].delta.content
+    for message_response in client.chat_completion(messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p):
+        token = message_response.choices[0].delta.content
         response += token
         yield response
 
 def cambiar_pestaña():
     return gr.update(visible=False), gr.update(visible=True)
 
-my_theme = gr.themes.Soft(
-    primary_hue="emerald",
-    secondary_hue="green",
-    neutral_hue="slate",
-    text_size="sm",
-    spacing_size="sm",
-    font=[gr.themes.GoogleFont('Nunito'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
-    font_mono=[gr.themes.GoogleFont('Nunito'), 'ui-monospace', 'Consolas', 'monospace'],
-    ).set(
-    body_background_fill='*neutral_50',
-    body_text_color='*neutral_600',
-    body_text_size='*text_sm',
-    embed_radius='*radius_md',
-    shadow_drop='*shadow_spread',
-    shadow_spread='*button_shadow_active'
-    )
-
 with gr.Blocks(theme=my_theme) as demo:
-    with gr.Column(visible=True, elem_id="pantalla-inicial") as pantalla_inicial:
-        gr.HTML(
-            """
-            <style>
-            @import url('https://fonts.googleapis.com/css2?family=Lobster&display=swap');
-            @import url('https://fonts.googleapis.com/css2?family=Roboto&display=swap');
-            
-            h1 {
-                font-family: 'Lobster', cursive;
-                font-size: 5em !important;
-                text-align: center;
-                margin: 0;
-            }
-            
-            .gr-button {
-                background-color: #4CAF50 !important; 
-                color: white !important; 
-                border: none;
-                padding: 15px 32px;
-                text-align: center;
-                text-decoration: none;
-                display: inline-block;
-                font-size: 16px;
-                margin: 4px 2px;
-                cursor: pointer;
-                border-radius: 12px;
-            }
-            
-            .gr-button:hover {
-                background-color: #45a049; 
-            }
-            h2 {
-                font-family: 'Lobster', cursive;
-                font-size: 3em !important;
-                text-align: center;
-                margin: 0;
-            }
-            p.slogan, h4, p, h3 {
-                font-family: 'Roboto', sans-serif;
-                text-align: center;
-            }
-            </style>
-            <h1>Iremia</h1>
-            <h4 style='text-align: center; font-size: 1.5em'>Tu aliado para el bienestar de tu bebé</h4>
-            """
-        )
-        gr.Markdown(
-            "<h4 style='text-align: left; font-size: 1.5em;'>¿Qué es Iremia?</h4>"
-        )
-        gr.Markdown(
-            "<p style='text-align: left'>Iremia es un proyecto llevado a cabo por un grupo de estudiantes interesados en el desarrollo de modelos de inteligencia artificial, enfocados específicamente en casos de uso relevantes para ayudar a cuidar a los más pequeños de la casa.</p>"
-        )
-        gr.Markdown(
-            "<h4 style='text-align: left; font-size: 1.5em;'>Nuestra misión</h4>"
-        )
-        gr.Markdown(
-            "<p style='text-align: left'>Sabemos que la paternidad puede suponer un gran desafío. Nuestra misión es brindarles a todos los padres unas herramientas de última tecnología que los ayuden a navegar esos primeros meses de vida tan cruciales en el desarrollo de sus pequeños.</p>"
-        )
-        gr.Markdown(
-            "<h4 style='text-align: left; font-size: 1.5em;'>¿Qué ofrece Iremia?</h4>"
-        )
-        gr.Markdown(
-            "<p style='text-align: left'>Iremia ofrece dos funcionalidades muy interesantes:</p>"
-        )
-        gr.Markdown(
-            "<p style='text-align: left'>Predictor: Con nuestro modelo de inteligencia artificial, somos capaces de predecir por qué tu hijo de menos de 2 años está llorando. Además, tendrás acceso a un asistente personal para consultar cualquier duda que tengas sobre el cuidado de tu pequeño.</p>"
-        )
-        gr.Markdown(
-            "<p style='text-align: left'>Monitor: Nuestro monitor no es como otros que hay en el mercado, ya que es capaz de reconocer si un sonido es un llanto del bebé o no, y si está llorando, predice automáticamente la causa, lo cual te brindará la tranquilidad de saber siempre qué pasa con tu pequeño y te ahorrará tiempo y muchas horas de sueño.</p>"
-        )
+    estilo()
+    with gr.Column(visible=True) as chatbot:    
+        gr.Markdown("<h2>Asistente</h2>")
+        gr.ChatInterface(
+            chatbot_config # TODO: Mirar argumentos
+            )
+        gr.Markdown("Este chatbot no sustituye a un profesional de la salud. Ante cualquier preocupación o duda, consulta con tu pediatra.")
         with gr.Row():
             with gr.Column():
                 gr.Markdown("<h2>Predictor</h2>")
@@ -134,74 +111,41 @@ with gr.Blocks(theme=my_theme) as demo:
                 gr.Markdown("<h2>Monitor</h2>")
                 boton_pagina_2 = gr.Button("Prueba el monitor")
                 gr.Markdown("<p>Un monitor inteligente que detecta si tu hijo está llorando y te indica el motivo antes de que puedas levantarte del sofá</p>")
-    with gr.Column(visible=False) as pagina_1:
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("<h2>Predictor</h2>")
-                audio_input = gr.Audio(
-                    min_length=1.0,
-                    format="wav",
-                    label="Baby recorder",
-                    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(
-                    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."
-                max_tokens = 512
-                temperature = 0.7
-                top_p = 0.95
-                chatbot = gr.ChatInterface(
-                    respond, # TODO: Cambiar para que argumentos estén aquí metidos
-                    additional_inputs=[
-                        gr.State(value=system_message),
-                        gr.State(value=max_tokens),
-                        gr.State(value=temperature),
-                        gr.State(value=top_p)
-                    ],
-                )
-                gr.Markdown("Este chatbot no sustituye a un profesional de la salud. Ante cualquier preocupación o duda, consulta con tu pediatra.")
-                boton_volver_inicio_1 = gr.Button("Volver a la pantalla inicial").click(cambiar_pestaña, outputs=[pagina_1, pantalla_inicial])
-    with gr.Column(visible=False) as pagina_2:
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("<h2>Monitor</h2>")
-                audio_input = gr.Audio(
-                    min_length=1.0,
-                    format="wav",
-                    label="Baby recorder",
-                    type="filepath", # Para no usar numpy y preprocesar siempre igual
-                    )
-                classify_btn = gr.Button("¿Por qué llora?")
-                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."
-                max_tokens = 512
-                temperature = 0.7
-                top_p = 0.95
-                chatbot = gr.ChatInterface(
-                    respond, # TODO: Cambiar para que argumentos estén aquí metidos
-                    additional_inputs=[
-                        gr.State(value=system_message),
-                        gr.State(value=max_tokens),
-                        gr.State(value=temperature),
-                        gr.State(value=top_p)
-                    ],
-                )
-                gr.Markdown("Este chatbot no sustituye a un profesional de la salud. Ante cualquier preocupación o duda, consulta con tu pediatra.")
-                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()
+    with gr.Column(visible=False) as pag_predictor:
+        gr.Markdown("<h2>Predictor</h2>")
+        audio_input = gr.Audio(
+            min_length=1.0,
+            format="wav",
+            label="Baby recorder",
+            type="filepath",
+            )
+        classify_btn = gr.Button("¿Por qué llora?")
+        classify_btn.click(
+            lambda audio: predict( # Mirar porque usar lambda
+                audio,
+                model_path="distilhubert-finetuned-mixed-data",
+                dataset_path="data/mixed_data",
+                filter_white_noise=True
+                ),
+            inputs=audio_input,
+            outputs=gr.Textbox(label="Tu bebé llora por:")
+            )
+        gr.Button("Volver a la pantalla inicial").click(cambiar_pestaña, outputs=[pag_predictor, chatbot])
+    with gr.Column(visible=False) as pag_monitor:
+        gr.Markdown("<h2>Monitor</h2>")
+        audio_stream = gr.Audio(
+                # min_length=1.0, # mirar por qué no va esto
+                format="wav",
+                label="Baby recorder",
+                type="filepath",
+                streaming=True
+            )
+        audio_stream.stream(
+            predict_stream,
+            inputs=audio_stream,
+            outputs=gr.Textbox(label="Tu bebé está:"),
+        )
+        gr.Button("Volver a la pantalla inicial").click(cambiar_pestaña, outputs=[pag_monitor, chatbot])
+    boton_pagina_1.click(cambiar_pestaña, outputs=[chatbot, pag_predictor])
+    boton_pagina_2.click(cambiar_pestaña, outputs=[chatbot, pag_monitor])
+demo.launch(share=True)

interfaz.py

@@ -0,0 +1,101 @@
+import gradio as gr
+
+my_theme = gr.themes.Soft(
+    primary_hue="emerald",
+    secondary_hue="green",
+    neutral_hue="slate",
+    text_size="sm",
+    spacing_size="sm",
+    font=[gr.themes.GoogleFont('Nunito'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
+    font_mono=[gr.themes.GoogleFont('Nunito'), 'ui-monospace', 'Consolas', 'monospace'],
+    ).set(
+    body_background_fill='*neutral_50',
+    body_text_color='*neutral_600',
+    body_text_size='*text_sm',
+    embed_radius='*radius_md',
+    shadow_drop='*shadow_spread',
+    shadow_spread='*button_shadow_active'
+    )
+def estilo():
+    gr.HTML(
+        """
+        <style>
+        @import url('https://fonts.googleapis.com/css2?family=Lobster&display=swap');
+        @import url('https://fonts.googleapis.com/css2?family=Roboto&display=swap');
+        
+        h1 {
+            font-family: 'Lobster', cursive;
+            font-size: 5em !important;
+            text-align: center;
+            margin: 0;
+        }
+        
+        .gr-button {
+            background-color: #4CAF50 !important; 
+            color: white !important; 
+            border: none;
+            padding: 15px 32px;
+            text-align: center;
+            text-decoration: none;
+            display: inline-block;
+            font-size: 16px;
+            margin: 4px 2px;
+            cursor: pointer;
+            border-radius: 12px;
+        }
+        
+        .gr-button:hover {
+            background-color: #45a049; 
+        }
+        h2 {
+            font-family: 'Lobster', cursive;
+            font-size: 3em !important;
+            text-align: center;
+            margin: 0;
+        }
+        p.slogan, h4, p, h3 {
+            font-family: 'Roboto', sans-serif;
+            text-align: center;
+        }
+        </style>
+        <h1>Iremia</h1>
+        <h4 style='text-align: center; font-size: 1.5em'>Tu aliado para el bienestar de tu bebé</h4>
+        """
+    )
+    return my_theme
+def inicio():
+    estilo()
+    gr.Markdown(
+        "<h4 style='text-align: left; font-size: 1.5em;'>¿Qué es Iremia?</h4>"
+    )
+    gr.Markdown(
+        "<p style='text-align: left'>Iremia es un proyecto llevado a cabo por un grupo de estudiantes interesados en el desarrollo de modelos de inteligencia artificial, enfocados específicamente en casos de uso relevantes para ayudar a cuidar a los más pequeños de la casa.</p>"
+    )
+    gr.Markdown(
+        "<h4 style='text-align: left; font-size: 1.5em;'>Nuestra misión</h4>"
+    )
+    gr.Markdown(
+        "<p style='text-align: left'>Sabemos que la paternidad puede suponer un gran desafío. Nuestra misión es brindarles a todos los padres unas herramientas de última tecnología que los ayuden a navegar esos primeros meses de vida tan cruciales en el desarrollo de sus pequeños.</p>"
+    )
+    gr.Markdown(
+        "<h4 style='text-align: left; font-size: 1.5em;'>¿Qué ofrece Iremia?</h4>"
+    )
+    gr.Markdown(
+        "<p style='text-align: left'>Iremia ofrece dos funcionalidades muy interesantes:</p>"
+    )
+    gr.Markdown(
+        "<p style='text-align: left'>Predictor: Con nuestro modelo de inteligencia artificial, somos capaces de predecir por qué tu hijo de menos de 2 años está llorando. Además, tendrás acceso a un asistente personal para consultar cualquier duda que tengas sobre el cuidado de tu pequeño.</p>"
+    )
+    gr.Markdown(
+        "<p style='text-align: left'>Monitor: Nuestro monitor no es como otros que hay en el mercado, ya que es capaz de reconocer si un sonido es un llanto del bebé o no, y si está llorando, predice automáticamente la causa, lo cual te brindará la tranquilidad de saber siempre qué pasa con tu pequeño y te ahorrará tiempo y muchas horas de sueño.</p>"
+    )
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("<h2>Predictor</h2>")
+            boton_pagina_1 = gr.Button("Prueba el predictor")
+            gr.Markdown("<p>Descubre por qué llora tu bebé y resuelve dudas sobre su cuidado con nuestro Iremia assistant</p>")
+        with gr.Column():
+            gr.Markdown("<h2>Monitor</h2>")
+            boton_pagina_2 = gr.Button("Prueba el monitor")
+            gr.Markdown("<p>Un monitor inteligente que detecta si tu hijo está llorando y te indica el motivo antes de que puedas levantarte del sofá</p>")
+    return boton_pagina_1, boton_pagina_2

model.py

@@ -23,11 +23,13 @@ config_file = "models_config.json"
 clasificador = "class"
 monitor = "mon"
 batch_size = 16
+num_workers = 12
 
 class AudioDataset(Dataset):
-    def __init__(self, dataset_path, label2id):
+    def __init__(self, dataset_path, label2id, filter_white_noise):
         self.dataset_path = dataset_path
         self.label2id = label2id
+        self.filter_white_noise = filter_white_noise
         self.file_paths = []
         self.labels = []
         for label_dir, label_id in self.label2id.items():
@@ -37,7 +39,7 @@ class AudioDataset(Dataset):
                     audio_path = os.path.join(label_path, file_name)
                     self.file_paths.append(audio_path)
                     self.labels.append(label_id)
-        self.file_paths.sort(key=lambda x: x.split('_part')[0])
+        self.file_paths.sort(key=lambda x: x.split('_part')[0]) # no sé si influye
 
     def __len__(self):
         return len(self.file_paths)
@@ -55,29 +57,33 @@ class AudioDataset(Dataset):
         waveform, sample_rate = torchaudio.load(
             audio_path,
             normalize=True, # Convierte a float32
-            # num_frames= # TODO: Probar para que no haga falta recortar los audios
             )
         if sample_rate != SAMPLING_RATE: # Resamplear si no es 16kHz
             resampler = torchaudio.transforms.Resample(sample_rate, SAMPLING_RATE)
             waveform = resampler(waveform)
         if waveform.shape[0] > 1: # Si es stereo, convertir a mono
             waveform = waveform.mean(dim=0, keepdim=True)
-        waveform = waveform / (torch.max(torch.abs(waveform)) + 1e-6) # Sin 1e-6 el accuracy es pésimo!!
+        waveform = waveform / (torch.max(torch.abs(waveform)) + 1e-6) # Normalizar, sin 1e-6 el accuracy es pésimo!!
         max_length = int(SAMPLING_RATE * MAX_DURATION)
         if waveform.shape[1] > max_length:
-            waveform = waveform[:, :max_length]
+            waveform = waveform[:, :max_length] # Truncar
         else:
-            waveform = torch.nn.functional.pad(waveform, (0, max_length - waveform.shape[1]))
+            waveform = torch.nn.functional.pad(waveform, (0, max_length - waveform.shape[1])) # Padding
         inputs = FEATURE_EXTRACTOR(
             waveform.squeeze(),
-            sampling_rate=SAMPLING_RATE,
+            sampling_rate=SAMPLING_RATE, # Hecho a mano, por si acaso
             return_tensors="pt",
             # max_length=int(SAMPLING_RATE * MAX_DURATION),
-            # truncation=True,
-            padding=True,
+            # truncation=True, # Hecho a mano
+            # padding=True, # Hecho a mano
         )
         return inputs.input_values.squeeze()
 
+def is_white_noise(audio):
+    mean = torch.mean(audio)
+    std = torch.std(audio)
+    return torch.abs(mean) < 0.001 and std < 0.01
+
 def seed_everything():
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)
@@ -96,9 +102,9 @@ def build_label_mappings(dataset_path):
             label_id += 1
     return label2id, id2label
 
-def create_dataloader(dataset_path, test_size=0.2, num_workers=12, shuffle=True, pin_memory=True):
+def create_dataloader(dataset_path, filter_white_noise, test_size=0.2, shuffle=True, pin_memory=True):
     label2id, id2label = build_label_mappings(dataset_path)
-    dataset = AudioDataset(dataset_path, label2id)
+    dataset = AudioDataset(dataset_path, label2id, filter_white_noise)
     dataset_size = len(dataset)
     indices = list(range(dataset_size))
     random.shuffle(indices)
@@ -115,9 +121,9 @@ def create_dataloader(dataset_path, test_size=0.2, num_workers=12, shuffle=True,
     )
     return train_dataloader, test_dataloader, label2id, id2label
 
-def load_model(num_labels, label2id, id2label):
+def load_model(model_path, num_labels, label2id, id2label):
     config = HubertConfig.from_pretrained(
-        MODEL,
+        pretrained_model_name_or_path=model_path,
         num_labels=num_labels,
         label2id=label2id,
         id2label=id2label,
@@ -125,18 +131,23 @@ def load_model(num_labels, label2id, id2label):
     )
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = HubertForSequenceClassification.from_pretrained( # TODO: mirar parámetros. Posibles optimizaciones
-        MODEL,
+        pretrained_model_name_or_path=model_path,
         config=config,
-        torch_dtype=torch.float32, # No afecta 1ª época, mejor ponerlo
+        torch_dtype=torch.float32,
     )
     model.to(device)
     return model
 
-def model_params(dataset_path):
-    train_dataloader, test_dataloader, label2id, id2label = create_dataloader(dataset_path)
-    model = load_model(num_labels=len(id2label), label2id=label2id, id2label=id2label)    
+def train_params(dataset_path, filter_white_noise):
+    train_dataloader, test_dataloader, label2id, id2label = create_dataloader(dataset_path, filter_white_noise)
+    model = load_model(model_path=MODEL, num_labels=len(id2label), label2id=label2id, id2label=id2label)    
     return model, train_dataloader, test_dataloader, id2label
 
+def predict_params(dataset_path, model_path, filter_white_noise):
+    _, _, label2id, id2label = create_dataloader(dataset_path, filter_white_noise)
+    model = load_model(model_path, num_labels=len(id2label), label2id=label2id, id2label=id2label)
+    return model, None, None, id2label
+
 def compute_metrics(eval_pred):
     predictions = torch.argmax(torch.tensor(eval_pred.predictions), dim=-1)
     references = torch.tensor(eval_pred.label_ids)
@@ -149,9 +160,9 @@ def compute_metrics(eval_pred):
         "f1": f1,
     }
 
-def main(training_args, output_dir, dataset_path):
+def main(training_args, output_dir, dataset_path, filter_white_noise):
     seed_everything()
-    model, train_dataloader, test_dataloader, _ = model_params(dataset_path)
+    model, train_dataloader, test_dataloader, _ = train_params(dataset_path, filter_white_noise)
     trainer = Trainer(
         model=model,
         args=training_args,
@@ -162,9 +173,10 @@ def main(training_args, output_dir, dataset_path):
     )
     torch.cuda.empty_cache() # liberar memoria de la GPU
     trainer.train() # se pueden modificar los parámetros para continuar el train
-    os.makedirs(output_dir, exist_ok=True) # Crear carpeta con el modelo si no existe
-    trainer.save_model(output_dir) # para subir el modelo a Hugging Face. Necesario para hacer la predicción, no sé por qué.
-    # upload_folder(repo_id=f"A-POR-LOS-8000/{output_dir}",folder_path=output_dir, token=token) # subir modelo a organización
+    # trainer.save_model(output_dir) # Guardar modelo local.
+    os.makedirs(output_dir, exist_ok=True) # Crear carpeta
+    trainer.push_to_hub(token=token) # Subir modelo a perfil
+    upload_folder(repo_id=f"A-POR-LOS-8000/{output_dir}", folder_path=output_dir, token=token) # subir a organización y local
 
 def load_config(model_name):
     with open(config_file, 'r') as f:
@@ -176,8 +188,10 @@ def load_config(model_name):
 
 if __name__ == "__main__":
     config = load_config(clasificador) # PARA CAMBIAR MODELOS
+    filter_white_noise = True
     # config = load_config(monitor) # PARA CAMBIAR MODELOS
+    # filter_white_noise = False
     training_args = config["training_args"]
     output_dir = config["output_dir"]
     dataset_path = config["dataset_path"]
-    main(training_args, output_dir, dataset_path)
+    main(training_args, output_dir, dataset_path, filter_white_noise)