Usar label2id menos

model.py

@@ -14,17 +14,15 @@ from transformers import (
     EarlyStoppingCallback
 )
 
-MODEL = "ntu-spml/distilhubert" # modelo base utilizado, para usar otro basta con cambiar esto
-FEATURE_EXTRACTOR = Wav2Vec2FeatureExtractor.from_pretrained(MODEL)
+MODEL = "ntu-spml/distilhubert" # modelo base
+FEATURE_EXTRACTOR = Wav2Vec2FeatureExtractor.from_pretrained(MODEL) # feature extractor del modelo base
 seed = 123
-MAX_DURATION = 1.00
-SAMPLING_RATE = FEATURE_EXTRACTOR.sampling_rate # 16000
+MAX_DURATION = 1.00 # Máxima duración de los audios
+SAMPLING_RATE = FEATURE_EXTRACTOR.sampling_rate # 16kHz
 token = os.getenv("HF_TOKEN") # TODO: probar a guardar el token en un archivo en local
 config_file = "models_config.json"
-clasificador = "class"
-monitor = "mon"
-batch_size = 16
-num_workers = 12
+batch_size = 1024 # TODO: repasar si sigue siendo necesario
+num_workers = 12 # Núcleos de la CPU
 
 class AudioDataset(Dataset):
     def __init__(self, dataset_path, label2id, filter_white_noise):
@@ -120,13 +118,12 @@ def create_dataloader(dataset_path, filter_white_noise, test_size=0.2, shuffle=T
     test_dataloader = DataLoader(
         test_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory
     )
-    return train_dataloader, test_dataloader, label2id, id2label
+    return train_dataloader, test_dataloader, id2label
 
-def load_model(model_path, label2id, id2label, num_labels):
+def load_model(model_path, id2label, num_labels):
     config = HubertConfig.from_pretrained(
         pretrained_model_name_or_path=model_path,
         num_labels=num_labels,
-        label2id=label2id,
         id2label=id2label,
         finetuning_task="audio-classification"
     )
@@ -140,13 +137,13 @@ def load_model(model_path, label2id, id2label, num_labels):
     return model
 
 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, label2id, id2label, num_labels=len(id2label))    
+    train_dataloader, test_dataloader, id2label = create_dataloader(dataset_path, filter_white_noise)
+    model = load_model(MODEL, id2label, num_labels=len(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, label2id, id2label, num_labels=len(id2label))
+    _, _, id2label = create_dataloader(dataset_path, filter_white_noise)
+    model = load_model(model_path, id2label, num_labels=len(id2label))
     return model, id2label
 
 def compute_metrics(eval_pred):
@@ -173,11 +170,11 @@ def main(training_args, output_dir, dataset_path, filter_white_noise):
         callbacks=[TensorBoardCallback(), EarlyStoppingCallback(early_stopping_patience=3)]
     )
     torch.cuda.empty_cache() # liberar memoria de la GPU
-    trainer.train() # se pueden modificar los parámetros para continuar el train
+    trainer.train() # resume_from_checkpoint para continuar el train
     # trainer.save_model(output_dir) # Guardar modelo local.
-    os.makedirs(output_dir, exist_ok=True) # Crear carpeta
+    os.makedirs(output_dir, exist_ok=True)
     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
+    upload_folder(repo_id=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: