Add training scripts

create_student_model.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Initialise a student Whisper model from a pre-trained teacher model for
+teacher-student distillation.
+"""
+
+import argparse
+import copy
+import logging
+
+import numpy as np
+import torch
+from transformers import GenerationConfig, WhisperForConditionalGeneration, WhisperProcessor
+
+
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Initialise a student Whisper model from a teacher model, copying the relevant layer weights and adjusting the processor as necessary."
+    )
+    parser.add_argument(
+        "--teacher_checkpoint",
+        type=str,
+        required=True,
+        help="The HF Hub ID of the teacher checkpoint.",
+    )
+    parser.add_argument(
+        "--subfolder",
+        type=str,
+        default="",
+        help="In case the relevant teacher weights are located inside a subfolder of the model repo on huggingface.co, you "
+        "can specify the folder name here.",
+    )
+    parser.add_argument(
+        "--encoder_layers",
+        type=int,
+        default=None,
+        help="Number of encoder layers to use in the student model. Defaults to all layers from the teacher.",
+    )
+    parser.add_argument(
+        "--decoder_layers",
+        type=int,
+        default=2,
+        help="Number of decoder layers to use in the student model. Defaults to 2 layers.",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        required=True,
+        help="Where to save the student weights and processor.",
+    )
+    parser.add_argument(
+        "--push_to_hub",
+        type=bool,
+        required=False,
+        default=False,
+        help="Whether to push the student weights and processor to the Hub.",
+    )
+    parser.add_argument(
+        "--cache_dir",
+        type=str,
+        default=None,
+        help="Where to store the pretrained models downloaded from huggingface.co",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def init_student_model_from_teacher(
+    teacher_checkpoint,
+    encoder_layers=None,
+    decoder_layers=2,
+    save_dir=None,
+    push_to_hub=None,
+    cache_dir=None,
+    subfolder="",
+):
+    teacher_model = WhisperForConditionalGeneration.from_pretrained(
+        teacher_checkpoint,
+        cache_dir=cache_dir,
+        subfolder=subfolder,
+        low_cpu_mem_usage=True,
+    )
+    processor = WhisperProcessor.from_pretrained(teacher_checkpoint)
+    generation_config = GenerationConfig.from_pretrained(teacher_checkpoint)
+    generation_config.forced_decoder_ids = None
+
+    teacher_config = teacher_model.config
+    teacher_encoder_layers = teacher_config.encoder_layers
+    teacher_decoder_layers = teacher_config.decoder_layers
+
+    student_config = copy.deepcopy(teacher_config)
+    student_config.update(
+        {
+            "encoder_layers": encoder_layers if encoder_layers is not None else teacher_encoder_layers,
+            "decoder_layers": decoder_layers,
+        }
+    )
+
+    encoder_mapping = np.linspace(0, teacher_encoder_layers - 1, student_config.encoder_layers, dtype=int)
+    encoder_mapping[-1] = teacher_encoder_layers - 1
+
+    encoder_map = {}
+    for student_layer, teacher_layer in enumerate(encoder_mapping):
+        encoder_map[teacher_layer] = student_layer
+
+    decoder_mapping = np.linspace(0, teacher_decoder_layers - 1, student_config.decoder_layers, dtype=int)
+    decoder_mapping[-1] = teacher_decoder_layers - 1
+
+    decoder_map = {}
+    for student_layer, teacher_layer in enumerate(decoder_mapping):
+        decoder_map[teacher_layer] = student_layer
+
+    # init the student params from the teacher model
+    student_model = WhisperForConditionalGeneration(student_config)
+    missing_keys, unexpected_keys = student_model.load_state_dict(teacher_model.state_dict(), strict=False)
+    if len(missing_keys) > 0:
+        raise RuntimeError(
+            "Error(s) in loading state_dict for WhisperForConditionalGeneration. \n"
+            f"Missing key(s) in state_dict: {missing_keys}"
+        )
+    if decoder_layers == teacher_decoder_layers:
+        decoder_keys = [key for key in unexpected_keys if "model.decoder.layers" in key]
+        if len(decoder_keys) > 0:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for WhisperForConditionalGeneration. \n"
+                f"Unexpected key(s) in state_dict: {decoder_keys}"
+            )
+    if encoder_layers == teacher_encoder_layers:
+        encoder_keys = [key for key in unexpected_keys if "model.encoder.layers" in key]
+        if len(encoder_keys) > 0:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for WhisperForConditionalGeneration. \n"
+                f"Unexpected key(s) in state_dict: {encoder_keys}"
+            )
+
+    for layer in range(teacher_decoder_layers):
+        if layer in decoder_map:
+            # re-introduce pre-defined layers from the teacher
+            student_model.model.decoder.layers[decoder_map[layer]].load_state_dict(
+                teacher_model.model.decoder.layers[layer].state_dict()
+            )
+
+    if encoder_layers is not None:
+        for layer in range(teacher_encoder_layers):
+            if layer in encoder_map:
+                # re-introduce pre-defined layers from the teacher
+                student_model.model.encoder.layers[encoder_map[layer]].load_state_dict(
+                    teacher_model.model.encoder.layers[layer].state_dict()
+                )
+
+    # remove the teacher params and model
+    del teacher_model
+
+    # save the converted weights and model
+    if save_dir is not None:
+        student_model.save_pretrained(save_dir)
+        # we also need to correctly save the processor and generation config
+        processor.save_pretrained(save_dir)
+        generation_config.save_pretrained(save_dir)
+
+    # check we can do a forward pass with the saved model - first load the weights and processor
+    logger.info("Checking we can load the saved model...")
+    student_model = WhisperForConditionalGeneration.from_pretrained(
+        save_dir,
+        low_cpu_mem_usage=True,
+    )
+    processor = WhisperProcessor.from_pretrained(save_dir)
+
+    # define some random inputs
+    input_features = processor(np.ones(16000), sampling_rate=16000, return_tensors="pt").input_features
+    decoder_start_token_id = student_model.config.decoder_start_token_id
+    decoder_input_ids = torch.ones((input_features.shape[0], 1), dtype=torch.long) * decoder_start_token_id
+
+    # do a forward pass - outputs will be gibberish for the initialised model so we can't check them
+    # but we make can sure the model runs as expected
+    logger.info("Checking we can run the converted model forward...")
+    _ = student_model(input_features, decoder_input_ids=decoder_input_ids).logits
+    logger.info("Conversion successful!")
+
+    if push_to_hub:
+        student_model.push_to_hub(save_dir)
+        processor.push_to_hub(save_dir)
+        generation_config.push_to_hub(save_dir)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    init_student_model_from_teacher(
+        teacher_checkpoint=args.teacher_checkpoint,
+        encoder_layers=args.encoder_layers,
+        decoder_layers=args.decoder_layers,
+        save_dir=args.save_dir,
+        push_to_hub=args.push_to_hub,
+        cache_dir=args.cache_dir,
+        subfolder=args.subfolder,
+    )

run.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+
+accelerate launch run_distillation.py \
+  --model_name_or_path "./distil-large-v3-init" \
+  --teacher_model_name_or_path "openai/whisper-large-v3" \
+  --train_dataset_name "../common_voice_16_1_de_pseudo_labelled" \
+  --train_split_name "train" \
+  --text_column_name "sentence" \
+  --eval_dataset_name "../common_voice_16_1_de_pseudo_labelled" \
+  --eval_split_name "validation" \
+  --eval_text_column_name "sentence" \
+  --eval_steps 5000 \
+  --save_steps 5000 \
+  --warmup_steps 500 \
+  --learning_rate 0.0001 \
+  --timestamp_probability 0.2 \
+  --condition_on_prev_probability 0.2 \
+  --language "de" \
+  --task "transcribe" \
+  --logging_steps 25 \
+  --save_total_limit 1 \
+  --max_steps 50000 \
+  --wer_threshold 10 \
+  --per_device_train_batch_size 32 \
+  --per_device_eval_batch_size 32 \
+  --dataloader_num_workers 8 \
+  --preprocessing_num_workers 8 \
+  --ddp_timeout 7200 \
+  --dtype "bfloat16" \
+  --attn_implementation "flash_attention_2" \
+  --output_dir "./" \
+  --do_train \
+  --do_eval \
+  --gradient_checkpointing \
+  --overwrite_output_dir \
+  --predict_with_generate \
+  --freeze_encoder \
+  --freeze_embed_positions \
+  --streaming False \
+  --push_to_hub
+

run_all.sh

@@ -0,0 +1,67 @@
+#!/usr/bin/env bash
+
+accelerate launch run_pseudo_labelling.py \
+  --model_name_or_path "openai/whisper-large-v3" \
+  --dataset_name "mozilla-foundation/common_voice_16_1" \
+  --dataset_config_name "de" \
+  --dataset_split_name "train+validation+test" \
+  --text_column_name "sentence" \
+  --id_column_name "path" \
+  --output_dir "../common_voice_16_1_de_pseudo_labelled" \
+  --wandb_project "distil-whisper-labelling" \
+  --per_device_eval_batch_size 64 \
+  --dtype "bfloat16" \
+  --attn_implementation "flash_attention_2" \
+  --logging_steps 500 \
+  --max_label_length 256 \
+  --concatenate_audio \
+  --preprocessing_batch_size 500 \
+  --preprocessing_num_workers 8 \
+  --dataloader_num_workers 8 \
+  --report_to "wandb" \
+  --language "de" \
+  --task "transcribe" \
+  --return_timestamps \
+  --streaming False \
+  --generation_num_beams 1 \
+  --push_to_hub
+
+accelerate launch run_distillation.py \
+  --model_name_or_path "./distil-large-v3-init" \
+  --teacher_model_name_or_path "openai/whisper-large-v3" \
+  --train_dataset_name "../common_voice_16_1_de_pseudo_labelled" \
+  --train_split_name "train" \
+  --text_column_name "sentence" \
+  --eval_dataset_name "mozilla-foundation/common_voice_16_1" \
+  --eval_split_name "validation" \
+  --eval_text_column_name "sentence" \
+  --eval_steps 5000 \
+  --save_steps 5000 \
+  --warmup_steps 500 \
+  --learning_rate 0.0001 \
+  --timestamp_probability 0.2 \
+  --condition_on_prev_probability 0.2 \
+  --language "de" \
+  --task "transcribe" \
+  --logging_steps 25 \
+  --save_total_limit 1 \
+  --max_steps 50000 \
+  --wer_threshold 10 \
+  --per_device_train_batch_size 32 \
+  --per_device_eval_batch_size 32 \
+  --dataloader_num_workers 8 \
+  --preprocessing_num_workers 8 \
+  --ddp_timeout 7200 \
+  --dtype "bfloat16" \
+  --attn_implementation "flash_attention_2" \
+  --output_dir "./" \
+  --do_train \
+  --do_eval \
+  --gradient_checkpointing \
+  --overwrite_output_dir \
+  --predict_with_generate \
+  --freeze_encoder \
+  --freeze_embed_positions \
+  --streaming False \
+  --push_to_hub
+

run_distillation.py

@@ -0,0 +1,1683 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Training the Whisper model for sequence to sequence speech recognition via teacher-student distillation.
+"""
+# You can also adapt this script for your own distillation tasks. Pointers for this are left as comments.
+
+import logging
+import os
+import re
+import shutil
+import sys
+import time
+from dataclasses import dataclass, field
+from functools import partial
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import datasets
+import evaluate
+import numpy as np
+import torch
+import torch.nn as nn
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from datasets import (
+    DatasetDict,
+    IterableDataset,
+    IterableDatasetDict,
+    concatenate_datasets,
+    interleave_datasets,
+    load_dataset,
+)
+from huggingface_hub import Repository, create_repo, get_full_repo_name, upload_folder
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import (
+    AddedToken,
+    HfArgumentParser,
+    Seq2SeqTrainingArguments,
+    WhisperConfig,
+    WhisperFeatureExtractor,
+    WhisperForConditionalGeneration,
+    WhisperProcessor,
+    WhisperTokenizerFast,
+    get_scheduler,
+    set_seed,
+)
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.models.whisper.english_normalizer import BasicTextNormalizer, EnglishTextNormalizer
+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.34.0.dev0")
+
+require_version("datasets>=2.14.6", "To fix: `pip install --upgrade datasets`")
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to distill from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained Whisper model or model identifier from huggingface.co/models"}
+    )
+    teacher_model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained teacher model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained config name or path if not the same as model_name"},
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"},
+    )
+    feature_extractor_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "feature extractor name or path if not the same as model_name"},
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    subfolder: str = field(
+        default="",
+        metadata={
+            "help": "In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can"
+            "specify the folder name here."
+        },
+    )
+    token: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
+                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
+            )
+        },
+    )
+    attn_implementation: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+            "Which attention implementation to use in the encoder and decoder attention layers. Can be one of:\n"
+            "1. `eager` or `None`: default Transformers attention implementation.\n"
+            "2. `sdpa`: Flash Attention through PyTorch SDPA. Requires `torch>=2.1`. Recommended for hardware where Flash Attention 2 is not supported, e.g. Turing GPUs, (T4, RTX 2080).\n"
+            "3. `flash_attn_2`: Flash Attention 2 through the Flash Attention package https://github.com/Dao-AILab/flash-attention. **Always** recommended on supported hardware (Ampere, Ada, or Hopper GPUs, e.g., A100, RTX 3090, RTX 4090, H100)."
+        )
+        },
+    )
+    def __post_init__(self):
+        if self.attn_implementation not in [None, "eager", "sdpa", "flash_attention_2"]:
+            raise ValueError(
+                f"Got `--attn_implementation={self.attn_implementation}`, which is an invalid attention type. Should be one of:\n"
+                "1. `eager` or `None`: default Transformers attention implementation.\n"
+                "2. `sdpa`: Flash Attention through PyTorch SDPA. Requires `torch>=2.1`. Recommended for hardware where Flash Attention 2 is not supported, e.g. Turing GPUs, (T4, RTX 2080).\n"
+                "3. `flash_attn_2`: Flash Attention 2 through the Flash Attention package https://github.com/Dao-AILab/flash-attention. **Always** recommended on supported hardware (Ampere, Ada, or Hopper GPUs, e.g., A100, RTX 3090, RTX 4090, H100)."
+            )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    train_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load LibriSpeech "
+            "and Common Voice, set `train_dataset_name='librispeech_asr+common_voice'`."
+        },
+    )
+    train_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the training dataset to use (via the datasets library). Load and combine "
+            "multiple datasets by separating dataset configs by a '+' symbol. Note that the order of the configs should "
+            "match the order of the datasets."
+        },
+    )
+    train_dataset_samples: str = field(
+        default=None,
+        metadata={
+            "help": "Number of samples in each dataset when loading multiple datasets with streaming mode. "
+            "Not required when using one dataset or non-streaming mode. The sample values provide the sampling "
+            "probability for each dataset. Setting them equal to the number of sample values ensures that every "
+            "sample from every dataset is used once per epoch."
+        },
+    )
+    eval_dataset_name: str = field(
+        default=None,
+        metadata={
+            "help": "The name of the evaluation dataset to use (via the datasets library). Defaults to the training "
+            "dataset name if unspecified. Load multiple evaluation datasets by separating dataset "
+            "ids by a '+' symbol."
+        },
+    )
+    eval_dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The configuration name of the evaluation dataset to use (via the datasets library). Defaults to the "
+            "training dataset config name if unspecified."
+        },
+    )
+    dataset_cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to cache directory for saving and loading datasets"},
+    )
+    overwrite_cache: bool = field(
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets"},
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing if using non-streaming mode."},
+    )
+    preprocessing_batch_size: Optional[int] = field(
+        default=256,
+        metadata={"help": "Number of examples per batch provided to the `prepare_dataset` function."},
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this value if set."
+            )
+        },
+    )
+    audio_column_name: str = field(
+        default="audio",
+        metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},
+    )
+    text_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the text data in the training set."},
+    )
+    eval_text_column_name: str = field(
+        default="text",
+        metadata={"help": ("The name of the dataset column containing the text data in the evaluation set.")},
+    )
+    max_duration_in_seconds: float = field(
+        default=30.0,
+        metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"},
+    )
+    min_duration_in_seconds: float = field(
+        default=0.0,
+        metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"},
+    )
+    max_label_length: int = field(
+        default=448,
+        metadata={"help": "Truncate transcriptions that are longer `max_label_length` tokens."},
+    )
+    pad_target_to_multiple_of: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "If set will pad the target sequence to a multiple of the provided"
+                " value. This is important to avoid triggering recompilations on TPU."
+                " If unspecified, will default to padding the targets to max length."
+            )
+        },
+    )
+    preprocessing_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to only do data preprocessing and skip training. This is"
+                " especially useful when data preprocessing errors out in distributed"
+                " training due to timeout. In this case, one should run the"
+                " preprocessing in a non-distributed setup with"
+                " `preprocessing_only=True` so that the cached datasets can"
+                " consequently be loaded in distributed training"
+            )
+        },
+    )
+    train_split_name: str = field(
+        default="train",
+        metadata={
+            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    eval_split_name: str = field(
+        default="validation",
+        metadata={
+            "help": (
+                "The name of the evaluation data set split to use (via the datasets library). Defaults to 'validation'"
+            )
+        },
+    )
+    streaming: bool = field(
+        default=True,
+        metadata={"help": "Whether to use Datasets' streaming mode to load and pre-process the data."},
+    )
+    wer_threshold: float = field(
+        default=None,
+        metadata={
+            "help": "Filter training data with Whisper transcriptions that have greater than `wer_threshold` "
+            "WER with the normalised transcriptions. This only takes effect if training on pseudo-labels targets."
+            "If `--use_pseudo_labels=False`, then no WER filtering is performed, since we train directly on the text"
+            "transcriptions."
+        },
+    )
+    use_pseudo_labels: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether or not to use pseudo-label transcriptions as the targets. If True, the pseudo-labels "
+            "must be in the dataset column `whisper_transcript` from the previous pseudo-labelling step. This is "
+            "not currently yet configurable."
+        },
+    )
+    timestamp_probability: float = field(
+        default=0.2, metadata={"help": "Probability for training on timestamped tokens if the data contains it."}
+    )
+    condition_on_prev_probability: float = field(
+        default=0.2, metadata={"help": "Probability for conditioning on the previous text example."}
+    )
+    return_timestamps: bool = field(
+        default=False, metadata={"help": "Whether or not to predict timestamps in the generation step."}
+    )
+    language: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "Language for multilingual distillation. This argument should be set for multilingual distillation "
+                "only. For English speech recognition, it should be left as `None`."
+            )
+        },
+    )
+    task: str = field(
+        default="transcribe",
+        metadata={
+            "help": "Task, either `transcribe` for speech recognition or `translate` for speech translation."
+            "This argument should be set for multilingual distillation only. For English speech recognition, it should be left as `None`."
+        },
+    )
+    wandb_project: str = field(
+        default="distil-whisper",
+        metadata={"help": "The name of the wandb project."},
+    )
+
+
+@dataclass
+class DistillationTrainingArguments(Seq2SeqTrainingArguments):
+    freeze_encoder: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to freeze the entire encoder model. Only recommended when the entire encoder has been "
+                "copied from the teacher model."
+            )
+        },
+    )
+    freeze_embed_positions: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to freeze the decoder embedding positions."},
+    )
+    temperature: Optional[float] = field(
+        default=2.0, metadata={"help": "Temperature to anneal the logits when computing the softmax."}
+    )
+    kl_weight: Optional[float] = field(
+        default=1.0,
+        metadata={
+            "help": (
+                "Weighting assigned to the MSE loss in the KD formulation. MSE loss is "
+                "computed between the teacher-student hidden states and attentions."
+            )
+        },
+    )
+    dtype: Optional[str] = field(
+        default="float32",
+        metadata={
+            "help": (
+                "The data type (dtype) in which to run training. One of `float32` (full-precision), "
+                "`float16` or `bfloat16` (both half-precision)."
+            )
+        },
+    )
+
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        processor ([`Wav2Vec2Processor`])
+            The processor used for proccessing the data.
+        decoder_start_token_id (:obj: `int`)
+            The start-of-sequence token id of the decoder.
+        decoder_prev_token_id (:obj: `int`)
+            The start-of-prompt token id of the decoder
+        input_padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned input sequences (according to the model's padding side and padding index)
+            among:
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        target_padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned target sequences (according to the model's padding side and padding index).
+            See above for details.
+        max_target_length (:obj:`int`, `optional`):
+            Maximum length of the ``labels`` of the returned list and optionally padding length (see above).
+    """
+
+    processor: Any
+    decoder_start_token_id: int
+    decoder_prev_token_id: int
+    input_padding: Union[bool, str] = "max_length"
+    target_padding: Union[bool, str] = "max_length"
+    max_target_length: Optional[int] = None
+
+    def __call__(self, features: List[Dict[str, Union[List[int], np.ndarray]]]) -> Dict[str, np.ndarray]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+
+        # dataloader returns a list of features which we convert to a dict
+        input_features = {"input_features": [feature["input_features"] for feature in features]}
+        label_features = {"input_ids": [feature["labels"] for feature in features]}
+
+        # reformat list to dict and set to pytorch format
+        batch = self.processor.feature_extractor.pad(
+            input_features,
+            padding=self.input_padding,
+            return_tensors="pt",
+        )
+
+        labels_batch = self.processor.tokenizer.pad(
+            label_features,
+            max_length=self.max_target_length,
+            padding=self.target_padding,
+            return_tensors="pt",
+        )
+
+        # shift labels to the right to get decoder input ids
+        labels = labels_batch["input_ids"]
+        decoder_input_ids = labels[:, :-1]
+        labels = labels[:, 1:]
+        labels_mask = labels_batch.attention_mask[:, 1:]
+
+        # replace padding with -100 to ignore correctly when computing the loss
+        labels = labels.masked_fill(labels_mask.ne(1), -100)
+
+        # replace initial prompt tokens with -100 to ignore correctly when computing the loss
+        bos_index = torch.argmax((labels == self.decoder_start_token_id).long(), dim=1)
+        bos_index = torch.where(bos_index > 0, bos_index + 1, bos_index)
+        prompt_mask = torch.arange(labels.shape[1]) < bos_index[:, None]
+        labels = torch.where(prompt_mask, -100, labels)
+
+        batch["labels"] = labels
+        batch["decoder_input_ids"] = decoder_input_ids
+
+        return batch
+
+
+def log_metric(
+    accelerator,
+    metrics: Dict,
+    train_time: float,
+    step: int,
+    epoch: int,
+    learning_rate: float = None,
+    prefix: str = "train",
+):
+    """Helper function to log all training/evaluation metrics with the correct prefixes and styling."""
+    log_metrics = {}
+    for k, v in metrics.items():
+        log_metrics[f"{prefix}/{k}"] = v
+    log_metrics[f"{prefix}/time"] = train_time
+    log_metrics[f"{prefix}/epoch"] = epoch
+    if learning_rate is not None:
+        log_metrics[f"{prefix}/learning_rate"] = learning_rate
+    accelerator.log(log_metrics, step=step)
+
+
+def log_pred(
+    accelerator,
+    pred_str: List[str],
+    label_str: List[str],
+    norm_pred_str: List[str],
+    norm_label_str: List[str],
+    step: int,
+    prefix: str = "eval",
+    num_lines: int = 200000,
+):
+    """Helper function to log target/predicted transcriptions to weights and biases (wandb)."""
+    if accelerator.is_main_process:
+        wandb_tracker = accelerator.get_tracker("wandb")
+        # pretty name for current step: step 50000 -> step 50k
+        cur_step_pretty = f"{int(step // 1000)}k" if step > 1000 else step
+        prefix_pretty = prefix.replace("/", "-")
+
+        # convert str data to a wandb compatible format
+        str_data = [[label_str[i], pred_str[i], norm_label_str[i], norm_pred_str[i]] for i in range(len(pred_str))]
+        # log as a table with the appropriate headers
+        wandb_tracker.log_table(
+            table_name=f"predictions/{prefix_pretty}-step-{cur_step_pretty}",
+            columns=["Target", "Pred", "Norm Target", "Norm Pred"],
+            data=str_data[:num_lines],
+            step=step,
+        )
+
+        # log incorrect normalised predictions
+        str_data = np.asarray(str_data)
+        str_data_incorrect = str_data[str_data[:, -2] != str_data[:, -1]]
+        # log as a table with the appropriate headers
+        wandb_tracker.log_table(
+            table_name=f"incorrect_predictions/{prefix_pretty}-step-{cur_step_pretty}",
+            columns=["Target", "Pred", "Norm Target", "Norm Pred"],
+            data=str_data_incorrect[:num_lines],
+            step=step,
+        )
+
+
+def convert_dataset_str_to_list(
+    dataset_names,
+    dataset_config_names,
+    splits=None,
+    text_column_names=None,
+    dataset_samples=None,
+    default_split="train",
+) -> List[Dict]:
+    """
+    Given three lists of dataset names, configs and splits, this function groups the corresponding
+    names/configs/splits. Each dataset is assigned a unique dictionary with these metadata values, and the
+    function returns a list of dictionaries, one for each dataset.
+    """
+    if isinstance(dataset_names, str):
+        dataset_names = dataset_names.split("+")
+        dataset_config_names = dataset_config_names.split("+") if dataset_config_names is not None else None
+        splits = splits.split("+") if splits is not None else None
+        text_column_names = text_column_names.split("+") if text_column_names is not None else None
+        dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
+
+    # basic checks to ensure we've got the right number of datasets/configs/splits/columns/probs
+    if dataset_config_names is not None and len(dataset_names) != len(dataset_config_names):
+        raise ValueError(
+            f"Ensure one config is passed for each dataset, got {len(dataset_names)} datasets and"
+            f" {len(dataset_config_names)} configs."
+        )
+
+    if splits is not None and len(splits) != len(dataset_names):
+        raise ValueError(
+            f"Ensure one split is passed for each dataset, got {len(dataset_names)} datasets and {len(splits)} splits."
+        )
+
+    if text_column_names is not None and len(text_column_names) != len(dataset_names):
+        raise ValueError(
+            f"Ensure one text column name is passed for each dataset, got {len(dataset_names)} datasets and"
+            f" {len(text_column_names)} text column names."
+        )
+
+    if dataset_samples is not None:
+        if len(dataset_samples) != len(dataset_names):
+            raise ValueError(
+                f"Ensure one sample is passed for each dataset, got {len(dataset_names)} datasets and "
+                f"{len(dataset_samples)} samples."
+            )
+        dataset_samples = [float(ds_sample) for ds_sample in dataset_samples]
+    else:
+        dataset_samples = [None] * len(dataset_names)
+
+    dataset_config_names = (
+        dataset_config_names if dataset_config_names is not None else ["default" for _ in range(len(dataset_names))]
+    )
+    text_column_names = (
+        text_column_names if text_column_names is not None else ["text" for _ in range(len(dataset_names))]
+    )
+    splits = splits if splits is not None else [default_split for _ in range(len(dataset_names))]
+
+    dataset_names_dict = []
+    for i, ds_name in enumerate(dataset_names):
+        dataset_names_dict.append(
+            {
+                "name": ds_name,
+                "config": dataset_config_names[i],
+                "split": splits[i],
+                "text_column_name": text_column_names[i],
+                "samples": dataset_samples[i],
+            }
+        )
+    return dataset_names_dict
+
+
+def load_multiple_datasets(
+    dataset_names: Union[List, str],
+    dataset_config_names: Union[List, str],
+    splits: Optional[Union[List, str]] = None,
+    text_column_names: Optional[List] = None,
+    sampling_rate: Optional[int] = 16000,
+    stopping_strategy: Optional[str] = "first_exhausted",
+    dataset_samples: Optional[Union[List, np.array]] = None,
+    streaming: Optional[bool] = True,
+    seed: Optional[int] = None,
+    accelerator: Optional[Accelerator] = None,
+    use_pseudo_labels: float = None,
+    **kwargs,
+) -> IterableDataset:
+    dataset_names_dict = convert_dataset_str_to_list(
+        dataset_names, dataset_config_names, splits, text_column_names, dataset_samples
+    )
+
+    if dataset_samples is not None:
+        dataset_samples = [ds_dict["samples"] for ds_dict in dataset_names_dict]
+        probabilities = np.array(dataset_samples) / np.sum(dataset_samples)
+    else:
+        probabilities = None
+
+    all_datasets = []
+    # iterate over the datasets we want to interleave
+    for dataset_dict in tqdm(
+        dataset_names_dict,
+        desc="Combining datasets...",
+        disable=not accelerator.is_local_main_process if accelerator is not None else False,
+    ):
+        dataset = load_dataset(
+            dataset_dict["name"],
+            dataset_dict["config"],
+            split=dataset_dict["split"],
+            streaming=streaming,
+            **kwargs,
+        )
+        # resample to specified sampling rate
+        dataset = dataset.cast_column("audio", datasets.features.Audio(sampling_rate))
+        dataset_features = dataset.features.keys()
+        columns_to_keep = {"audio", "text"}
+
+        if dataset_dict["text_column_name"] not in dataset_features:
+            raise ValueError(
+                f"Text column name {dataset_dict['text_column_name']} not found in dataset"
+                f" '{dataset_dict['name']}'. Make sure to set `--text_column_name` to the"
+                f" correct text column - one of {', '.join(dataset_features)}."
+            )
+
+        # blanket renaming of all transcription columns to text
+        if dataset_dict["text_column_name"] != "text":
+            dataset = dataset.rename_column(dataset_dict["text_column_name"], "text")
+
+        if use_pseudo_labels:
+            if "whisper_transcript" not in dataset_features:
+                raise ValueError(
+                    f"Pseudo-label column `whisper_transcript` not found in dataset {dataset_dict['name']}. Ensure"
+                    "pseudo-labels are present in the dataset under this column name, or train directly on the text "
+                    "labels by setting `--use_pseudo_labels=False` and defining the appropriate `--text_column_name`."
+                )
+            columns_to_keep.add("whisper_transcript")
+
+        if "condition_on_prev" in dataset_features:
+            columns_to_keep.add("condition_on_prev")
+
+        dataset_features = dataset.features.keys()
+        dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
+        all_datasets.append(dataset)
+
+    if len(all_datasets) == 1:
+        # we have a single dataset so just return it as is
+        return all_datasets[0]
+
+    if streaming:
+        interleaved_dataset = interleave_datasets(
+            all_datasets,
+            stopping_strategy=stopping_strategy,
+            probabilities=probabilities,
+            seed=seed,
+        )
+    else:
+        interleaved_dataset = concatenate_datasets(all_datasets)
+
+    return interleaved_dataset
+
+
+def sorted_checkpoints(output_dir=None, checkpoint_prefix="checkpoint") -> List[str]:
+    """Helper function to sort saved checkpoints from oldest to newest."""
+    ordering_and_checkpoint_path = []
+
+    glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]
+
+    for path in glob_checkpoints:
+        regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
+        if regex_match is not None and regex_match.groups() is not None:
+            ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
+
+    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
+    checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
+    return checkpoints_sorted
+
+
+def rotate_checkpoints(save_total_limit=None, output_dir=None, checkpoint_prefix="checkpoint") -> None:
+    """Helper function to delete old checkpoints."""
+    if save_total_limit is None or save_total_limit <= 0:
+        return
+    # Check if we should delete older checkpoint(s)
+    checkpoints_sorted = sorted_checkpoints(output_dir=output_dir, checkpoint_prefix=checkpoint_prefix)
+    if len(checkpoints_sorted) <= save_total_limit:
+        return
+
+    number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
+    checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
+    for checkpoint in checkpoints_to_be_deleted:
+        logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
+        shutil.rmtree(checkpoint, ignore_errors=True)
+
+
+_RE_CHECKPOINT = re.compile(r"^checkpoint-(\d+)-epoch-(\d+)$")
+
+
+def get_last_checkpoint(folder):
+    content = os.listdir(folder)
+    checkpoints = [
+        path
+        for path in content
+        if _RE_CHECKPOINT.search(path) is not None and os.path.isdir(os.path.join(folder, path))
+    ]
+    if len(checkpoints) == 0:
+        return
+    return os.path.join(folder, max(checkpoints, key=lambda x: int(_RE_CHECKPOINT.search(x).groups()[0])))
+
+
+def get_parameter_names(model, forbidden_layer_types, forbidden_module=None):
+    """
+    Returns the names of the model parameters that are not inside a forbidden layer or forbidden module.
+    Can be used to get a subset of parameter names for decay masks, or to exclude parameters from an optimiser
+    (e.g. if the module is frozen).
+    """
+    result = []
+    for name, child in model.named_children():
+        result += [
+            f"{name}.{n}"
+            for n in get_parameter_names(child, forbidden_layer_types, forbidden_module)
+            if not (
+                isinstance(child, tuple(forbidden_layer_types))
+                or (child in tuple(forbidden_module) if forbidden_module is not None else False)
+            )
+        ]
+    # Add model specific parameters (defined with nn.Parameter) since they are not in any child.
+    result += list(model._parameters.keys())
+    return result
+
+
+def main():
+    # 1. Parse input arguments
+    # We keep distinct sets of args, for cleaner separation of model/data/training related args
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, DistillationTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # 2. Initialize the accelerator
+    # We will let the accelerator handle device placement for us in this example
+    # We simply have to specify the training precision and any trackers being used
+    # We'll use the same dtype arguments as our JAX/Flax training script and convert
+    # it to accelerate format
+    if training_args.dtype == "float16":
+        mixed_precision = "fp16"
+        teacher_dtype = torch.float16
+    elif training_args.dtype == "bfloat16":
+        mixed_precision = "bf16"
+        teacher_dtype = torch.bfloat16
+    else:
+        mixed_precision = "no"
+        teacher_dtype = torch.float32
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        mixed_precision=mixed_precision,
+        log_with=training_args.report_to,
+        project_dir=training_args.output_dir,
+    )
+
+    accelerator.init_trackers(project_name=data_args.wandb_project)
+
+    # 3. Set-up basic logging
+    # Create one log on every process with the configuration for debugging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    # Log a small summary on each proces
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
+        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set the verbosity to info of the Transformers logger (on main process only)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # 4. Detecting last checkpoint and eventually continue from last checkpoint
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # 5. Handle the repository creation
+    if accelerator.is_main_process:
+        if training_args.push_to_hub:
+            if training_args.hub_model_id is None:
+                repo_name = get_full_repo_name(
+                    Path(training_args.output_dir).absolute().name,
+                    token=training_args.hub_token,
+                )
+            else:
+                repo_name = training_args.hub_model_id
+            create_repo(repo_name, exist_ok=True, token=training_args.hub_token)
+
+            with open(os.path.join(training_args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "wandb" not in gitignore:
+                    gitignore.write("wandb\n")
+        elif training_args.output_dir is not None:
+            os.makedirs(training_args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # 6. Load dataset - either streaming or non-streaming (offline)
+    raw_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict()
+
+    # set seed for determinism
+    set_seed(training_args.seed)
+
+    if training_args.do_train:
+        raw_datasets["train"] = load_multiple_datasets(
+            data_args.train_dataset_name,
+            data_args.train_dataset_config_name,
+            splits=data_args.train_split_name,
+            text_column_names=data_args.text_column_name,
+            use_pseudo_labels=data_args.use_pseudo_labels,
+            streaming=data_args.streaming,
+            dataset_samples=data_args.train_dataset_samples,
+            seed=training_args.seed,
+            accelerator=accelerator,
+            cache_dir=data_args.dataset_cache_dir,
+            token=model_args.token,
+        )
+        raw_datasets_train_features = list(raw_datasets["train"].features.keys())
+
+    if training_args.do_eval:
+        dataset_names_dict = convert_dataset_str_to_list(
+            data_args.eval_dataset_name if data_args.eval_dataset_name else data_args.train_dataset_name,
+            data_args.eval_dataset_config_name
+            if data_args.eval_dataset_config_name
+            else data_args.train_dataset_config_name,
+            splits=data_args.eval_split_name,
+            text_column_names=data_args.eval_text_column_name,
+        )
+        all_eval_splits = []
+        if len(dataset_names_dict) == 1:
+            # load a single eval set
+            dataset_dict = dataset_names_dict[0]
+            all_eval_splits.append("eval")
+            raw_datasets["eval"] = load_dataset(
+                dataset_dict["name"],
+                dataset_dict["config"],
+                split=dataset_dict["split"],
+                cache_dir=data_args.dataset_cache_dir,
+                token=model_args.token,
+                streaming=data_args.streaming,
+            )
+            if data_args.eval_text_column_name != "text":
+                raw_datasets["eval"] = raw_datasets["eval"].rename_column(data_args.eval_text_column_name, "text")
+        else:
+            # load multiple eval sets
+            for dataset_dict in dataset_names_dict:
+                if dataset_dict["name"] == "esb/diagnostic-dataset":
+                    # for the ESB diagnostic dataset, the dataset name is effectively the config
+                    pretty_name = f"{dataset_dict['config']}-diagnostic/{dataset_dict['split']}"
+                else:
+                    pretty_name = f"{dataset_dict['name'].split('/')[-1]}/{dataset_dict['split'].replace('.', '-')}"
+                all_eval_splits.append(pretty_name)
+                raw_datasets[pretty_name] = load_dataset(
+                    dataset_dict["name"],
+                    dataset_dict["config"],
+                    split=dataset_dict["split"],
+                    cache_dir=data_args.dataset_cache_dir,
+                    token=model_args.token,
+                    streaming=data_args.streaming,
+                )
+                # make column names consistent (text, audio)
+                if dataset_dict["text_column_name"] != "text":
+                    raw_datasets[pretty_name] = raw_datasets[pretty_name].rename_column(
+                        dataset_dict["text_column_name"], "text"
+                    )
+                raw_datasets[pretty_name] = raw_datasets[pretty_name].remove_columns(
+                    set(raw_datasets[pretty_name].features.keys()) - {"audio", "text"}
+                )
+
+    if not training_args.do_train and not training_args.do_eval:
+        raise ValueError(
+            "Cannot not train and not do evaluation. At least one of training or evaluation has to be performed."
+        )
+
+    # 7. Load pretrained model, tokenizer, and feature extractor
+    config = WhisperConfig.from_pretrained(
+        (model_args.config_name if model_args.config_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        token=model_args.token,
+    )
+    feature_extractor = WhisperFeatureExtractor.from_pretrained(
+        (model_args.feature_extractor_name if model_args.feature_extractor_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        token=model_args.token,
+    )
+    tokenizer = WhisperTokenizerFast.from_pretrained(
+        (model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+        token=model_args.token,
+    )
+
+    # override timestamp tokens until tokenizer issues are fixed in transformers
+    timestamps = [AddedToken("<|%.2f|>" % (i * 0.02), lstrip=False, rstrip=False) for i in range(1500 + 1)]
+    tokenizer.add_tokens(timestamps)
+
+    # The teacher model can safely be cast to the dtype of training since we don't
+    # update the params
+    teacher_model = WhisperForConditionalGeneration.from_pretrained(
+        model_args.teacher_model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        token=model_args.token,
+        low_cpu_mem_usage=True,
+        torch_dtype=teacher_dtype,
+        attn_implementation=model_args.attn_implementation,
+    )
+
+    student_model = WhisperForConditionalGeneration.from_pretrained(
+        model_args.model_name_or_path,
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        subfolder=model_args.subfolder,
+        token=model_args.token,
+        low_cpu_mem_usage=True,
+        attn_implementation=model_args.attn_implementation,
+    )
+
+    if student_model.config.decoder_start_token_id is None or teacher_model.config.decoder_start_token_id is None:
+        raise ValueError(
+            f"Make sure that `config.decoder_start_token_id` is correctly defined for both the "
+            f"student and teacher model. Got {student_model.config.decoder_start_token_id} for the "
+            f"student and {teacher_model.config.decoder_start_token_id} for the teacher."
+        )
+
+    # enable gradient checkpointing if necessary
+    if training_args.gradient_checkpointing:
+        student_model.gradient_checkpointing_enable()
+
+    def set_trainable_parameters(module, requires_grad=False):
+        for param in module.parameters():
+            param.requires_grad = requires_grad
+        module._requires_grad = requires_grad
+
+    # freeze student encoder if necessary
+    if training_args.freeze_encoder:
+        set_trainable_parameters(student_model.model.encoder, requires_grad=False)
+        student_model.model.encoder.gradient_checkpointing = False
+
+    if training_args.freeze_embed_positions:
+        # set_trainable_parameters(student_model.model.decoder.embed_tokens, requires_grad=False)
+        set_trainable_parameters(student_model.model.decoder.embed_positions, requires_grad=False)
+        if student_model.model.decoder.gradient_checkpointing:
+            logger.info(
+                "Disabling gradient checkpointing in the decoder since it's incompatible with `freeze_embed_positions`."
+            )
+
+    share_hidden_states = training_args.freeze_encoder and student_model.config.d_model == teacher_model.config.d_model
+    if share_hidden_states:
+        # tie the weights for the teacher encoder if we're freezing the student and it's the same as the teacher
+        teacher_model.model.encoder = student_model.model.encoder
+
+    if hasattr(teacher_model.generation_config, "is_multilingual") and teacher_model.generation_config.is_multilingual:
+        # We need to set the language and task ids for previously multilingual checkpoints
+        is_multilingual = True
+        tokenizer.set_prefix_tokens(language=data_args.language, task=data_args.task, predict_timestamps=False)
+        student_model.generation_config.update(
+            **{
+                "language": data_args.language,
+                "task": data_args.task,
+            }
+        )
+    elif data_args.language is not None:
+        raise ValueError(
+            "Setting language token for an English-only checkpoint is not permitted. The language argument should "
+            "only be set for multilingual checkpoints."
+        )
+    else:
+        is_multilingual = False
+
+    # 8. Create a single speech processor - make sure all processes wait until data is saved
+    if accelerator.is_main_process:
+        feature_extractor.save_pretrained(training_args.output_dir)
+        tokenizer.save_pretrained(training_args.output_dir)
+        # save the config and generation config as well
+        config.save_pretrained(training_args.output_dir)
+        student_model.generation_config.save_pretrained(training_args.output_dir)
+
+    accelerator.wait_for_everyone()
+    processor = WhisperProcessor.from_pretrained(training_args.output_dir)
+
+    # 9. Resample speech dataset: `datasets` takes care of automatically loading and resampling the audio,
+    # so we just need to set the correct target sampling rate.
+    sampling_rate = feature_extractor.sampling_rate
+    raw_datasets = raw_datasets.cast_column(
+        data_args.audio_column_name,
+        datasets.features.Audio(sampling_rate=sampling_rate),
+    )
+
+    # 10. Preprocessing the datasets: we need to read the audio files as arrays and tokenize the targets.
+    # 10.1: Define the pre-processing constants
+    max_input_length = int(data_args.max_duration_in_seconds * sampling_rate)
+    min_input_length = int(data_args.min_duration_in_seconds * sampling_rate)
+    max_label_length = (
+        data_args.max_label_length if data_args.max_label_length is not None else student_model.config.max_length
+    )
+
+    timestamp_probability = data_args.timestamp_probability
+    condition_on_prev_probability = data_args.condition_on_prev_probability
+    return_timestamps = data_args.return_timestamps if timestamp_probability > 0 else False
+
+    timestamp_ids = tokenizer.timestamp_ids()
+    timestamp_begin = tokenizer.all_special_ids[-1]
+    timestamp_position = 3 if is_multilingual else 1
+
+    decoder_start_token_id = student_model.config.decoder_start_token_id  # <|startoftranscript|>
+    decoder_prev_token_id = tokenizer.all_special_ids[-3]  # <|startofprev|>
+    prompt_cutoff_length = max_label_length // 2
+
+    num_workers = data_args.preprocessing_num_workers
+    dataloader_num_workers = training_args.dataloader_num_workers
+    prefetch_factor = training_args.dataloader_prefetch_factor
+
+    metric = evaluate.load("wer")
+    normalizer = (
+        BasicTextNormalizer() if data_args.language is not None else EnglishTextNormalizer(tokenizer.english_spelling_normalizer)
+    )
+    wer_threshold = data_args.wer_threshold
+    use_pseudo_labels = data_args.use_pseudo_labels
+    train_text_column_name = "whisper_transcript" if use_pseudo_labels else "text"
+
+    # 10.2: filter based on maximum number of training/evaluation samples
+    if training_args.do_train and data_args.max_train_samples is not None:
+        raw_datasets["train"] = (
+            raw_datasets["train"].take(data_args.max_train_samples)
+            if data_args.streaming
+            else raw_datasets["train"].select(range(data_args.max_train_samples))
+        )
+
+    if training_args.do_eval and data_args.max_eval_samples is not None:
+        for eval_split in all_eval_splits:
+            raw_datasets[eval_split] = (
+                raw_datasets[eval_split].take(data_args.max_eval_samples)
+                if data_args.streaming
+                else raw_datasets[eval_split].select(range(data_args.max_eval_samples))
+            )
+
+    # 10.3: filter training data based on WER threshold -> this is KEY to good distillation performance
+    def is_wer_in_range(ground_truth, whisper_transcript):
+        norm_ground_truth = normalizer(ground_truth)
+        if whisper_transcript is not None and whisper_transcript.upper() == whisper_transcript:
+            # filter entirely upper-case transcriptions: these are erroneous generations from large-v3
+            return False
+        elif len(norm_ground_truth) > 0 and whisper_transcript is not None:
+            norm_whisper_transcript = normalizer(whisper_transcript)
+            wer = 100 * metric.compute(predictions=[norm_whisper_transcript], references=[norm_ground_truth])
+            return wer < wer_threshold
+        else:
+            # filter automatically since we can't know the WER
+            return False
+
+    filter_by_wer_threshold = partial(
+        raw_datasets["train"].filter,
+        function=is_wer_in_range,
+        input_columns=["text", "whisper_transcript"],
+    )
+
+    if wer_threshold is not None and use_pseudo_labels:
+        raw_datasets["train"] = (
+            filter_by_wer_threshold(num_proc=num_workers, desc="filtering train dataset by wer")
+            if not data_args.streaming
+            else filter_by_wer_threshold()
+        )
+
+    # 10.4: pre-process training/evaluation datasets
+    def prepare_train_dataset(batch):
+        """
+        Pre-process the raw dataset in a three stage process:
+            1. Convert the audio arrays to log-mel spectrogram inputs
+            2. Possibly filter the timestamp tokens from the token ids (depending on the timestamp probability)
+            3. Possibly add prompt tokens if conditioning on previous text (depending on the conditioning probability)
+        """
+        # process audio input
+        audio = [sample["array"] for sample in batch["audio"]]
+        inputs = feature_extractor(audio, sampling_rate=sampling_rate)
+        batch["input_features"] = inputs.input_features
+        batch["input_length"] = [len(sample) for sample in audio]
+
+        # process text targets - for training these are the Whisper-generated pseudo-labels
+        input_str_batched = batch[train_text_column_name]
+        condition_on_prev_batched = batch.get("condition_on_prev", len(input_str_batched) * [None])
+
+        all_token_ids = []
+        all_token_ids_unprompted = []
+        for prev_ids, input_str in zip(condition_on_prev_batched, input_str_batched):
+            token_ids = tokenizer(input_str, add_special_tokens=not use_pseudo_labels).input_ids
+
+            # check whether we have timestamps in the PLs and filter if required
+            has_timestamps = len(set(token_ids) & set(timestamp_ids)) > 0
+            if has_timestamps:
+                # sample from binomial distribution to get probability of training on timestamps
+                predict_timestamps = bool(np.random.binomial(1, timestamp_probability))
+                if not predict_timestamps:
+                    # filter timestamps and insert the <|notimestamps|> task token
+                    token_ids = [token for token in token_ids if token < timestamp_begin]
+                    token_ids.insert(timestamp_position, timestamp_begin)
+
+            all_token_ids_unprompted.append(token_ids)
+            # check whether to condition on previous text - we do this with probability condition_on_prev_probability
+            condition_on_prev = bool(np.random.binomial(1, condition_on_prev_probability))
+            if not condition_on_prev:
+                prev_ids = None
+            elif "condition_on_prev" not in batch and len(all_token_ids_unprompted) > 1:
+                # prompt ids are the penultimate token ids in the batch
+                prev_ids = all_token_ids_unprompted[-2]
+
+            if prev_ids is not None:
+                if has_timestamps and not predict_timestamps:
+                    # filter timestamp ids from prompt when not predicting timestamps
+                    prev_ids = [token for token in prev_ids if token < timestamp_begin]
+
+                # check that the length of the prompt does not exceed more than half the max label length (224)
+                if len(prev_ids) > prompt_cutoff_length:
+                    prev_ids = prev_ids[-prompt_cutoff_length + 1 :]
+                    prev_ids = [decoder_prev_token_id] + prev_ids
+
+                # and that the total length of the labels does not exceed the max label length (448)
+                if len(prev_ids + token_ids) > max_label_length:
+                    trim_length = len(prev_ids + token_ids) - max_label_length + 1
+                    prev_ids = prev_ids[trim_length:]
+                    prev_ids = [decoder_prev_token_id] + prev_ids
+
+                token_ids = prev_ids + token_ids
+
+            all_token_ids.append(token_ids)
+
+        batch["labels"] = all_token_ids
+        return batch
+
+    def prepare_eval_dataset(batch):
+        # process audio input
+        sample = batch["audio"]
+        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
+        batch["input_features"] = inputs.input_features[0]
+        batch["input_length"] = len(sample["array"])
+
+        # process targets - for evaluation these are the ground-truth transcriptions
+        input_str = batch["text"]
+        batch["labels"] = tokenizer(input_str).input_ids
+        return batch
+
+    vectorized_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict()
+    if training_args.do_train:
+        # with streaming mode we can only have 1 worker, whereas with non-streaming
+        # we can use `num_workers` (which is much faster)
+        # We gate the pre-processing function accordingly
+        map_fn_train = partial(
+            raw_datasets["train"].map,
+            function=prepare_train_dataset,
+            remove_columns=raw_datasets_train_features,
+            batched=True,
+            batch_size=data_args.preprocessing_batch_size,
+        )
+        if accelerator.is_main_process:
+            vectorized_datasets["train"] = (
+                map_fn_train(num_proc=num_workers, desc="preprocess train dataset")
+                if not data_args.streaming
+                else map_fn_train()
+            )
+    if training_args.do_eval:
+        for eval_split in all_eval_splits:
+            raw_datasets_eval_features = list(raw_datasets[eval_split].features.keys())
+            map_fn_eval = partial(
+                raw_datasets[eval_split].map, function=prepare_eval_dataset, remove_columns=raw_datasets_eval_features
+            )
+            if accelerator.is_main_process:
+                vectorized_datasets[eval_split] = (
+                    map_fn_eval(num_proc=num_workers, desc="preprocess eval dataset")
+                    if not data_args.streaming
+                    else map_fn_eval()
+                )
+
+    # 10.5: Filter training data with inputs longer than `max_input_length`
+    def is_audio_in_length_range(length):
+        return min_input_length < length < max_input_length
+
+    filter_by_audio_fn = partial(
+        vectorized_datasets.filter, function=is_audio_in_length_range, input_columns=["input_length"]
+    )
+    if accelerator.is_main_process:
+        vectorized_datasets = (
+            filter_by_audio_fn(num_proc=num_workers, desc="filtering train dataset by audio length")
+            if not data_args.streaming
+            else filter_by_audio_fn()
+        )
+
+    # 10.6: Filter training data with labels longer than `max_label_length`
+    def is_labels_in_length_range(labels):
+        return 0 < len(labels) <= max_label_length
+
+    filter_by_labels_fn = partial(
+        vectorized_datasets.filter, function=is_labels_in_length_range, input_columns=["labels"]
+    )
+    if accelerator.is_main_process:
+        vectorized_datasets = (
+            filter_by_labels_fn(num_proc=num_workers, desc="filtering train dataset")
+            if not data_args.streaming
+            else filter_by_labels_fn()
+        )
+
+    # Pre-processing complete!
+    # For large datasets it is advised to run the preprocessing on a
+    # single machine first with `--preprocessing_only` since there will mostly likely
+    # be a timeout when running the script in distributed mode.
+    # In a second step, `--preprocessing_only` can then be set to `False` to load the
+    # cached dataset
+    if data_args.preprocessing_only:
+        if data_args.streaming:
+            raise ValueError(
+                "When using streaming mode, dataset pre-processing is performed on the fly, hence there is no notion"
+                "of a cached pre-processed dataset. Remove the argument `--preprocessing_only` to run pre-processing "
+                "on the fly with streaming mode."
+            )
+        cache = {k: v.cache_files for k, v in vectorized_datasets.items()}
+        logger.info(f"Data preprocessing finished. Files cached at {cache}.")
+        return
+
+    # 11. Define Evaluation Metrics
+    def compute_metrics(preds, labels):
+        # replace padded labels by the padding token
+        for idx in range(len(labels)):
+            labels[idx][labels[idx] == -100] = tokenizer.pad_token_id
+
+        pred_str = tokenizer.batch_decode(preds, skip_special_tokens=True, decode_with_timestamps=return_timestamps)
+        # we do not want to group tokens when computing the metrics
+        label_str = tokenizer.batch_decode(labels, skip_special_tokens=True)
+        wer_ortho = 100 * metric.compute(predictions=pred_str, references=label_str)
+
+        # normalize everything and re-compute the WER
+        norm_pred_str = [normalizer(pred) for pred in pred_str]
+        norm_label_str = [normalizer(label) for label in label_str]
+        # for logging, we need the pred/labels to match the norm_pred/norm_labels, so discard any filtered samples here
+        pred_str = [pred_str[i] for i in range(len(norm_pred_str)) if len(norm_label_str[i]) > 0]
+        label_str = [label_str[i] for i in range(len(norm_label_str)) if len(norm_label_str[i]) > 0]
+        # filtering step to only evaluate the samples that correspond to non-zero normalized references:
+        norm_pred_str = [norm_pred_str[i] for i in range(len(norm_pred_str)) if len(norm_label_str[i]) > 0]
+        norm_label_str = [norm_label_str[i] for i in range(len(norm_label_str)) if len(norm_label_str[i]) > 0]
+
+        wer = 100 * metric.compute(predictions=norm_pred_str, references=norm_label_str)
+        return {"wer": wer, "wer_ortho": wer_ortho}, pred_str, label_str, norm_pred_str, norm_label_str
+
+    # 12. Define Training Schedule
+    # Store some constants
+    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
+    train_batch_size = per_device_train_batch_size * accelerator.num_processes
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+
+    if not data_args.streaming and training_args.max_steps < 0:
+        num_epochs = int(training_args.num_train_epochs)
+        steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
+        total_train_steps = steps_per_epoch * num_epochs
+    elif training_args.max_steps > 0:
+        logger.info("max_steps is given, it will override any value given in num_train_epochs")
+        total_train_steps = int(training_args.max_steps)
+        if not data_args.streaming:
+            steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
+            num_epochs = int(np.ceil(total_train_steps / steps_per_epoch))
+        else:
+            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
+            num_epochs = sys.maxsize
+            steps_per_epoch = total_train_steps
+    else:
+        raise ValueError("max_steps must be specified when training with a streaming (iterable) dataset")
+
+    if training_args.eval_steps is None:
+        logger.info(
+            f"eval_steps is not set, evaluating at the end of {'each epoch' if not data_args.streaming else 'training'}"
+        )
+        eval_steps = steps_per_epoch
+    else:
+        eval_steps = training_args.eval_steps
+
+    # 13. Define optimizer, LR scheduler, collator
+    decay_parameters = get_parameter_names(
+        student_model,
+        [nn.LayerNorm],
+        forbidden_module=[student_model.model.encoder] if training_args.freeze_encoder else None,
+    )
+    decay_parameters = [name for name in decay_parameters if "bias" not in name]
+    optimizer_grouped_parameters = [
+        {
+            "params": [param for name, param in student_model.named_parameters() if name in decay_parameters],
+            "weight_decay": training_args.weight_decay,
+        },
+        {
+            "params": [param for name, param in student_model.named_parameters() if name not in decay_parameters],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = torch.optim.AdamW(
+        params=optimizer_grouped_parameters,
+        lr=training_args.learning_rate,
+        betas=(training_args.adam_beta1, training_args.adam_beta2),
+        eps=training_args.adam_epsilon,
+    )
+
+    # LR scheduler gets stepped by `num_processes` each time -> account for this in warmup / total steps
+    lr_scheduler = get_scheduler(
+        name=training_args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=training_args.warmup_steps * accelerator.num_processes,
+        num_training_steps=total_train_steps * accelerator.num_processes,
+    )
+
+    data_collator = DataCollatorSpeechSeq2SeqWithPadding(
+        processor=processor,
+        decoder_start_token_id=decoder_start_token_id,
+        decoder_prev_token_id=decoder_prev_token_id,
+        input_padding="longest",
+        target_padding="max_length",
+        max_target_length=max_label_length,
+    )
+
+    # 14. Define generation arguments - we need to do this before we wrap the models in DDP
+    # so that we can still access the configs
+    num_beams = (
+        training_args.generation_num_beams
+        if training_args.generation_num_beams is not None
+        else getattr(student_model.generation_config, "num_beams", 1)
+    )
+
+    gen_kwargs = {
+        "max_length": max_label_length,
+        "num_beams": num_beams,
+        "return_timestamps": return_timestamps,
+    }
+    if is_multilingual:
+        # forcing the language and task tokens helps multilingual models in their generations
+        gen_kwargs.update(
+            {
+                "language": data_args.language,
+                "task": data_args.task,
+            }
+        )
+
+    # 15. Prepare everything with accelerate
+    student_model, teacher_model, optimizer, lr_scheduler = accelerator.prepare(
+        student_model, teacher_model, optimizer, lr_scheduler
+    )
+
+    def kl_divergence(target_distribution, log_predicted_distribution, labels):
+        kl_loss = nn.KLDivLoss(reduction="none")
+        divergence = kl_loss(log_predicted_distribution, target_distribution)
+        # ignore padded tokens from divergence, i.e. where labels are not set to -100
+        padding_mask = labels >= 0
+        padding_mask = padding_mask.unsqueeze(-1)
+        divergence = divergence * padding_mask
+        # take the average over the mini-batch
+        divergence = divergence.sum() / padding_mask.sum()
+        return divergence
+
+    # Define gradient update step fn
+    def train_step(
+        batch,
+        temperature=2.0,
+    ):
+        student_model.train()
+        teacher_model.eval()
+
+        student_outputs = student_model(**batch)
+        with torch.no_grad():
+            if share_hidden_states:
+                # if the student and teacher share the same frozen encoder then we don't have to recompute the
+                # encoder hidden-states for the teacher model, we can just re-use from the student
+                encoder_outputs = BaseModelOutput(student_outputs.encoder_last_hidden_state.to(dtype=teacher_dtype))
+                teacher_outputs = teacher_model(encoder_outputs=encoder_outputs, labels=batch["labels"])
+            else:
+                # do the full forward pass for the teacher model (encoder + decoder)
+                teacher_outputs = teacher_model(**batch)
+
+        # CE (data) loss
+        ce_loss = student_outputs.loss
+        # rescale distribution by temperature to ensure gradients scale correctly
+        teacher_distribution = nn.functional.softmax(teacher_outputs.logits / temperature, dim=-1)
+        # log softmax of student predictions for numerical stability
+        student_distribution = nn.functional.log_softmax(student_outputs.logits / temperature, dim=-1)
+        # KL-divergence loss (scaled by temperature)
+        kl_loss = kl_divergence(teacher_distribution, student_distribution, batch["labels"]) * temperature**2
+
+        # use Distil-Whisper formulation (fix weight of CE loss and tune KL weight)
+        loss = 0.8 * ce_loss + training_args.kl_weight * kl_loss
+        metrics = {"loss": loss, "ce_loss": ce_loss, "kl_loss": kl_loss}
+        return loss, metrics
+
+    # Define eval fn
+    def eval_step(batch):
+        student_model.eval()
+        teacher_model.eval()
+
+        with torch.no_grad():
+            student_outputs = student_model(**batch)
+            if share_hidden_states:
+                encoder_outputs = BaseModelOutput(student_outputs.encoder_last_hidden_state.to(dtype=teacher_dtype))
+                teacher_outputs = teacher_model(encoder_outputs=encoder_outputs, labels=batch["labels"])
+            else:
+                teacher_outputs = teacher_model(**batch)
+
+        # CE (data) loss
+        ce_loss = student_outputs.loss
+
+        # log softmax / softmax for numerical stability
+        student_distribution = nn.functional.log_softmax(student_outputs.logits, dim=-1)
+        teacher_distribution = nn.functional.softmax(teacher_outputs.logits, dim=-1)
+        # temperature is always 1 for eval
+        kl_loss = kl_divergence(teacher_distribution, student_distribution, batch["labels"])
+
+        # use Distil-Whisper formulation (fix weight of CE loss and tune KL weight)
+        loss = 0.8 * ce_loss + training_args.kl_weight * kl_loss
+        metrics = {"loss": loss, "ce_loss": ce_loss, "kl_loss": kl_loss}
+        return metrics
+
+    def generate_step(batch):
+        student_model.eval()
+        output_ids = accelerator.unwrap_model(student_model).generate(batch["input_features"], **gen_kwargs)
+        output_ids = accelerator.pad_across_processes(output_ids, dim=1, pad_index=tokenizer.pad_token_id)
+        return output_ids
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {total_train_steps * train_batch_size * gradient_accumulation_steps}")
+    if not data_args.streaming:
+        logger.info(f"  Num epochs = {num_epochs}")
+    logger.info("  Instantaneous batch size per device =" f" {training_args.per_device_train_batch_size}")
+    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
+    )
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+
+    # ======================== Training ================================
+    train_time = 0
+    train_start = time.time()
+    steps_trained_progress_bar = tqdm(
+        range(total_train_steps), desc="Train steps ... ", position=0, disable=not accelerator.is_local_main_process
+    )
+    continue_training = True
+    epochs_trained = 0
+    cur_step = 0
+
+    checkpoint = None
+    if training_args.resume_from_checkpoint is not None:
+        checkpoint = training_args.resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+
+    if checkpoint is not None:
+        accelerator.load_state(checkpoint)
+        # Find num steps and epoch from saved state string pattern
+        pattern = r"checkpoint-(\d+)-epoch-(\d+)"
+        match = re.search(pattern, checkpoint)
+        cur_step = int(match.group(1))
+        epochs_trained = int(match.group(2))
+
+        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
+        logger.info(f"  Continuing training from epoch {epochs_trained}")
+        logger.info(f"  Continuing training from global step {cur_step}")
+
+        steps_trained_progress_bar.update(cur_step)
+
+        for epoch in range(0, epochs_trained):
+            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+
+        if not data_args.streaming and training_args.max_steps < 0:
+            # we know exactly the number of steps per epoch, so can skip through the required number of batches
+            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
+        else:
+            # Currently we don't know how many steps we've taken in the current epoch
+            # So we just shuffle the dataset one extra time and start from a fresh epoch
+            # This is "good enough" for our purposes but not fully correct
+            resume_step = None
+            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+    else:
+        resume_step = None
+
+    for epoch in range(epochs_trained, num_epochs):
+        vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
+        train_dataloader = DataLoader(
+            vectorized_datasets["train"],
+            collate_fn=data_collator,
+            batch_size=per_device_train_batch_size,
+            num_workers=dataloader_num_workers,
+            prefetch_factor=prefetch_factor,
+            pin_memory=training_args.dataloader_pin_memory,
+
+        )
+        train_dataloader = accelerator.prepare(train_dataloader)
+        if hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDataset):
+            train_dataloader.dataset.set_epoch(epoch)
+
+        if resume_step is not None:
+            # Skip the first N batches in the dataloader when resuming from a checkpoint
+            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            resume_step = None
+
+        for batch in train_dataloader:
+            with accelerator.accumulate(student_model):
+                loss, train_metric = train_step(batch, temperature=training_args.temperature)
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(student_model.parameters(), training_args.max_grad_norm)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            # Check if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                steps_trained_progress_bar.update(1)
+                cur_step += 1
+
+                if cur_step % training_args.logging_steps == 0:
+                    steps_trained_progress_bar.write(
+                        f"Step... ({cur_step} / {total_train_steps} | Loss:"
+                        f" {train_metric['loss']}, Learning Rate:"
+                        f" {lr_scheduler.get_last_lr()[0]})"
+                    )
+                    log_metric(
+                        accelerator,
+                        metrics=train_metric,
+                        learning_rate=lr_scheduler.get_last_lr()[0],
+                        train_time=train_time + time.time() - train_start,
+                        step=cur_step,
+                        epoch=epoch,
+                        prefix="train",
+                    )
+
+                # save checkpoint and weights after each save_steps and at the end of training
+                if (cur_step % training_args.save_steps == 0) or cur_step == total_train_steps:
+                    intermediate_dir = os.path.join(training_args.output_dir, f"checkpoint-{cur_step}-epoch-{epoch}")
+                    accelerator.save_state(output_dir=intermediate_dir)
+                    accelerator.wait_for_everyone()
+                    if accelerator.is_main_process:
+                        rotate_checkpoints(training_args.save_total_limit, output_dir=training_args.output_dir)
+
+                        if cur_step == total_train_steps:
+                            # un-wrap student model for save
+                            student_model = accelerator.unwrap_model(student_model)
+                            student_model.save_pretrained(training_args.output_dir)
+                            # re-wrap student model for final eval
+                            student_model = accelerator.prepare(student_model)
+
+                        if training_args.push_to_hub:
+                            upload_folder(
+                                folder_path=training_args.output_dir,
+                                repo_id=repo_name,
+                                repo_type="model",
+                                commit_message=f"Saving train state of step {cur_step}",
+                            )
+
+                if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
+                    train_time += time.time() - train_start
+                    student_model.eval()
+                    # ======================== Evaluating ==============================
+                    for eval_split in all_eval_splits:
+                        eval_metrics = []
+                        eval_preds = []
+                        eval_labels = []
+                        eval_start = time.time()
+
+                        validation_dataloader = DataLoader(
+                            vectorized_datasets[eval_split],
+                            collate_fn=data_collator,
+                            batch_size=per_device_eval_batch_size,
+                            drop_last=False,
+                            num_workers=dataloader_num_workers,
+                            prefetch_factor=prefetch_factor,
+                            pin_memory=training_args.dataloader_pin_memory,
+                        )
+                        validation_dataloader = accelerator.prepare(validation_dataloader)
+
+                        for batch in tqdm(
+                            validation_dataloader,
+                            desc=f"Evaluating {eval_split}...",
+                            position=2,
+                            disable=not accelerator.is_local_main_process,
+                        ):
+                            # Model forward
+                            eval_metric = eval_step(batch)
+                            eval_metric = accelerator.gather_for_metrics(eval_metric)
+                            eval_metrics.append(eval_metric)
+
+                            # generation
+                            if training_args.predict_with_generate:
+                                generated_ids = generate_step(batch)
+                                # Gather all predictions and targets
+                                generated_ids, labels = accelerator.gather_for_metrics(
+                                    (generated_ids, batch["labels"])
+                                )
+                                eval_preds.extend(generated_ids)
+                                eval_labels.extend(labels)
+
+                        eval_time = time.time() - eval_start
+                        # normalize eval metrics
+                        eval_metrics = {
+                            key: torch.mean(torch.stack([d[key] for d in eval_metrics])) for key in eval_metrics[0]
+                        }
+
+                        # compute WER metric
+                        wer_desc = ""
+                        if training_args.predict_with_generate:
+                            wer_metric, pred_str, label_str, norm_pred_str, norm_label_str = compute_metrics(
+                                eval_preds, eval_labels
+                            )
+                            eval_metrics.update(wer_metric)
+                            wer_desc = " ".join([f"Eval {key}: {value} |" for key, value in wer_metric.items()])
+                            log_pred(
+                                accelerator,
+                                pred_str,
+                                label_str,
+                                norm_pred_str,
+                                norm_label_str,
+                                step=cur_step,
+                                prefix=eval_split,
+                            )
+
+                        # Print metrics and update progress bar
+                        steps_trained_progress_bar.write(
+                            f"Eval results for step ({cur_step} / {total_train_steps} | Eval Loss: {eval_metrics['loss']} |"
+                            f" {wer_desc})"
+                        )
+
+                        log_metric(
+                            accelerator,
+                            metrics=eval_metrics,
+                            train_time=eval_time,
+                            step=cur_step,
+                            epoch=epoch,
+                            prefix=eval_split,
+                        )
+
+                    # flush the train metrics
+                    train_start = time.time()
+
+                # break condition
+                if cur_step == total_train_steps:
+                    continue_training = False
+                    break
+
+        if not continue_training:
+            break
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()
+

run_init.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+python create_student_model.py \
+  --teacher_checkpoint "openai/whisper-large-v3" \
+  --encoder_layers 32 \
+  --decoder_layers 2 \
+  --save_dir "./distil-large-v3-init"
+

run_labelling.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+
+accelerate launch run_pseudo_labelling.py \
+  --model_name_or_path "openai/whisper-large-v3" \
+  --dataset_name "mozilla-foundation/common_voice_16_1" \
+  --dataset_config_name "hi" \
+  --dataset_split_name "train+validation+test" \
+  --text_column_name "sentence" \
+  --id_column_name "path" \
+  --output_dir "../common_voice_16_1_de_pseudo_labelled" \
+  --wandb_project "distil-whisper-labelling" \
+  --per_device_eval_batch_size 64 \
+  --dtype "bfloat16" \
+  --attn_implementation "flash_attention_2" \
+  --logging_steps 500 \
+  --max_label_length 256 \
+  --concatenate_audio \
+  --preprocessing_batch_size 500 \
+  --preprocessing_num_workers 8 \
+  --dataloader_num_workers 8 \
+  --report_to "wandb" \
+  --language "de" \
+  --task "transcribe" \
+  --return_timestamps \
+  --streaming False \
+  --generation_num_beams 1 \
+  --push_to_hub
+

run_pseudo_labelling.py

@@ -0,0 +1,1015 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Pseudo-labelling audio data using the Whisper model in preparation for distillation.
+"""
+import csv
+
+# You can also adapt this script for your own pseudo-labelling tasks. Pointers for this are left as comments.
+import logging
+import os
+import string
+import sys
+import time
+import warnings
+from dataclasses import dataclass, field
+from datetime import timedelta
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import datasets
+import evaluate
+import numpy as np
+import torch
+import transformers
+from accelerate import Accelerator, InitProcessGroupKwargs
+from accelerate.logging import get_logger
+from datasets import (
+    DatasetDict,
+    IterableDatasetDict,
+    load_dataset,
+)
+from huggingface_hub import HfFolder, Repository, create_repo, get_full_repo_name
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import (
+    HfArgumentParser,
+    Seq2SeqTrainingArguments,
+    WhisperConfig,
+    WhisperFeatureExtractor,
+    WhisperForConditionalGeneration,
+    WhisperProcessor,
+    WhisperTokenizerFast,
+)
+from transformers.models.whisper.english_normalizer import EnglishTextNormalizer, BasicTextNormalizer
+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.34.0.dev0")
+
+require_version("datasets>=2.14.6", "To fix: `pip install --upgrade datasets`")
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to distill from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained Whisper model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained config name or path if not the same as model_name"},
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"},
+    )
+    feature_extractor_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "feature extractor name or path if not the same as model_name"},
+    )
+    processor_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "processor name or path if not the same as model_name"},
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    subfolder: str = field(
+        default="",
+        metadata={
+            "help": "In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can"
+            "specify the folder name here."
+        },
+    )
+    token: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
+                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
+            )
+        },
+    )
+    dtype: Optional[str] = field(
+        default="float32",
+        metadata={
+            "help": (
+                "The data type (dtype) in which to load the model weights. One of `float32` (full-precision), "
+                "`float16` or `bfloat16` (both half-precision)."
+            )
+        },
+    )
+    attn_implementation: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+            "Which attention implementation to use in the encoder and decoder attention layers. Can be one of:\n"
+            "1. `eager` or `None`: default Transformers attention implementation.\n"
+            "2. `sdpa`: Flash Attention through PyTorch SDPA. Requires `torch>=2.1`. Recommended for hardware where Flash Attention 2 is not supported, e.g. Turing GPUs, (T4, RTX 2080).\n"
+            "3. `flash_attn_2`: Flash Attention 2 through the Flash Attention package https://github.com/Dao-AILab/flash-attention. **Always** recommended on supported hardware (Ampere, Ada, or Hopper GPUs, e.g., A100, RTX 3090, RTX 4090, H100)."
+        )
+        },
+    )
+    attn_type: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "Deprecated. Use `attn_implementation` instead."
+        },
+    )
+    def __post_init__(self):
+        if self.attn_type is not None and self.attn_implementation is None:
+            # set attn_implementation in a backwards compatible way
+            if self.attn_type == "flash_attn":
+                self.attn_implementation = "sdpa"
+            elif self.attn_type == "flash_attn_2":
+                self.attn_implementation = "flash_attention_2"
+            elif self.attn_type in [None, "eager", "sdpa", "flash_attention_2"]:
+                self.attn_implementation = self.attn_type
+            else:
+                raise ValueError(
+                    f"Argument `--attn_type` is deprecated, and set to an invalid option `{self.attn_type}`. You should omit the argument `--attn_type`, and instead set `-attention_implementation` to one of the following:\n"
+                    "1. `eager` or `None`: default Transformers attention implementation.\n"
+                    "2. `sdpa`: Flash Attention through PyTorch SDPA. Requires `torch>=2.1`. Recommended for hardware where Flash Attention 2 is not supported, e.g. Turing GPUs, (T4, RTX 2080).\n"
+                    "3. `flash_attn_2`: Flash Attention 2 through the Flash Attention package https://github.com/Dao-AILab/flash-attention. **Always** recommended on supported hardware (Ampere, Ada, or Hopper GPUs, e.g., A100, RTX 3090, RTX 4090, H100)."
+                )
+            warnings.warn(f"Argument `--attn_type` is deprecated. Use `--attn_implementation` instead. Inferring `--attn_implementation={self.attn_implementation} from argument `--attn_type={self.attn_type}`.")
+        elif self.attn_type is not None and self.attn_implementation is not None:
+            raise ValueError("`--attn_type` and `--attn_implementation` are both specified. Only the argument `--attn_implementation`.")
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    dataset_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset to use (via the datasets library)."},
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "The configuration name of the dataset to use (via the datasets library)."},
+    )
+    dataset_cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to cache directory for saving and loading datasets"},
+    )
+    overwrite_cache: bool = field(
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets"},
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    preprocessing_batch_size: Optional[int] = field(
+        default=500,
+        metadata={"help": "The batch size to use for the dataset pre-processing."},
+    )
+    audio_column_name: str = field(
+        default="audio",
+        metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},
+    )
+    text_column_name: str = field(
+        default="text",
+        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'."},
+    )
+    id_column_name: str = field(
+        default="id",
+        metadata={"help": "The name of the dataset column containing the id data. Defaults to 'id'"},
+    )
+    speaker_id_column_name: str = field(
+        default=None,
+        metadata={"help": "The name of the dataset column containing the speaker id data. Defaults to None."},
+    )
+    max_duration_in_seconds: float = field(
+        default=30.0,
+        metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"},
+    )
+    max_label_length: int = field(
+        default=256,
+        metadata={"help": "Truncate transcriptions that are longer `max_label_length` tokens."},
+    )
+    concatenate_audio: bool = field(
+        default=True,
+        metadata={"help": "Whether or not to concatenate the audio samples to `max_duration_in_seconds`."},
+    )
+    preprocessing_only: bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to only do data preprocessing and skip training. This is"
+                " especially useful when data preprocessing errors out in distributed"
+                " training due to timeout. In this case, one should run the"
+                " preprocessing in a non-distributed setup with"
+                " `preprocessing_only=True` so that the cached datasets can"
+                " consequently be loaded in distributed training"
+            )
+        },
+    )
+    dataset_split_name: str = field(
+        default="train+validation+test",
+        metadata={
+            "help": (
+                "The name of the data set splits to use (via the datasets library)."
+                " Defaults to 'train+validation+test'. Multiple splits can be passed by splitting a"
+                " list through the '+' character, e.g. 'train+validation' will"
+                " pseudo-label both the 'train' and 'validation' splits sequentially."
+            )
+        },
+    )
+    wandb_project: str = field(
+        default="distil-whisper",
+        metadata={"help": "The name of the wandb project."},
+    )
+    streaming: bool = field(
+        default=False,
+        metadata={"help": "Whether to use dataset's streaming mode to load and pre-process the data."},
+    )
+    max_samples_per_split: Optional[int] = field(
+        default=None,
+        metadata={"help": "For debugging purposes, truncate the number of examples per split to this value if set."},
+    )
+    return_timestamps: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to return the timestamps with the text. This enables the `FlaxWhisperTimestampsLogitsProcessor`."
+        },
+    )
+    language: str = field(
+        default=None,
+        metadata={
+            "help": (
+                "Language for multilingual distillation. This argument should be set for multilingual distillation "
+                "only. For English speech recognition, it should be left as `None`."
+            )
+        },
+    )
+    task: str = field(
+        default="transcribe",
+        metadata={
+            "help": "Task, either `transcribe` for speech recognition or `translate` for speech translation."
+            "This argument should be set for multilingual distillation only. For English speech recognition, it should be left as `None`."
+        },
+    )
+    decode_token_ids: bool = field(
+        default=True,
+        metadata={"help": "Deprecated. The predicted token ids should always be decoded to text transcriptions."},
+    )
+    private_dataset: bool = field(
+        default=False,
+        metadata={"help": "Whether or not to create a private dataset for the pseudo-labelled data."},
+    )
+
+    def __post_init__(self):
+        if not self.decode_token_ids:
+            raise ValueError(
+                "The argument `--decode_token_ids` is deprecated. The token ids are now always decoded to "
+                "their corresponding text string. This is following a fix to the merges of the Whisper tokenizer"
+                "on the Hugging Face Hub: https://huggingface.co/openai/whisper-large-v2/discussions/100. "
+                "You should either omit the argument `--decode_token_ids`, or set it to True explicitly."
+            )
+
+def shift_tokens_right(label_ids: np.array, decoder_start_token_id: int) -> np.ndarray:
+    """
+    Shift label ids one token to the right.
+    """
+    shifted_label_ids = np.zeros_like(label_ids)
+    shifted_label_ids[:, 1:] = label_ids[:, :-1]
+    shifted_label_ids[:, 0] = decoder_start_token_id
+
+    return shifted_label_ids
+
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    """
+    Data collator that will dynamically pad the inputs received.
+    Args:
+        processor ([`Wav2Vec2Processor`])
+            The processor used for proccessing the data.
+        decoder_start_token_id (:obj: `int`)
+            The start-of-sequence token id of the decoder.
+        input_padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned input sequences (according to the model's padding side and padding index)
+            among:
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        target_padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned target sequences (according to the model's padding side and padding index).
+            See above for details.
+        max_target_length (:obj:`int`, `optional`):
+            Maximum length of the ``labels`` of the returned list and optionally padding length (see above).
+    """
+
+    processor: Any
+    decoder_start_token_id: int
+    input_padding: Union[bool, str] = "max_length"
+    target_padding: Union[bool, str] = "max_length"
+    max_target_length: Optional[int] = None
+
+    def __call__(self, features: List[Dict[str, Union[List[int], np.ndarray]]]) -> Dict[str, np.ndarray]:
+        # split inputs and labels since they have to be of different lengths and need
+        # different padding methods
+        model_input_name = self.processor.model_input_names[0]
+
+        # dataloader returns a list of features which we convert to a dict
+        input_features = {model_input_name: [feature[model_input_name] for feature in features]}
+        label_features = {"input_ids": [feature["labels"] for feature in features]}
+        file_ids = [feature["file_id"] for feature in features]
+
+        # reformat list to dict and set to pytorch format
+        batch = self.processor.feature_extractor.pad(
+            input_features,
+            padding=self.input_padding,
+            return_tensors="pt",
+        )
+
+        labels_batch = self.processor.tokenizer.pad(
+            label_features,
+            max_length=self.max_target_length,
+            padding=self.target_padding,
+            return_tensors="pt",
+        )
+
+        # replace padding with -100 to ignore correctly when computing the loss
+        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 (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+        batch["file_ids"] = file_ids
+
+        return batch
+
+
+def log_metric(
+    accelerator,
+    metrics: Dict,
+    train_time: float,
+    prefix: str = "eval",
+):
+    """Helper function to log all evaluation metrics with the correct prefixes and styling."""
+    log_metrics = {}
+    for k, v in metrics.items():
+        log_metrics[f"{prefix}/{k}"] = v
+    log_metrics[f"{prefix}/time"] = train_time
+    accelerator.log(log_metrics)
+
+
+def log_pred(
+    accelerator,
+    pred_str: List[str],
+    label_str: List[str],
+    norm_pred_str: List[str],
+    norm_label_str: List[str],
+    prefix: str = "eval",
+    num_lines: int = 200000,
+):
+    """Helper function to log target/predicted transcriptions to weights and biases (wandb)."""
+    if accelerator.is_main_process:
+        wandb_tracker = accelerator.get_tracker("wandb")
+        # pretty name for split
+        prefix = prefix.replace("/", "-")
+
+        # convert str data to a wandb compatible format
+        str_data = [[label_str[i], pred_str[i], norm_label_str[i], norm_pred_str[i]] for i in range(len(pred_str))]
+        # log as a table with the appropriate headers
+        wandb_tracker.log_table(
+            table_name=f"{prefix}/all_predictions",
+            columns=["Target", "Pred", "Norm Target", "Norm Pred"],
+            data=str_data[:num_lines],
+        )
+
+        # log incorrect normalised predictions
+        str_data = np.asarray(str_data)
+        str_data_incorrect = str_data[str_data[:, -2] != str_data[:, -1]]
+        # log as a table with the appropriate headers
+        wandb_tracker.log_table(
+            table_name=f"{prefix}/incorrect_predictions",
+            columns=["Target", "Pred", "Norm Target", "Norm Pred"],
+            data=str_data_incorrect[:num_lines],
+        )
+
+
+def main():
+    # 1. Parse input arguments
+    # We keep distinct sets of args, for cleaner separation of model/data/training related args
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # 2. Initialize the accelerator
+    # We will let the accelerator handle device placement for us in this example
+    # We simply have to specify the training precision and any trackers being used
+    # We'll use the same dtype arguments as our JAX/Flax training script and convert
+    # it to accelerate format
+    if model_args.dtype == "float16":
+        mixed_precision = "fp16"
+        torch_dtype = torch.float16
+    elif model_args.dtype == "bfloat16":
+        mixed_precision = "bf16"
+        torch_dtype = torch.bfloat16
+    else:
+        mixed_precision = "no"
+        torch_dtype = torch.float32
+
+    kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=7200))
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
+        mixed_precision=mixed_precision,
+        log_with=training_args.report_to,
+        project_dir=training_args.output_dir,
+        kwargs_handlers=[kwargs],
+    )
+
+    accelerator.init_trackers(project_name=data_args.wandb_project)
+
+    # 3. Set-up basic logging
+    # Create one log on every process with the configuration for debugging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    # Log a small summary on each proces
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
+        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set the verbosity to info of the Transformers logger (on main process only)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # 3. Load dataset
+    raw_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict()
+    token = model_args.token if model_args.token is not None else HfFolder().get_token()
+
+    data_splits = data_args.dataset_split_name.split("+")
+    for split in data_splits:
+        if data_args.streaming:
+            raw_datasets[split] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=split,
+                cache_dir=data_args.dataset_cache_dir,
+                token=token,
+                streaming=True,
+            )
+        else:
+            raw_datasets[split] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=split,
+                cache_dir=data_args.dataset_cache_dir,
+                token=token,
+                streaming=False,
+                num_proc=data_args.preprocessing_num_workers,
+            )
+
+    if data_args.audio_column_name not in next(iter(raw_datasets.values())).column_names:
+        raise ValueError(
+            f"--audio_column_name '{data_args.audio_column_name}' not found in dataset"
+            f" '{data_args.dataset_name}'. Make sure to set `--audio_column_name` to"
+            " the correct audio column - one of"
+            f" {', '.join(next(iter(raw_datasets.values())).column_names)}."
+        )
+
+    if data_args.text_column_name not in next(iter(raw_datasets.values())).column_names:
+        raise ValueError(
+            f"--text_column_name {data_args.text_column_name} not found in dataset"
+            f" '{data_args.dataset_name}'. Make sure to set `--text_column_name` to the"
+            " correct text column - one of"
+            f" {', '.join(next(iter(raw_datasets.values())).column_names)}."
+        )
+
+    # 7. Load pretrained model, tokenizer, and feature extractor
+    config = WhisperConfig.from_pretrained(
+        (model_args.config_name if model_args.config_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        token=token,
+    )
+    feature_extractor = WhisperFeatureExtractor.from_pretrained(
+        (model_args.feature_extractor_name if model_args.feature_extractor_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        token=token,
+    )
+    tokenizer = WhisperTokenizerFast.from_pretrained(
+        (model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+        token=token,
+    )
+    processor = WhisperProcessor.from_pretrained(
+        (model_args.processor_name if model_args.processor_name else model_args.model_name_or_path),
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        token=token,
+    )
+
+    model = WhisperForConditionalGeneration.from_pretrained(
+        model_args.model_name_or_path,
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        subfolder=model_args.subfolder,
+        token=token,
+        low_cpu_mem_usage=True,
+        torch_dtype=torch_dtype,
+        attn_implementation=model_args.attn_implementation,
+    )
+    model.eval()
+
+    if model.config.decoder_start_token_id is None:
+        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
+
+    return_timestamps = data_args.return_timestamps
+    if hasattr(model.generation_config, "is_multilingual") and model.generation_config.is_multilingual:
+        is_multilingual = True
+        # We need to set the language and task ids for multilingual checkpoints
+        tokenizer.set_prefix_tokens(
+            language=data_args.language, task=data_args.task, predict_timestamps=return_timestamps
+        )
+    elif data_args.language is not None:
+        raise ValueError(
+            "Setting language token for an English-only checkpoint is not permitted. The language argument should "
+            "only be set for multilingual checkpoints."
+        )
+    else:
+        is_multilingual = False
+
+    # 6. Resample speech dataset: `datasets` takes care of automatically loading and resampling the audio,
+    # so we just need to set the correct target sampling rate.
+    raw_datasets = raw_datasets.cast_column(
+        data_args.audio_column_name,
+        datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate),
+    )
+
+    # 7. Preprocessing the datasets.
+    # We need to read the audio files as arrays and tokenize the targets.
+    max_input_length = int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate)
+    max_label_length = (
+        data_args.max_label_length if data_args.max_label_length is not None else model.config.max_length
+    )
+    audio_column_name = data_args.audio_column_name
+    sampling_rate = feature_extractor.sampling_rate
+
+    preprocessing_batch_size = data_args.preprocessing_batch_size
+    num_workers = data_args.preprocessing_num_workers
+    dataloader_num_workers = training_args.dataloader_num_workers
+
+    text_column_name = data_args.text_column_name
+    model_input_name = feature_extractor.model_input_names[0]
+    id_column_name = data_args.id_column_name
+    speaker_id_column_name = data_args.speaker_id_column_name
+    normalizer = (
+        BasicTextNormalizer() if data_args.language is not None else EnglishTextNormalizer(tokenizer.english_spelling_normalizer)
+    )
+
+    timestamp_position = 3 if is_multilingual else 1
+    decoder_prev_token_id = tokenizer.convert_tokens_to_ids("<|startofprev|>")
+    decoder_eot_token_id = tokenizer.eos_token_id
+
+    if data_args.max_samples_per_split is not None:
+        for split in data_splits:
+            raw_datasets[split] = (
+                raw_datasets[split].take(data_args.max_samples_per_split)
+                if data_args.streaming
+                else raw_datasets[split].select(range(data_args.max_samples_per_split))
+            )
+
+    if speaker_id_column_name is not None:
+        raw_datasets = raw_datasets.sort(speaker_id_column_name)
+
+    def concatenate_dataset(batch):
+        audio = [sample["array"] for sample in batch[audio_column_name]]
+        input_lengths = [len(sample) for sample in audio]
+
+        text = batch[text_column_name]
+        speaker_id = batch[speaker_id_column_name] if speaker_id_column_name else len(text) * [None]
+
+        concatenated_audio = []
+        concatenated_text = []
+        concatenated_speaker = []
+        condition_on_prev = []
+        audio_sample = audio[0]
+        text_sample = text[0]
+
+        for idx in range(1, len(audio)):
+            prev_speaker = speaker_id[idx - 1]
+            speaker = speaker_id[idx]
+
+            if len(audio_sample) + input_lengths[idx] < max_input_length:
+                if speaker == prev_speaker:
+                    # we have no information about whether the segments follow on sequentially
+                    # so we just ensure the same speaker as we concatenate across files
+                    audio_sample = np.append(audio_sample, audio[idx])
+                    # extra spaces in the text transcription don't matter, since we only use it for the WER computation
+                    text_sample += " " + text[idx]
+                else:
+                    # speakers do not follow sequentially, save the audio and start looping again
+                    concatenated_audio.append(audio_sample)
+                    concatenated_text.append(text_sample)
+                    concatenated_speaker.append(speaker)
+                    condition_on_prev.append(0)
+                    audio_sample = audio[idx]
+                    text_sample = text[idx]
+
+            else:
+                # concatenated audio exceeds max length, save the audio and start looping again
+                concatenated_audio.append(audio_sample)
+                concatenated_text.append(text_sample)
+                concatenated_speaker.append(speaker)
+                condition_on_prev.append(1)
+                audio_sample = audio[idx]
+                text_sample = text[idx]
+
+        batch[audio_column_name] = [{"array": array, "sampling_rate": sampling_rate} for array in concatenated_audio]
+        batch[text_column_name] = concatenated_text
+        batch[id_column_name] = concatenated_speaker
+        batch["condition_on_prev"] = condition_on_prev
+
+        return batch
+
+    raw_datasets_features = list(next(iter(raw_datasets.values())).features.keys())
+    if data_args.concatenate_audio and not data_args.streaming:
+        raw_datasets = raw_datasets.map(
+            concatenate_dataset,
+            batched=True,
+            batch_size=preprocessing_batch_size,
+            num_proc=num_workers,
+            remove_columns=set(raw_datasets_features) - {audio_column_name, text_column_name, id_column_name, "condition_on_prev"},
+            desc="Concatenating dataset...",
+        )
+
+        raw_datasets = raw_datasets.cast_column(audio_column_name, datasets.features.Audio(sampling_rate=sampling_rate))
+        pretty_name = data_args.dataset_name.split("/")[-1]
+
+        def postprocess_ids(speaker_ids, indices):
+            speaker_ids_formatted = []
+            for speaker, idx in zip(speaker_ids, indices):
+                formatted_idx = f"{pretty_name}-{speaker}-{idx}" if speaker is not None else f"{pretty_name}-{idx}"
+                speaker_ids_formatted.append(formatted_idx)
+            return {id_column_name: speaker_ids_formatted}
+
+        raw_datasets = raw_datasets.map(
+            postprocess_ids,
+            input_columns=[id_column_name],
+            with_indices=True,
+            desc="Setting sample idxs...",
+            batched=True,
+            batch_size=preprocessing_batch_size,
+            num_proc=num_workers,
+        )
+    else:
+        raise ValueError(
+            "Streaming mode is not yet compatible with concatenating audios to `max_duration_in_seconds`."
+            "Either set `--streaming=False` and download the audios locally, or open an issue on the Distil-Whisper repo to request this feature."
+        )
+
+    def prepare_dataset(batch):
+        # process audio
+        sample = batch[audio_column_name]
+        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
+        # process audio length
+        batch[model_input_name] = inputs.get(model_input_name)[0]
+
+        # process targets
+        input_str = batch[text_column_name]
+        batch["labels"] = tokenizer(input_str, max_length=max_label_length, truncation=True).input_ids
+
+        # record the id of the sample as token ids
+        batch["file_id"] = batch[id_column_name]
+        return batch
+
+    raw_datasets_features = list(next(iter(raw_datasets.values())).features.keys())
+    if data_args.streaming:
+        vectorized_datasets = raw_datasets.map(prepare_dataset, remove_columns=raw_datasets_features)
+    else:
+        vectorized_datasets = raw_datasets.map(
+            prepare_dataset,
+            remove_columns=raw_datasets_features,
+            num_proc=num_workers,
+            desc="preprocess dataset",
+        )
+
+    # for large datasets it is advised to run the preprocessing on a
+    # single machine first with `args.preprocessing_only` since there will mostly likely
+    # be a timeout when running the script in distributed mode.
+    # In a second step `args.preprocessing_only` can then be set to `False` to load the
+    # cached dataset
+    if data_args.preprocessing_only:
+        cache = {k: v.cache_files for k, v in vectorized_datasets.items()}
+        logger.info(f"Data preprocessing finished. Files cached at {cache}.")
+        return
+
+    if data_args.streaming and dataloader_num_workers > 0:
+        logger.warning(
+            "Using multiple dataloader num workers with streaming mode will result in different shards of "
+            "data being transcribed in parallel. This is not advised if you want to preserve the order of the "
+            "audio-text data."
+        )
+
+    # Handle the repository creation
+    output_dir = training_args.output_dir
+    if training_args.push_to_hub:
+        if training_args.hub_model_id is None:
+            repo_name = get_full_repo_name(
+                Path(output_dir).absolute().name,
+                token=token,
+            )
+        else:
+            repo_name = training_args.hub_model_id
+        create_repo(repo_name, exist_ok=True, token=token, repo_type="dataset", private=data_args.private_dataset)
+        repo = Repository(
+            output_dir,
+            clone_from=repo_name,
+            token=token,
+            repo_type="dataset",
+        )
+        # Ensure large txt files can be pushed to the Hub with git-lfs
+        with open(os.path.join(output_dir, ".gitattributes"), "r+") as f:
+            git_lfs_extensions = f.read()
+            if "*.csv" not in git_lfs_extensions:
+                f.write("*.csv filter=lfs diff=lfs merge=lfs -text")
+    else:
+        # this is where we'll save our transcriptions
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+    # 8. Load Metric
+    metric = evaluate.load("wer")
+
+    def compute_metrics(preds, labels, file_ids):
+        # replace padded labels by the padding token
+        for idx in range(len(labels)):
+            labels[idx][labels[idx] == -100] = tokenizer.pad_token_id
+
+        pred_str = tokenizer.batch_decode(preds, skip_special_tokens=False, decode_with_timestamps=return_timestamps)
+        # we do not want to group tokens when computing the metrics
+        label_str = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        # normalize everything and re-compute the WER
+        norm_pred_str = [normalizer(pred) for pred in pred_str]
+        norm_label_str = [normalizer(label) for label in label_str]
+        # for logging, we need the pred/labels to match the norm_pred/norm_labels, so discard any filtered samples here
+        pred_str = [pred_str[i] for i in range(len(norm_pred_str)) if len(norm_label_str[i]) > 0]
+        label_str = [label_str[i] for i in range(len(norm_label_str)) if len(norm_label_str[i]) > 0]
+        file_ids = [file_ids[i] for i in range(len(file_ids)) if len(norm_label_str[i]) > 0]
+        # filtering step to only evaluate the samples that correspond to non-zero normalized references:
+        norm_pred_str = [norm_pred_str[i] for i in range(len(norm_pred_str)) if len(norm_label_str[i]) > 0]
+        norm_label_str = [norm_label_str[i] for i in range(len(norm_label_str)) if len(norm_label_str[i]) > 0]
+
+        wer = 100 * metric.compute(predictions=norm_pred_str, references=norm_label_str)
+
+        return {"wer": wer}, pred_str, label_str, norm_pred_str, norm_label_str, file_ids
+
+    def filter_eot_tokens(preds):
+        for idx in range(len(preds)):
+            # remove the EOT tokens to get the 'true' token length
+            token_ids = [token for token in preds[idx] if token != decoder_eot_token_id]
+            token_ids = token_ids + [decoder_eot_token_id]
+            preds[idx] = token_ids
+        return preds
+
+    # 12. Define Training Schedule
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+
+    data_collator = DataCollatorSpeechSeq2SeqWithPadding(
+        processor=processor,
+        decoder_start_token_id=model.config.decoder_start_token_id,  # <|startoftranscript|>
+        input_padding="longest",
+        target_padding="max_length",
+        max_target_length=max_label_length,
+    )
+
+    # 14. Define generation arguments - we need to do this before we wrap the models in DDP
+    # so that we can still access the configs
+    num_beams = (
+        training_args.generation_num_beams
+        if training_args.generation_num_beams is not None
+        else getattr(model.generation_config, "num_beams", 1)
+    )
+
+    gen_kwargs = {
+        "max_length": max_label_length,
+        "num_beams": num_beams,
+        "return_timestamps": return_timestamps,
+    }
+    if hasattr(model.generation_config, "is_multilingual") and model.generation_config.is_multilingual:
+        # forcing the language and task tokens helps multilingual models in their generations
+        gen_kwargs.update(
+            {
+                "language": data_args.language,
+                "task": data_args.task,
+            }
+        )
+    model.generation_config.forced_decoder_ids = None
+
+    # 15. Prepare everything with accelerate
+    model = accelerator.prepare(model)
+
+    def eval_step_with_save(split="eval"):
+        # ======================== Evaluating ==============================
+        eval_preds = []
+        eval_labels = []
+        eval_ids = []
+        pred_str = []
+        eval_start = time.time()
+
+        eval_loader = DataLoader(
+            vectorized_datasets[split],
+            batch_size=per_device_eval_batch_size,
+            collate_fn=data_collator,
+            num_workers=dataloader_num_workers,
+            pin_memory=True,
+        )
+
+        eval_loader = accelerator.prepare(eval_loader)
+        batches = tqdm(eval_loader, desc=f"Evaluating {split}...", disable=not accelerator.is_local_main_process)
+
+        # make the split name pretty for librispeech etc
+        split = split.replace(".", "-").split("/")[-1]
+        output_csv = os.path.join(output_dir, f"{split}-transcription.csv")
+
+        for step, batch in enumerate(batches):
+            file_ids = batch.pop("file_ids")
+            # Generate predictions and pad to max generated length
+            generate_fn = model.module.generate if accelerator.num_processes > 1 else model.generate
+            generated_ids = generate_fn(batch["input_features"].to(dtype=torch_dtype), **gen_kwargs)
+            generated_ids = accelerator.pad_across_processes(generated_ids, dim=1, pad_index=tokenizer.pad_token_id)
+            # Gather all predictions and targets
+            file_ids, generated_ids, labels = accelerator.gather_for_metrics(
+                (file_ids, generated_ids, batch["labels"])
+            )
+            eval_preds.extend(generated_ids.cpu().numpy())
+            eval_labels.extend(labels.cpu().numpy())
+            eval_ids.extend(file_ids)
+
+            if step % training_args.logging_steps == 0 and step > 0:
+                batches.write(f"Saving transcriptions for split {split} step {step}")
+                accelerator.wait_for_everyone()
+                pred_ids = eval_preds[-(len(eval_preds) - len(pred_str)):]
+                pred_ids = filter_eot_tokens(pred_ids)
+                pred_str.extend(
+                    tokenizer.batch_decode(pred_ids, skip_special_tokens=False,decode_with_timestamps=return_timestamps)
+                )
+                csv_data = [[eval_ids[i], pred_str[i]] for i in range(len(eval_preds))]
+
+                with open(output_csv, "w", encoding="UTF8", newline="") as f:
+                    writer = csv.writer(f)
+                    # write multiple rows
+                    writer.writerow(["file_id", "whisper_transcript"])
+                    writer.writerows(csv_data)
+
+                if training_args.push_to_hub and accelerator.is_main_process:
+                    repo.push_to_hub(
+                        commit_message=f"Saving transcriptions for split {split} step {step}.",
+                        blocking=False,
+                    )
+
+        accelerator.wait_for_everyone()
+        eval_time = time.time() - eval_start
+
+        # compute WER metric for eval sets
+        wer_desc = ""
+        if "validation" in split or "test" in split:
+            eval_preds = filter_eot_tokens(eval_preds)
+            wer_metric, pred_str, label_str, norm_pred_str, norm_label_str, eval_ids = compute_metrics(
+                eval_preds, eval_labels, eval_ids
+            )
+            wer_desc = " ".join([f"Eval {key}: {value} |" for key, value in wer_metric.items()])
+            # Save metrics + predictions
+            log_metric(
+                accelerator,
+                metrics=wer_metric,
+                train_time=eval_time,
+                prefix=split,
+            )
+            log_pred(
+                accelerator,
+                pred_str,
+                label_str,
+                norm_pred_str,
+                norm_label_str,
+                prefix=split,
+            )
+        else:
+            pred_ids = eval_preds[-(len(eval_preds) - len(pred_str)):]
+            pred_ids = filter_eot_tokens(pred_ids)
+            pred_str.extend(
+                tokenizer.batch_decode(pred_ids, skip_special_tokens=False, decode_with_timestamps=return_timestamps)
+            )
+
+        batches.write(f"Saving final transcriptions for split {split}.")
+        csv_data = [[eval_ids[i], eval_preds[i]] for i in range(len(eval_preds))]
+        with open(output_csv, "w", encoding="UTF8", newline="") as f:
+            writer = csv.writer(f)
+            # write multiple rows
+            writer.writerow(["file_id", "whisper_transcript"])
+            writer.writerows(csv_data)
+
+        # Print metrics
+        logger.info(wer_desc)
+
+        if not data_args.streaming and accelerator.is_main_process:
+            raw_datasets[split] = raw_datasets[split].add_column("whisper_transcript", pred_str)
+            raw_datasets[split] = raw_datasets[split].add_column("eval_preds", eval_preds)
+
+            def add_concatenated_text(eval_preds, condition_on_prev):
+                concatenated_prev = [None]
+                for token_ids, condition in zip(eval_preds[:-1], condition_on_prev[1:]):
+                    if condition is False:
+                        concatenated_prev.append(None)
+                    else:
+                        prompt_ids = [token for token in token_ids if token != decoder_eot_token_id]
+                        prompt_ids = [decoder_prev_token_id] + prompt_ids[timestamp_position:]
+                        concatenated_prev.append(prompt_ids)
+                return {"condition_on_prev": concatenated_prev}
+
+            raw_datasets[split] = raw_datasets[split].map(
+                add_concatenated_text,
+                input_columns=["eval_preds", "condition_on_prev"],
+                remove_columns=["eval_preds"],
+                desc="Setting condition on prev...",
+                batched=True,
+                batch_size=preprocessing_batch_size,
+                num_proc=num_workers,
+            )
+
+    logger.info("***** Running Labelling *****")
+    logger.info("  Instantaneous batch size per device =" f" {training_args.per_device_eval_batch_size}")
+    logger.info(
+        f"  Total eval batch size (w. parallel & distributed) = {training_args.per_device_eval_batch_size * accelerator.num_processes}"
+    )
+    logger.info(f"  Predict labels with timestamps = {return_timestamps}")
+    for split in data_splits:
+        eval_step_with_save(split=split)
+        accelerator.wait_for_everyone()
+        if training_args.push_to_hub and accelerator.is_main_process:
+            repo.push_to_hub(
+                commit_message=f"Saving final transcriptions for split {split.replace('.', '-').split('/')[-1]}",
+                blocking=False,
+            )
+    if not data_args.streaming and accelerator.is_main_process:
+        raw_datasets.save_to_disk(output_dir, num_proc=num_workers)
+        if training_args.push_to_hub:
+            raw_datasets.push_to_hub(repo_name, config_name=data_args.dataset_config_name)
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()