allow for multiple datasets from hf in run

run_speech_recognition_ctc.py

@@ -30,7 +30,7 @@ import datasets
 import numpy as np
 import torch
 import wandb
-from datasets import DatasetDict, load_dataset, load_metric
+from datasets import DatasetDict, concatenate_datasets, load_dataset, load_metric
 
 import transformers
 from transformers import (
@@ -140,21 +140,33 @@ class DataTrainingArguments:
     """
 
     dataset_name: str = field(
-        metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+        metadata={
+            "help": "The name of the dataset to use (via the datasets library)."
+            " To use multiple datasets, specify them separated by a comma."
+            " e.g.: 'mozilla-foundation/common_voice_7_0,marinone94/nst_sv'"
+        }
     )
     dataset_config_name: str = field(
-        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+        default=None, metadata={
+            "help": "The configuration name of the dataset to use (via the datasets library)."
+            " To use multiple datasets, specify them separated by a comma."
+            " e.g.: 'sv-SE,sv'"    
+        }
     )
     train_split_name: str = field(
         default="train+validation",
         metadata={
-            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
+            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train+validation'"
+            " To use multiple datasets, specify them separated by a comma."
+            " e.g.: 'train+validation,all'"
         },
     )
     eval_split_name: str = field(
         default="test",
         metadata={
-            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
+            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'test'"
+            " To use multiple datasets, specify them separated by a comma."
+            " e.g.: 'test,None'"
         },
     )
     audio_column_name: str = field(
@@ -407,12 +419,36 @@ def main():
     raw_datasets = DatasetDict()
 
     if training_args.do_train:
-        raw_datasets["train"] = load_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
-            split=data_args.train_split_name,
-            use_auth_token=data_args.use_auth_token,
-        )
+
+        # Multiple datasets might need to be loaded from HF
+        # It assumes they all follow the common voice format
+        for (dataset_name, dataset_config_name, train_split_name) in zip(
+            data_args.dataset_name.split(","),
+            data_args.dataset_config_name.split(","),
+            data_args.train_split_name.split(","),
+        ):
+
+                
+            if train_split_name != "None":
+                if "train" not in raw_datasets:
+                    raw_datasets["train"] = load_dataset(
+                        dataset_name,
+                        dataset_config_name,
+                        split=train_split_name,
+                        use_auth_token=data_args.use_auth_token,
+                    )
+                else:
+                    raw_datasets["train"] = concatenate_datasets(
+                        [
+                            raw_datasets["train"],
+                            load_dataset(
+                                dataset_name,
+                                dataset_config_name,
+                                split=train_split_name,
+                                use_auth_token=data_args.use_auth_token,
+                            )
+                        ]
+                    )
 
         if data_args.audio_column_name not in raw_datasets["train"].column_names:
             raise ValueError(
@@ -432,12 +468,34 @@ def main():
             raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
 
     if training_args.do_eval:
-        raw_datasets["eval"] = load_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
-            split=data_args.eval_split_name,
-            use_auth_token=data_args.use_auth_token,
-        )
+        # Multiple datasets might need to be loaded from HF
+        # It assumes they all follow the common voice format
+        for (dataset_name, dataset_config_name, eval_split_name) in zip(
+            data_args.dataset_name.split(","),
+            data_args.dataset_config_name.split(","),
+            data_args.eval_split_name.split(","),
+        ):
+
+            if train_split_name != "None":
+                if "eval" not in raw_datasets:
+                    raw_datasets["eval"] = load_dataset(
+                        dataset_name,
+                        dataset_config_name,
+                        split=eval_split_name,
+                        use_auth_token=data_args.use_auth_token,
+                    )
+                else:
+                    raw_datasets["eval"] = concatenate_datasets(
+                        [
+                            raw_datasets["eval"],
+                            load_dataset(
+                                dataset_name,
+                                dataset_config_name,
+                                split=train_split_name,
+                                use_auth_token=data_args.use_auth_token,
+                            )
+                        ]
+                    )
 
         if data_args.max_eval_samples is not None:
             raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))