allow multiple languages and datasets

run_speech_recognition_seq2seq_streaming.py

@@ -49,6 +49,7 @@ from transformers import (
     set_seed,
 )
 from transformers.models.whisper.english_normalizer import BasicTextNormalizer
+from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE
 from transformers.trainer_pt_utils import IterableDatasetShard
 from transformers.trainer_utils import get_last_checkpoint, is_main_process
 from transformers.utils import check_min_version, send_example_telemetry
@@ -61,6 +62,9 @@ require_version("datasets>=1.18.2", "To fix: pip install -r examples/pytorch/spe
 
 logger = logging.getLogger(__name__)
 
+SENDING_NOTIFICATION = "*** Sending notification to email ***"
+RECIPIENT_ADDRESS = "[email protected]"
+
 wandb_token = os.environ.get("WANDB_TOKEN", "None")
 hf_token = os.environ.get("HF_TOKEN", None)
 if (hf_token is None or wandb_token == "None") and os.path.exists("./creds.txt"):
@@ -160,10 +164,16 @@ class DataTrainingArguments:
     Arguments pertaining to what data we are going to input our model for training and eval.
     """
 
-    dataset_name: str = field(
+    dataset_train_name: str = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_train_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    dataset_eval_name: str = field(
         default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
     )
-    dataset_config_name: Optional[str] = field(
+    dataset_eval_config_name: Optional[str] = field(
         default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
     )
     text_column: Optional[str] = field(
@@ -232,7 +242,16 @@ class DataTrainingArguments:
         default=True,
         metadata={"help": "Whether to normalise the references and predictions in the eval WER calculation."},
     )
-    language: str = field(
+    language_train: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "Language for multilingual fine-tuning. This argument should be set for multilingual fine-tuning "
+                "only. For English speech recognition, it should be set to `None`."
+            )
+        },
+    )
+    language_eval: str = field(
         default=None,
         metadata={
             "help": (
@@ -273,6 +292,7 @@ class DataCollatorSpeechSeq2SeqWithPadding:
 
     processor: Any
     decoder_start_token_id: int
+    task_id: int
 
     def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
         # split inputs and labels since they have to be of different lengths and need
@@ -280,6 +300,7 @@ class DataCollatorSpeechSeq2SeqWithPadding:
         model_input_name = self.processor.model_input_names[0]
         input_features = [{model_input_name: feature[model_input_name]} for feature in features]
         label_features = [{"input_ids": feature["labels"]} for feature in features]
+        lang_features = [f"<|{TO_LANGUAGE_CODE[feature['language']]}|>" for feature in features]
 
         batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
 
@@ -292,6 +313,15 @@ class DataCollatorSpeechSeq2SeqWithPadding:
         # cut bos token here as it's append later anyways
         if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
             labels = labels[:, 1:]
+        lang_token_ids = self.processor.tokenizer(lang_features).input_ids
+        # Replace language and task if they are in the beginning, otherwise add them
+        if (labels[:, 1] == self.task_id).all().cpu().item():
+            labels[:, 0] = lang_token_ids
+            labels[:, 1] = torch.full_like(labels[:, 1], self.task_id)
+        else:
+            # convert task id to tensor of labels dim to concatenate
+            task_id = torch.full_like(labels[:, 0], self.task_id)
+            labels = torch.cat((lang_token_ids, task_id, labels), dim=1)
 
         batch["labels"] = labels
 
@@ -316,7 +346,7 @@ def notify_me(recipient, message=None):
         from email.mime.text import MIMEText
 
         msg = MIMEText(message)
-        msg["Subject"] = "Training is finished!"
+        msg["Subject"] = "Training updates..."
         msg["From"] = "[email protected]"
         msg["To"] = recipient
 
@@ -334,16 +364,26 @@ def load_maybe_streaming_dataset(dataset_names, dataset_config_names, split="tra
     each split is loaded individually and then splits combined by taking alternating examples from
     each (interleaving).
     """
+    column_names = None
+    if "column_names" in kwargs:
+        column_names = kwargs.pop("column_names").split(",")
+
     if "," in dataset_names or "+" in split:
         # load multiple splits separated by the `+` symbol with streaming mode
         dataset_splits = []
-        for dataset_name, dataset_config_name, split_names in zip(
-            dataset_names.split(","), dataset_config_names.split(","), split.split(",")
+        for dataset_name, dataset_config_name, split_names, lang in zip(
+            dataset_names.split(","), dataset_config_names.split(","), split.split(","), kwargs.pop("language").split(",")
         ):
             for split_name in split_names.split("+"):
-                dataset = load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=streaming, **kwargs)
+                dataset = load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=streaming, **kwargs)                    
+                raw_datasets_features = list(next(iter(dataset.values())).features.keys())
+                if column_names[0] not in raw_datasets_features:
+                    if len(column_names) == 1 or column_names[1] not in raw_datasets_features:
+                        raise ValueError("Column name not found in dataset.")
+                    dataset = dataset.rename_columns(column_names[1], column_names[0])
+                dataset["language"] = lang
                 dataset_splits.append(dataset)
-    
+
         # interleave multiple splits to form one dataset
         interleaved_dataset = interleave_datasets(dataset_splits)
         return interleaved_dataset
@@ -426,20 +466,23 @@ def main():
 
     if training_args.do_train:
         raw_datasets["train"] = load_maybe_streaming_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
+            data_args.dataset_train_name,
+            data_args.dataset_train_config_name,
             split=data_args.train_split_name,
             use_auth_token=hf_token if model_args.use_auth_token else None,
             streaming=data_args.streaming,
+            column_names=data_args.text_column_name,
+            language=data_args.language_train
         )
 
     if training_args.do_eval:
         raw_datasets["eval"] = load_maybe_streaming_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
+            data_args.dataset_eval_name,
+            data_args.dataset_eval_config_name,
             split=data_args.eval_split_name,
             use_auth_token=hf_token if model_args.use_auth_token else None,
             streaming=data_args.streaming,
+            language=data_args.language_eval
         )
 
     raw_datasets_features = list(next(iter(raw_datasets.values())).features.keys())
@@ -451,6 +494,7 @@ def main():
             f"{', '.join(raw_datasets_features)}."
         )
 
+    data_args.text_column_name = data_args.text_column_name.split(",")[0]
     if data_args.text_column_name not in raw_datasets_features:
         raise ValueError(
             f"--text_column_name {data_args.text_column_name} not found in dataset '{data_args.dataset_name}'. "
@@ -504,9 +548,13 @@ def main():
     if model_args.freeze_encoder:
         model.freeze_encoder()
 
-    if data_args.language is not None:
+    if data_args.language is not None and len(data_args.language.split(",")) == 1:
         # We only need to set the task id when the language is specified (i.e. in a multilingual setting)
+        # If more than a langugae is specified, it will be specified in the data collator
         tokenizer.set_prefix_tokens(language=data_args.language, task=data_args.task)
+    elif data_args.language is not None and len(data_args.language.split(",")) > 1:
+        # make sure language and task are not stored in the model config
+        model.config.forced_decoder_ids = None
 
     # 6. Resample speech dataset if necessary
     logger.info("*** Resample dataset ***")
@@ -558,6 +606,7 @@ def main():
         return batch
 
     with training_args.main_process_first(desc="dataset map pre-processing"):
+        raw_datasets_features.remove("language")
         vectorized_datasets = raw_datasets.map(
             prepare_dataset,
             remove_columns=raw_datasets_features,
@@ -617,9 +666,14 @@ def main():
     processor = AutoProcessor.from_pretrained(training_args.output_dir)
 
     # 10. Define data collator
+    task_token = data_args.task
+    if not task_token.startswith('<|'):
+        task_token = f'<{task_token}>'
+    task_id = tokenizer(task_token).input_ids[0]
     data_collator = DataCollatorSpeechSeq2SeqWithPadding(
         processor=processor,
         decoder_start_token_id=model.config.decoder_start_token_id,
+        task_id=task_id
     )
 
     # 11. Configure Trainer
@@ -716,20 +770,24 @@ def main():
         if model_args.model_index_name is not None:
             kwargs["model_name"] = model_args.model_index_name
 
+    # Training complete notification
+    logger.info(SENDING_NOTIFICATION)
+    notify_me(recipient=RECIPIENT_ADDRESS, message=json.dumps(kwargs, indent=4))
+    logger.info("*** Training complete!!! ***")
+
+
     if training_args.push_to_hub:
         logger.info("*** Pushing to hub ***")
         trainer.push_to_hub(**kwargs)
         logger.info("*** Pushed to hub ***")
+        logger.info(SENDING_NOTIFICATION)
+        notify_me(recipient=RECIPIENT_ADDRESS, message="Model 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 ***")
-    notify_me(recipient="[email protected]", message=json.dumps(kwargs, indent=4))
-
-    logger.info("*** Training complete!!! ***")
+        logger.info(SENDING_NOTIFICATION)
+        notify_me(recipient=RECIPIENT_ADDRESS, message="Model card created")
 
     return results
 

test_run_nordic.sh

@@ -1,9 +1,12 @@
 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" \
+	--dataset_train_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,google/fleurs,google/fleurs,google/fleurs" \
+	--dataset_train_config_name="sv-SE,da,nn-NO,,no-distant,,16k_mp3_nynorsk,sv_se,da_dk,nb_no" \
+	--language_train="swedish,danish,norwegian,swedish,norwegian,danish,norwegian,swedish,danish,norwegian" \
+	--train_split_name="train+validation,train+validation,train+validation,train,train+test,train,train+validation,train+validation,train+validation,train+validation" \
+	--dataset_eval_name="mozilla-foundation/common_voice_11_0,mozilla-foundation/common_voice_11_0,mozilla-foundation/common_voice_11_0" \
+	--dataset_eval_config_name="sv-SE,da,nn-NO" \
+	--language_eval="swedish,danish,norwegian" \
 	--eval_split_name="test" \
 	--model_index_name="Whisper Tiny Swedish" \
 	--max_train_samples="64" \
@@ -22,7 +25,7 @@ python $1run_speech_recognition_seq2seq_streaming.py \
 	--generation_max_length="225" \
 	--length_column_name="input_length" \
 	--max_duration_in_seconds="30" \
-	--text_column_name="sentence" \
+	--text_column_name="sentence,text" \
 	--freeze_feature_encoder="False" \
 	--report_to="wandb" \
 	--metric_for_best_model="wer" \