allowing multiple datasets

run_speech_recognition_seq2seq_streaming.py

@@ -328,24 +328,28 @@ def notify_me(recipient, message=None):
         smtp_obj.quit()
 
 
-def load_maybe_streaming_dataset(dataset_name, dataset_config_name, split="train", streaming=True, **kwargs):
+def load_maybe_streaming_dataset(dataset_names, dataset_config_names, split="train", streaming=True, **kwargs):
     """
     Utility function to load a dataset in streaming mode. For datasets with multiple splits,
     each split is loaded individually and then splits combined by taking alternating examples from
     each (interleaving).
     """
-    if "+" in split:
+    if "," in dataset_names or "+" in split:
         # load multiple splits separated by the `+` symbol with streaming mode
-        dataset_splits = [
-            load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=streaming, **kwargs)
-            for split_name in split.split("+")
-        ]
+        dataset_splits = []
+        for dataset_name, dataset_config_name, split_names in zip(
+            dataset_names.split(","), dataset_config_names.split(","), split.split(",")
+        ):
+            for split_name in split_names.split("+"):
+                dataset = load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=streaming, **kwargs)
+                dataset_splits.append(dataset)
+    
         # interleave multiple splits to form one dataset
         interleaved_dataset = interleave_datasets(dataset_splits)
         return interleaved_dataset
     else:
         # load a single split *with* streaming mode
-        dataset = load_dataset(dataset_name, dataset_config_name, split=split, streaming=streaming, **kwargs)
+        dataset = load_dataset(dataset_names, dataset_config_names, split=split, streaming=streaming, **kwargs)
         return dataset
 
 
@@ -652,14 +656,22 @@ def main():
         elif last_checkpoint is not None:
             checkpoint = last_checkpoint
         train_result = trainer.train(resume_from_checkpoint=checkpoint)
+        logger.info("*** Training completed ***")
+        logger.info("*** Saving model ***")
         trainer.save_model()  # Saves the feature extractor too for easy upload
-
+        logger.info("*** Model saves ***")
         metrics = train_result.metrics
         if data_args.max_train_samples:
             metrics["train_samples"] = data_args.max_train_samples
+        logger.info("*** Logging metrics ***")
         trainer.log_metrics("train", metrics)
+        logger.info("*** Metrics logged ***")
+        logger.info("*** Saving metrics ***")
         trainer.save_metrics("train", metrics)
+        logger.info("*** Metrics saved ***")
+        logger.info("*** Saving state ***")
         trainer.save_state()
+        logger.info("*** State saved ***")
 
     # 13. Evaluation
     results = {}
@@ -670,13 +682,18 @@ def main():
             max_length=training_args.generation_max_length,
             num_beams=training_args.generation_num_beams,
         )
+        logger.info("*** Evaluation done ***")
         if data_args.max_eval_samples:
             metrics["eval_samples"] = data_args.max_eval_samples
-
+        logger.info("*** Logging metrics ***")
         trainer.log_metrics("eval", metrics)
+        logger.info("*** Metrics logged ***")
+        logger.info("*** Saving metrics ***")
         trainer.save_metrics("eval", metrics)
+        logger.info("*** Metrics saved ***")
 
     # 14. Write Training Stats
+    logger.info("*** Writing training stats ***")
     kwargs = {
         "finetuned_from": model_args.model_name_or_path,
         "tasks": "automatic-speech-recognition",
@@ -693,11 +710,14 @@ def main():
         if model_args.model_index_name is not None:
             kwargs["model_name"] = model_args.model_index_name
 
-    logger.info("*** Pushing to hub ***")
     if training_args.push_to_hub:
+        logger.info("*** Pushing to hub ***")
         trainer.push_to_hub(**kwargs)
+        logger.info("*** Pushed to hub ***")
     else:
+        logger.info("*** Creating model card ***")
         trainer.create_model_card(**kwargs)
+        logger.info("*** Model card created ***")
     
     # Training complete notification
     logger.info("*** Sending notification ***")

test_run_nordic.sh

@@ -0,0 +1,39 @@
+python $1run_speech_recognition_seq2seq_streaming.py \
+	--model_name_or_path="openai/whisper-tiny" \
+	--dataset_name="mozilla-foundation/common_voice_11_0,mozilla-foundation/common_voice_11_0,mozilla-foundation/common_voice_11_0,babelbox/babelbox_voice,NbAiLab/NST,arpelarpe/nota,NbAiLab/NPSC" \
+	--dataset_config_name="sv-SE,da,nn-NO,,no-distant,,16k_mp3_nynorsk" \
+	--language="swedish" \
+	--train_split_name="train+validation,train+validation,train+validation,train,train+test, train,train+validation" \
+	--eval_split_name="test" \
+	--model_index_name="Whisper Tiny Swedish" \
+	--max_train_samples="64" \
+	--max_eval_samples="32" \
+	--max_steps="5000" \
+	--output_dir="./" \
+	--per_device_train_batch_size="8" \
+	--per_device_eval_batch_size="4" \
+	--logging_steps="25" \
+	--learning_rate="1e-5" \
+	--warmup_steps="500" \
+	--evaluation_strategy="steps" \
+	--eval_steps="1000" \
+	--save_strategy="steps" \
+	--save_steps="1000" \
+	--generation_max_length="225" \
+	--length_column_name="input_length" \
+	--max_duration_in_seconds="30" \
+	--text_column_name="sentence" \
+	--freeze_feature_encoder="False" \
+	--report_to="wandb" \
+	--metric_for_best_model="wer" \
+	--greater_is_better="False" \
+	--load_best_model_at_end \
+	--gradient_checkpointing \
+	--overwrite_output_dir \
+	--do_train \
+	--do_eval \
+	--predict_with_generate \
+	--do_normalize_eval \
+	--streaming \
+	--use_auth_token \
+	--push_to_hub