Finetuning bugs

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ repos:
     hooks:
       - id: python-use-type-annotations
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: end-of-file-fixer
       - id: trailing-whitespace

config/datasets/common_voice_17.yaml

@@ -5,4 +5,3 @@ common_voice_17:
   text_column: sentence
   audio_column: audio
   filter_dataset: false
-  process_dataset: true

config/datasets/common_voice_9.yaml

@@ -5,4 +5,3 @@ common_voice_9:
   text_column: sentence
   audio_column: audio
   filter_dataset: true
-  process_dataset: true

config/datasets/coral.yaml

@@ -5,4 +5,3 @@ coral:
   text_column: text
   audio_column: audio
   filter_dataset: false
-  process_dataset: true

config/datasets/fleurs.yaml

@@ -5,4 +5,3 @@ fleurs:
   text_column: raw_transcription
   audio_column: audio
   filter_dataset: true
-  process_dataset: true

config/datasets/ftspeech.yaml

@@ -5,4 +5,3 @@ ftspeech:
   text_column: sentence
   audio_column: audio
   filter_dataset: true
-  process_dataset: true

config/datasets/nota.yaml

@@ -5,4 +5,3 @@ nota:
   text_column: text
   audio_column: audio
   filter_dataset: true
-  process_dataset: true

config/datasets/nst.yaml

@@ -5,4 +5,3 @@ nst:
   text_column: text
   audio_column: audio
   filter_dataset: true
-  process_dataset: true

config/datasets/test_dataset.yaml

@@ -5,4 +5,3 @@ test_dataset:
   text_column: sentence
   audio_column: audio
   filter_dataset: true
-  process_dataset: true

makefile

@@ -53,6 +53,7 @@ install-pipx:
 	@if [ "$(shell which pipx)" = "" ]; then \
 		uname=$$(uname); \
 			case $${uname} in \
+				(*Linux*) installCmd='sudo apt install pipx'; ;; \
 				(*Darwin*) installCmd='brew install pipx'; ;; \
 				(*CYGWIN*) installCmd='py -3 -m pip install --upgrade --user pipx'; ;; \
 				(*) installCmd='python3 -m pip install --upgrade --user pipx'; ;; \
@@ -118,6 +119,8 @@ type-check:  ## Run type checking
 		--show-error-codes \
 		--check-untyped-defs
 
+check: lint format type-check  ## Check the code
+
 roest-315m:  ## Train the Røst-315M model
 	@accelerate launch \
 		--use-deepspeed \

src/coral/data.py

@@ -171,35 +171,25 @@ def load_data_for_finetuning(
                 num_proc=config.dataset_num_workers,
             )
 
-        if dataset_config.process_dataset:
-            ds = ds.remove_columns(
-                column_names=[
-                    column
-                    for column in ds.column_names or list()
-                    if column not in ["audio", "text"]
-                ]
-            ).shuffle(seed=config.seed)
-            ds = process_dataset(
-                dataset=ds,
-                clean_text=config.model.clean_text,
-                lower_case=config.model.lower_case,
-                characters_to_keep=config.characters_to_keep,
-                text_column="text",
-                audio_column="audio",
-                convert_numerals=False,
-                remove_input_dataset_columns=True,
-                cast_to_sampling_rate=config.model.sampling_rate,
-                processor=processor,
-                num_proc=config.dataset_num_workers,
-            )
+        ds = ds.remove_columns(
+            column_names=[
+                column
+                for column in ds.column_names or list()
+                if column not in ["audio", "text"]
+            ]
+        ).shuffle(seed=config.seed)
+
+        ds = ds.cast_column(
+            column="audio", feature=Audio(sampling_rate=config.model.sampling_rate)
+        )
 
         all_datasets.append(ds)
 
     assert len(all_datasets) > 0, "No datasets were loaded"
 
     if len(all_datasets) > 1:
         if is_main_process:
-            logger.info("Interleaving datasets")
+            logger.info("Interleaving datasets...")
             if config.dataset_probabilities is None and len(all_datasets) > 1:
                 logger.warning(
                     "No dataset probabilities were specified for the training split. "
@@ -228,6 +218,19 @@ def load_data_for_finetuning(
     else:
         train = all_datasets[0]
 
+    train = process_dataset(
+        dataset=train,
+        clean_text=config.model.clean_text,
+        lower_case=config.model.lower_case,
+        characters_to_keep=config.characters_to_keep,
+        text_column="text",
+        audio_column="audio",
+        convert_numerals=False,
+        remove_input_dataset_columns=True,
+        processor=processor,
+        num_proc=config.dataset_num_workers,
+    )
+
     data_dict = dict(train=train)
     dataset = IterableDatasetDict(data_dict)
 
@@ -256,6 +259,10 @@ def load_data_for_finetuning(
     if config.evaluation_dataset.audio_column != "audio":
         val = val.rename_column(config.evaluation_dataset.audio_column, "audio")
 
+    val = val.cast_column(
+        column="audio", feature=Audio(sampling_rate=config.model.sampling_rate)
+    )
+
     val = process_dataset(
         dataset=val,
         clean_text=config.model.clean_text,
@@ -265,7 +272,6 @@ def load_data_for_finetuning(
         audio_column="audio",
         convert_numerals=False,
         remove_input_dataset_columns=True,
-        cast_to_sampling_rate=config.model.sampling_rate,
         processor=processor,
         num_proc=config.dataset_num_workers,
     )
@@ -328,6 +334,9 @@ def load_dataset_for_evaluation(config: DictConfig) -> Dataset:
         max_seconds_per_example=config.max_seconds_per_example,
         is_main_process=is_main_process,
     )
+    dataset = dataset.cast_column(
+        column=config.audio_column, feature=Audio(sampling_rate=config.sampling_rate)
+    )
     dataset = process_dataset(
         dataset=dataset,
         clean_text=config.clean_text,
@@ -336,7 +345,6 @@ def load_dataset_for_evaluation(config: DictConfig) -> Dataset:
         text_column=config.text_column,
         audio_column=config.audio_column,
         remove_input_dataset_columns=False,
-        cast_to_sampling_rate=config.sampling_rate,
         convert_numerals=True,
     )
 
@@ -450,7 +458,6 @@ def process_dataset(
     audio_column: str | None,
     convert_numerals: bool,
     num_proc: int | None = None,
-    cast_to_sampling_rate: int | None = None,
     processor: Callable | None = None,
 ) -> Data:
     """Process the dataset.
@@ -479,21 +486,13 @@ def process_dataset(
         num_proc (optional):
             The number of processes to use for processing the dataset. If `None`, then
             no multiprocessing is used. Defaults to `None`.
-        cast_to_sampling_rate (optional):
-            The sampling rate to cast the audio to. If `None`, then the audio is not
-            cast. Defaults to `None`.
         processor (optional):
             The processor to use for processing the audio and transcriptions. If `None`,
             then the processor is not used. Defaults to `None`.
 
     Returns:
         The cleaned dataset.
     """
-    if audio_column is not None:
-        dataset = dataset.cast_column(
-            column=audio_column, feature=Audio(sampling_rate=cast_to_sampling_rate)
-        )
-
     if isinstance(dataset, Dataset) or isinstance(dataset, IterableDataset):
         column_names = dataset.column_names
     elif isinstance(dataset, DatasetDict) or isinstance(dataset, IterableDatasetDict):

src/coral/data_collators.py

@@ -70,7 +70,7 @@ def torch_call(self, features: list[dict]) -> BatchFeature:
             labels=label_features,
             padding=self.padding,
             return_tensors=self.return_tensors,
-            max_length=512,
+            max_length=min(self.processor.tokenizer.model_max_length, 512),
         )
 
         # Replace padding with -100 to ignore loss correctly
@@ -147,16 +147,16 @@ def torch_call(self, features: list[dict]) -> BatchFeature:
             label_features,
             padding=self.padding,
             return_tensors=self.return_tensors,
-            max_length=512,
+            max_length=min(self.processor.tokenizer.model_max_length, 512),
         )
 
-        # replace padding with -100 to ignore loss correctly
+        # Replace padding with -100 to ignore loss correctly
         labels = labels_batch["input_ids"].masked_fill(
             labels_batch.attention_mask.ne(1), -100
         )
 
-        # if bos token is appended in previous tokenization step,
-        # cut bos token here as it's append later anyways
+        # If bos token is appended in previous tokenization step, cut BOS token here as
+        # it's appended later anyway
         if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
             labels = labels[:, 1:]
 

src/coral/validation.py

@@ -5,7 +5,7 @@
 from typing import TypeVar
 
 import torch
-from datasets import Dataset, DatasetDict
+from datasets import Audio, Dataset, DatasetDict
 from transformers import AutomaticSpeechRecognitionPipeline, pipeline
 
 from .compute_metrics import compute_metrics_of_dataset_using_pipeline
@@ -60,6 +60,10 @@ def add_validations(
     if input_is_single_split:
         dataset = DatasetDict(dict(train=dataset))
 
+    dataset = dataset.cast_column(
+        column=audio_column, feature=Audio(sampling_rate=sampling_rate)
+    )
+
     processed_dataset = process_dataset(
         dataset=dataset,
         clean_text=clean_text,
@@ -69,7 +73,6 @@ def add_validations(
         convert_numerals=False,
         remove_input_dataset_columns=True,
         lower_case=lower_case,
-        cast_to_sampling_rate=sampling_rate,
     )
 
     logger.info(f"Loading the {model_id!r} ASR model...")

src/coral/wav2vec2.py

@@ -209,6 +209,7 @@ def load_training_arguments(self) -> TrainingArguments:
             use_cpu=hasattr(sys, "_called_from_test"),
             dataloader_num_workers=self.config.dataloader_num_workers,
             ddp_find_unused_parameters=False,
+            dispatch_batches=False,
         )
         return args
 

src/coral/whisper.py

@@ -12,6 +12,8 @@
 from omegaconf import DictConfig
 from torch.backends.mps import is_available as mps_is_available
 from transformers import (
+    AutoConfig,
+    AutoModelForSpeechSeq2Seq,
     EvalPrediction,
     Seq2SeqTrainer,
     Seq2SeqTrainingArguments,
@@ -51,12 +53,20 @@ def load_processor(self) -> WhisperProcessor:
         )
         assert isinstance(processor_or_tup, WhisperProcessor)
         self.processor = processor_or_tup
+
+        # Whisper tokenizers are misconfigured with a max_length that is too high, but
+        # the correct max_length is stored in the model config, so we'll update it here.
+        hf_config = AutoConfig.from_pretrained(self.config.model.pretrained_model_id)
+        self.processor.tokenizer.model_max_length = min(
+            self.processor.tokenizer.model_max_length, hf_config.max_length
+        )
+
         return self.processor
 
     def load_model(self) -> WhisperForConditionalGeneration:
         """Return the model for the setup."""
         with transformers_output_ignored():
-            model = WhisperForConditionalGeneration.from_pretrained(
+            model = AutoModelForSpeechSeq2Seq.from_pretrained(
                 self.config.model.pretrained_model_id,
                 dropout=self.config.model.dropout,
                 activation_dropout=self.config.model.activation_dropout,
@@ -179,6 +189,7 @@ def load_training_arguments(self) -> TrainingArguments:
             use_cpu=hasattr(sys, "_called_from_test"),
             dataloader_num_workers=self.config.dataloader_num_workers,
             ddp_find_unused_parameters=False,
+            dispatch_batches=False,
         )
         return args
 

src/scripts/finetune_asr_model.py

@@ -52,13 +52,16 @@ def main(config: DictConfig) -> None:
                     "training"
                 )
             config.model.layerdrop = 0.0
-        if config.padding != "max_length":
-            if is_main_process:
-                logger.info(
-                    "Forcing `padding` to be 'max_length' as this is required in a "
-                    "multi-GPU training"
-                )
-            config.padding = "max_length"
+
+        # TODO: This doesn't seem to be changed anymore, but keeping it here for some
+        #       time in case we need to re-enable it.
+        # if config.padding != "max_length":
+        #    if is_main_process:
+        #        logger.info(
+        #            "Forcing `padding` to be 'max_length' as this is required in a "
+        #            "multi-GPU training"
+        #        )
+        #    config.padding = "max_length"
 
     elif torch.cuda.device_count() > 1:
         if is_main_process: