test

run_continue_nst.sh

@@ -14,9 +14,9 @@ python run_whisper_finetuning.py \
 	--do_eval=True \
 	--audio_column_name="audio" \
 	--text_column_name="text" \
-	--per_device_train_batch_size=48 \
-    	--per_device_train_batch_size=48 \
-	--learning_rate=4e-5 \
+	--per_device_train_batch_size=24 \
+    	--per_device_train_batch_size=24 \
+	--learning_rate=2e-5 \
 	--warmup_steps=500 \
 	--max_steps=10000 \
 	--gradient_checkpointing=True \

run_whisper_finetuning.py

@@ -345,7 +345,6 @@ def main():
     model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
     # Metrics
-
     def compute_metrics(pred):
         pred_ids = pred.predictions
         label_ids = pred.label_ids
@@ -383,30 +382,40 @@ def main():
         print("\n* Training arguments")
         pprint(vars(training_args), indent=2)
 
-    def rename_column(ds, old_name, new_name):
-        feats = ds.info.features
-        ds = ds.rename_column(old_name, new_name)
-        feats[new_name] = feats.pop(old_name)
-        ds.info.features = feats
-        return ds
-
-    def remove_columns(ds, column_name):
-        feats = ds.info.features
-        ds = ds.remove_columns(column_name)
-        feats.pop(column_name)
-        ds.info.features = feats
-        return ds
 
     # Print training arguments
     if data_args.print_training_arguments:
         print_training_arguments(model_args, data_args, training_args)
 
+    
+  
+    # Initialise the model
+    feature_extractor = WhisperFeatureExtractor.from_pretrained(
+        model_args.model_name_or_path)
+    tokenizer = WhisperTokenizer.from_pretrained(
+        model_args.model_name_or_path, language=model_args.language, task=model_args.task)
+    processor = WhisperProcessor.from_pretrained(
+        model_args.model_name_or_path, language=model_args.language, task=model_args.task)
+    data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+    # Saving the processor and the tokenizer
+    processor.save_pretrained(training_args.output_dir)
+    tokenizer.save_pretrained(training_args.output_dir)
+
+
     # Load dataset
     train_dataset = load_dataset(data_args.dataset_name, data_args.dataset_config_name,
                                  split="train", streaming=True, use_auth_token=True)
     eval_dataset = load_dataset(data_args.dataset_name, data_args.dataset_config_name,
                                 split="test", streaming=True, use_auth_token=True)
-
+    
+    # Because a bug in Dataset (https://github.com/huggingface/datasets/issues/3888) we need to read the columns and keep them for later
+    column_names=[x for x in train_dataset.info.features]
+    
+    # Make sure everything is in 16K
+    train_dataset = train_dataset.cast_column(data_args.audio_column_name, Audio(sampling_rate=16000))
+    eval_dataset = eval_dataset.cast_column(data_args.audio_column_name, Audio(sampling_rate=16000))
+    
     # Rename columns
     if data_args.audio_column_name != "audio":
         train_dataset = train_dataset.rename_column(
@@ -420,42 +429,13 @@ def main():
         eval_dataset = eval_dataset.rename_column(
             data_args.text_column_name, "sentence")
 
-    # Initialise
-    feature_extractor = WhisperFeatureExtractor.from_pretrained(
-        model_args.model_name_or_path)
-    tokenizer = WhisperTokenizer.from_pretrained(
-        model_args.model_name_or_path, language=model_args.language, task=model_args.task)
-    processor = WhisperProcessor.from_pretrained(
-        model_args.model_name_or_path, language=model_args.language, task=model_args.task)
-    data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
-
-    # Saving the processor and the tokenizer
-    processor.save_pretrained(training_args.output_dir)
-    tokenizer.save_pretrained(training_args.output_dir)
-
-    # Prepare data
-    # TODO The casting of the  not working on the NPSC in 48K. It seems to be working for Common Voice
-    # The issue is that the dataset features returns None. But for me thay seem to have been set correctly
-    # In our case this is not needed, since the datasets already is available as 16K. But it would be great to solve this bug
-    # train_dataset = train_dataset.cast_column("audio", Audio(sampling_rate=16000))
-    # eval_dataset = eval_dataset.cast_column("audio", Audio(sampling_rate=16000))
-
-    # Remove non needed columns
-    #column_names=[x for x in train_dataset.info.features]
-
-    # for c in column_names:
-    #    if c not in ["audio", "text"]:
-    #        train_dataset = remove_columns(train_dataset, c)
-    #        eval_dataset = remove_columns(eval_dataset, c)
-
-    # TODO I would really like to remove the non needed columns here. At least this cleans up the output.
-    # I am unable to figure out how to do this Streaming mode. Can not find a way to list columns.
-    # train_data = train_data.map(prepare_dataset, remove_columns=train_data.column_names, num_proc=1)
-
-    train_dataset = train_dataset.map(prepare_dataset)
-    eval_dataset = eval_dataset.map(prepare_dataset)
+    
+    # Prepare the dataset
+    column_names.extend(['sentence','audio'])
+    train_dataset = train_dataset.map(prepare_dataset, remove_columns=column_names)
+    eval_dataset = eval_dataset.map(prepare_dataset, remove_columns=column_names)
 
-    # Metrics
+    # Define metrics
     metric = evaluate.load("wer")
 
     # Detecting last checkpoint.
@@ -547,13 +527,10 @@ def main():
     return train_result
 
 # XLA hook
-
-
 def _mp_fn(index):
     # For xla_spawn (TPUs)
     print("The XLA is initiated")
     main()
 
-
 if __name__ == "__main__":
     main()