train code

run.sh

@@ -1,4 +1,39 @@
 
-python xla_spawn.py --num_cores=4 run_whisper.py 
+python run_speech_recognition_whisper_pere.py \
+	--model_name_or_path="openai/whisper-small" \
+	--output_dir="../whisper-testrun1" \
+	--overwrite_output_dir=True \
+	--language="Norwegian" \
+	--task="transcribe" \
+	--dataset_name="mozilla-foundation/common_voice_11_0" \
+	--dataset_config="nn-NO" \
+	--output_dir="./whisper-small-hi" \
+	--do_train=True \
+	--do_eval=True \
+	--audio_column_name="audio" \
+	--text_column_name="sentence" \
+	--per_device_train_batch_size=16 \
+       	--per_device_train_batch_size=16 \
+	--learning_rate=2e-5 \
+	--warmup_steps=500 \
+	--max_steps=5000 \
+	--gradient_checkpointing=True \
+	--gradient_accumulation_steps=1 \
+	--group_by_length=True \
+	--evaluation_strategy="steps" \
+	--save_steps=1000 \
+	--eval_steps=1000 \
+	--logging_steps=25 \
+	--fp16=True \
+	--save_steps=1000 \
+	--load_best_model_at_end=True \
+	--metric_for_best_model="wer" \
+	--greater_is_better=False \
+	--report_to="tensorboard" \
+	--predict_with_generate=True \
+	--generation_max_length=225 \
+	--print_training_arguments=True \
+	--push_to_hub=True
 
 
+	

run_speech_recognition_whisper_pere.py

@@ -22,27 +22,20 @@ import re
 import sys
 import warnings
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional,Union
+import evaluate
 
-import datasets
 import numpy as np
 import torch
+from pprint import pprint
 import evaluate
 from datasets import DatasetDict, load_dataset
+from datasets import Audio
 
-import transformers
 from transformers import (
-    AutoConfig,
-    AutoFeatureExtractor,
-    AutoModelForCTC,
-    AutoProcessor,
-    AutoTokenizer,
     HfArgumentParser,
-    Trainer,
     TrainingArguments,
-    Wav2Vec2Processor,
     set_seed,
-
     WhisperFeatureExtractor,
     WhisperTokenizer,
     WhisperForConditionalGeneration,
@@ -54,18 +47,9 @@ from transformers.trainer_utils import get_last_checkpoint, is_main_process
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
 
-# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
-# check_min_version("4.24.0.dev0")
-
-# require_version("datasets>=2.6.1", "To fix: pip install -r examples/pytorch/text-classification/requirements.txt")
-
-logger = logging.getLogger(__name__)
-
-
 def list_field(default=None, metadata=None):
     return field(default_factory=lambda: default, metadata=metadata)
 
-
 @dataclass
 class ModelArguments:
     """
@@ -243,6 +227,14 @@ class DataTrainingArguments:
         default="|",
         metadata={"help": "The word delimiter token for the tokenizer"},
     )
+    predict_with_generate: bool = field(
+        default=True,
+        metadata={"help": "Output tokens in addition to loss and digits for calculating metrics"},
+    )
+    generation_max_length: int = field(
+        default=225,
+        metadata={"help": "Maximum number of tokens generated"},
+    )
     phoneme_language: Optional[str] = field(
         default=None,
         metadata={
@@ -252,6 +244,12 @@ class DataTrainingArguments:
                     " input audio to a sequence of phoneme sequences."
         },
     )
+    print_training_arguments: bool = field(
+        default=True,
+        metadata={
+            "help": "Prints the training arguments. For debugging"
+        },
+    )
 
 
 @dataclass
@@ -278,20 +276,17 @@ class DataCollatorSpeechSeq2SeqWithPadding:
             labels = labels[:, 1:]
 
         batch["labels"] = labels
-
         return batch
 
 
-
-
 def main():    
     # See all possible arguments in src/transformers/training_args.py
     # or by passing the --help flag to this script.
     # We now keep distinct sets of args, for a cleaner separation of concerns.
     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
     model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+ 
 
-    
     # Metrics
     def compute_metrics(pred):
         pred_ids = pred.predictions
@@ -309,8 +304,6 @@ def main():
         return {"wer": wer}
 
     # Prepare dataset
-
-
     def prepare_dataset(batch):
         # load and resample audio data from 48 to 16kHz
         audio = batch["audio"]
@@ -323,45 +316,58 @@ def main():
         batch["labels"] = tokenizer(batch["sentence"]).input_ids
         return batch
     
-    def make_dataset(training_args, data_args):
-        seed = training_args.seed or 42
-        dataset = datasets.load_dataset(training_args.dataset_name, training_args.dataset_config_name, use_auth_token=data_args.use_auth_token)
-        return dataset
+    def print_training_arguments(model_args, data_args, training_args):
+        print("Starting with the following parameters:")
+        print("\n* Model arguments:")
+        pprint(vars(model_args),indent=2)
+        print("\n* Data arguments")
+        pprint(vars(data_args),indent=2)
+        print("\n* Training arguments")
+        pprint(vars(training_args),indent=2)
+    
+    # TODO - Might use this function later
+    # def make_dataset(training_args, data_args):
+    #     seed = training_args.seed or 42
+    #     dataset = datasets.load_dataset(training_args.dataset_name, training_args.dataset_config_name, use_auth_token=data_args.use_auth_token)
+    #     return dataset
+
+    # Print training arguments
+    if data_args.print_training_arguments:
+        print_training_arguments(model_args, data_args, training_args)
 
-    # PERE - SHOULD BE CHANGED TO STREAMING LATER
     # Load dataset
     speech_data = DatasetDict()
-
-    # The smallest dataset I found
     speech_data["train"] = load_dataset(
-        "mozilla-foundation/common_voice_11_0", "nn-NO", split="train", use_auth_token=True)
+        data_args.dataset_name, data_args.dataset_config_name, split="train", use_auth_token=True)
     speech_data["test"] = load_dataset(
-        "mozilla-foundation/common_voice_11_0", "nn-NO", split="test", use_auth_token=True)
+        data_args.dataset_name, data_args.dataset_config_name, split="test", use_auth_token=True)
 
-    #  PERE - REPLACE WITH THIS 
+    # TODO - Implement streaming and include this
     # speech_data = make_dataset(training_args, data_args)
 
-    # Rename columns
+
+    # Adapt dataset - Change column names and delete extra data
+    # Map columns
     if "audio" not in speech_data.column_names["train"]:
         speech_data = speech_data.rename_column(source, "audio")
 
     if "sentence" not in speech_data.column_names["train"]:
         speech_data = speech_data.rename_column(target, "sentence")
 
-    # Remove not needed columns - Not really sure if this is necessary
+    # Remove not needed columns 
     remove_list = [i for i in speech_data.column_names["train"]
                 if i not in ["audio", "sentence"]]
 
     speech_data = speech_data.remove_columns(remove_list)
 
-    # PERE - NEEDS TO BE PARAMETERIZED
+
     # Initialise
     feature_extractor = WhisperFeatureExtractor.from_pretrained(
-        "openai/whisper-small")
+        model_args.model_name_or_path)
     tokenizer = WhisperTokenizer.from_pretrained(
-        "openai/whisper-small", language="Norwegian", task="transcribe")
+        model_args.model_name_or_path, language=model_args.language, task=model_args.task)
     processor = WhisperProcessor.from_pretrained(
-        "openai/whisper-small", language="Norwegian", task="transcribe")
+        model_args.model_name_or_path, language=model_args.language, task=model_args.task)
     data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
 
     # Prepare data
@@ -369,6 +375,8 @@ def main():
     speech_data = speech_data.map(
         prepare_dataset, remove_columns=speech_data.column_names["train"], num_proc=1)
 
+
+
     # Metrics
     metric = evaluate.load("wer")
 
@@ -395,88 +403,47 @@ def main():
         if last_checkpoint is not None:
             checkpoint = last_checkpoint
         elif os.path.isdir(model_args.model_name_or_path):
-            checkpoint = model_args.model_name_or_path
-            # Initialise a Pretrained model
-            # We need to set use_cache=False here if we want to use gradient accumulation
-            # PERE - For the test this is set static
-            
-            model = WhisperForConditionalGeneration.from_pretrained(
-                "openai/whisper-small", use_cache=False)
-            
+            checkpoint = model_args.model_name_or_path         
         else:
             checkpoint = None
+            
+         # We need to set use_cache=False here if we want to use gradient accumulation
+        model = WhisperForConditionalGeneration.from_pretrained(
+                "openai/whisper-small", use_cache=False)
 
+        # Overriding generation arguments - no tokens are forced as decoder outputs (see [`forced_decoder_ids`](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.generation_utils.GenerationMixin.generate.forced_decoder_ids)), no tokens are suppressed during generation (see [`suppress_tokens`](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.generation_utils.GenerationMixin.generate.suppress_tokens)):
+        model.config.forced_decoder_ids = None
+        model.config.suppress_tokens = []
+        
+        # Set seed before initializing model.
+        set_seed(training_args.seed)
+
+
+        trainer = Seq2SeqTrainer(
+            args=training_args,
+            model=model,
+            train_dataset=speech_data["train"],
+            eval_dataset=speech_data["test"],
+            data_collator=data_collator,
+            compute_metrics=compute_metrics,
+            tokenizer=processor.feature_extractor,
+        )
+ 
         train_result = trainer.train(resume_from_checkpoint=checkpoint)
         trainer.save_model()
 
         metrics = train_result.metrics
-        max_train_samples = (
-            data_args.max_train_samples
-            if data_args.max_train_samples is not None
-            else len(vectorized_datasets["train"])
-        )
-        metrics["train_samples"] = min(max_train_samples, len(vectorized_datasets["train"]))
-
         trainer.log_metrics("train", metrics)
         trainer.save_metrics("train", metrics)
         trainer.save_state()
-
-    # Overriding generation arguments - no tokens are forced as decoder outputs (see [`forced_decoder_ids`](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.generation_utils.GenerationMixin.generate.forced_decoder_ids)), no tokens are suppressed during generation (see [`suppress_tokens`](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.generation_utils.GenerationMixin.generate.suppress_tokens)):
-    model.config.forced_decoder_ids = None
-    model.config.suppress_tokens = []
-
-    # Set seed before initializing model.
-    set_seed(training_args.seed)
-
-    # Training arguments
-    training_args = Seq2SeqTrainingArguments(
-        output_dir="../whisper-testrun1",  # change to a repo name of your choice
-        per_device_train_batch_size=16,
-        gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
-        learning_rate=2e-5,
-        warmup_steps=500,
-        max_steps=5000,  # Changed from 4000
-        gradient_checkpointing=True,
-        group_by_length=True,
-        evaluation_strategy="steps",
-        per_device_eval_batch_size=8,
-        predict_with_generate=True,
-        generation_max_length=225,
-        save_steps=500,
-        eval_steps=500,
-        logging_steps=25,
-        report_to=["tensorboard"],
-        load_best_model_at_end=True,
-        metric_for_best_model="wer",
-        greater_is_better=False,
-        push_to_hub=True,
-    )
-
-    trainer = Seq2SeqTrainer(
-        args=training_args,
-        model=model,
-        train_dataset=speech_data["train"],
-        eval_dataset=speech_data["test"],
-        data_collator=data_collator,
-        compute_metrics=compute_metrics,
-        tokenizer=processor.feature_extractor,
-    )
-
- 
-    # Initialize Trainer
-    trainer = Seq2SeqTrainer(
-        model=model,
-        data_collator=data_collator,
-        args=training_args,
-        compute_metrics=compute_metrics,
-        train_dataset=vectorized_datasets["train"] if training_args.do_train else None,
-        eval_dataset=vectorized_datasets["validation"] if training_args.do_eval else None,
-        tokenizer=feature_extractor,
-    )
-
-    # 8. Finally, we can start training
-
+        
+        if training_args.push_to_hub:
+            trainer.push_to_hub(**kwargs)
+        else:
+            trainer.create_model_card(**kwargs)
  
+    # TODO - Look closer into the evaluation and the model card writing.
+    
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -500,14 +467,7 @@ def main():
         "dataset": f"{data_args.dataset_name.upper()} - {config_name.upper()}",
         "language": model_args.language,
     }
-    if "common_voice" in data_args.dataset_name:
-        kwargs["language"] = config_name
-
-    if training_args.push_to_hub:
-        trainer.push_to_hub(**kwargs)
-    else:
-        trainer.create_model_card(**kwargs)
-
+    
     return results
 
 
@@ -517,7 +477,5 @@ def _mp_fn(index):
     print("The XLA is initiated")
     main()
 
-
-
 if __name__ == "__main__":
     main()

run_xla.sh

@@ -0,0 +1,4 @@
+
+python xla_spawn.py --num_cores=4 run_whisper.py 
+
+