fix: Set Whisper TrainingArguments as wav2vec2

alexandrainst · Oct 24, 2024 · 2eff85b · 2eff85b
1 parent b85036c
commit 2eff85b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/coral/whisper.py b/src/coral/whisper.py
@@ -15,6 +15,7 @@
     AutoConfig,
     AutoModelForSpeechSeq2Seq,
     EvalPrediction,
+    SchedulerType,
     Seq2SeqTrainer,
     Seq2SeqTrainingArguments,
     Trainer,
@@ -167,6 +168,7 @@ def load_training_arguments(self) -> TrainingArguments:
             per_device_eval_batch_size=self.config.per_device_batch_size,
             gradient_accumulation_steps=gradient_accumulation_steps,
             learning_rate=self.config.learning_rate,
+            lr_scheduler_type=SchedulerType.COSINE,
             warmup_steps=self.config.warmup_steps,
             max_steps=self.config.max_steps,
             fp16=fp16,
@@ -178,14 +180,16 @@ def load_training_arguments(self) -> TrainingArguments:
             save_strategy="no" if self.config.save_total_limit == 0 else "steps",
             logging_steps=self.config.logging_steps,
             length_column_name="input_length",
-            gradient_checkpointing=True,
+            gradient_checkpointing=self.config.gradient_checkpointing,
             save_total_limit=self.config.save_total_limit,
             load_best_model_at_end=self.config.early_stopping,
             metric_for_best_model="wer",
             greater_is_better=False,
             seed=self.config.seed,
             remove_unused_columns=False,
             optim=OptimizerNames.ADAMW_TORCH,
+            adam_beta1=self.config.adam_first_momentum,
+            adam_beta2=self.config.adam_second_momentum,
             report_to=["wandb"] if self.config.wandb else [],
             ignore_data_skip=self.config.ignore_data_skip,
             save_safetensors=True,