Update model.py

model.py

@@ -6,6 +6,7 @@ import torchaudio
 from torch.utils.data import Dataset, DataLoader
 from huggingface_hub import login, upload_folder
 from transformers.integrations import TensorBoardCallback
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
 from transformers import (
     Wav2Vec2FeatureExtractor, HubertConfig, HubertForSequenceClassification,
     Trainer, TrainingArguments,
@@ -17,11 +18,11 @@ FEATURE_EXTRACTOR = Wav2Vec2FeatureExtractor.from_pretrained(MODEL)
 seed = 123
 MAX_DURATION = 1.00
 SAMPLING_RATE = FEATURE_EXTRACTOR.sampling_rate # 16000
-token = os.getenv('MODEL_REPO_ID')
+token = os.getenv("HF_TOKEN")
 config_file = "models_config.json"
 clasificador = "class"
 monitor = "mon"
-batch_size = 4096
+batch_size = 16
 
 class AudioDataset(Dataset):
     def __init__(self, dataset_path, label2id):
@@ -60,13 +61,19 @@ class AudioDataset(Dataset):
             waveform = resampler(waveform)
         if waveform.shape[0] > 1: # Si es stereo, convertir a mono
             waveform = waveform.mean(dim=0)
-        waveform = waveform / torch.max(torch.abs(waveform))
+        waveform = waveform / (torch.max(torch.abs(waveform)) + 1e-6) # Sin 1e-6 el accuracy es pésimo!!
+        max_length = int(SAMPLING_RATE * MAX_DURATION)
+        if waveform.shape[0] > max_length:
+            waveform = waveform[:max_length]
+        else:
+            # Pad the waveform if it's shorter than max length
+            waveform = torch.nn.functional.pad(waveform, (0, max_length - waveform.shape[0]))
         inputs = FEATURE_EXTRACTOR(
             waveform,
             sampling_rate=SAMPLING_RATE,
             return_tensors="pt",
-            max_length=int(SAMPLING_RATE * MAX_DURATION),
-            truncation=True,
+            # max_length=int(SAMPLING_RATE * MAX_DURATION),
+            # truncation=True,
             padding=True,
         )
         return inputs.input_values.squeeze()
@@ -131,10 +138,15 @@ def model_params(dataset_path):
     return model, train_dataloader, test_dataloader, id2label
 
 def compute_metrics(eval_pred):
-    predictions = torch.argmax(input=eval_pred.predictions)
-    references = eval_pred.label_ids
+    predictions = torch.argmax(torch.tensor(eval_pred.predictions), dim=-1)
+    references = torch.tensor(eval_pred.label_ids)
+    accuracy = accuracy_score(references, predictions)
+    precision, recall, f1, _ = precision_recall_fscore_support(references, predictions, average='weighted')
     return {
-        "accuracy": torch.mean(predictions == references),
+        "accuracy": accuracy,
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
     }
 
 def main(training_args, output_dir, dataset_path):