Merge branch 'main' of https://huggingface.co/NbAiLab/whisper-flaxtest into main

backup_run_flax_speech_recognition_seq2seq.py

@@ -1,880 +0,0 @@
-#!/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.
-"""
-Fine-tuning the Flax library models for sequence to sequence speech recognition.
-"""
-# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
-
-import logging
-import math
-import os
-import sys
-import time
-from dataclasses import field
-from functools import partial
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import datasets
-import flax
-import jax
-import jax.numpy as jnp
-import numpy as np
-import optax
-from datasets import Dataset, DatasetDict, load_dataset, load_metric
-from flax import jax_utils, traverse_util
-from flax.jax_utils import pad_shard_unpad, unreplicate
-from flax.training import train_state
-from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
-from huggingface_hub import Repository, create_repo
-from tqdm import tqdm
-
-import transformers
-from transformers import (
-    AutoConfig,
-    AutoFeatureExtractor,
-    AutoProcessor,
-    AutoTokenizer,
-    FlaxAutoModelForSpeechSeq2Seq,
-    HfArgumentParser,
-    Seq2SeqTrainingArguments,
-    is_tensorboard_available,
-)
-from transformers.file_utils import get_full_repo_name
-from transformers.utils import check_min_version, send_example_telemetry
-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.27.0.dev0")
-
-require_version("datasets>=1.18.0", "To fix: pip install -r examples/flax/speech-recogintion/requirements.txt")
-
-logger = logging.getLogger(__name__)
-
-
-@flax.struct.dataclass
-class ModelArguments:
-    """
-    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
-    """
-
-    model_name_or_path: str = field(
-        metadata={"help": "Path to pretrained 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)."},
-    )
-    use_auth_token: bool = field(
-        default=False,
-        metadata={
-            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
-            "with private models)."
-        },
-    )
-    dtype: Optional[str] = field(
-        default="float32",
-        metadata={
-            "help": (
-                "Floating-point format in which the model weights should be initialized and trained. Choose one of"
-                " `[float32, float16, bfloat16]`."
-            )
-        },
-    )
-    num_beams: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": (
-                "Number of beams to use for evaluation. This argument will be passed to `model.generate`, "
-                "which is used during evaluation."
-            )
-        },
-    )
-
-
-@flax.struct.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)."}
-    )
-    text_column: Optional[str] = field(
-        default=None,
-        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
-    )
-    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."},
-    )
-    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="text",
-        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},
-    )
-    max_duration_in_seconds: float = field(
-        default=20.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: float = field(
-        default=128,
-        metadata={"help": "Truncate transcriptions that are longer `max_eval_length` tokens."},
-    )
-    pad_input_to_multiple_of: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "If set will pad the input sequence to a multiple of the provided value. "
-            "This is important to avoid triggering recompilations on TPU. If unspecified, will default to padding the inputs to max length."
-        },
-    )
-    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'"
-        },
-    )
-    do_lower_case: bool = field(
-        default=True,
-        metadata={"help": "Whether the target text should be lower cased."},
-    )
-    language: str = field(
-        default=None,
-        metadata={
-            "help": (
-                "Language for multilingual fine-tuning. This argument should be set for multilingual fine-tuning "
-                "only. For English speech recognition, it should be set to `None`."
-            )
-        },
-    )
-    task: str = field(
-        default="transcribe",
-        metadata={"help": "Task, either `transcribe` for speech recognition or `translate` for speech translation."},
-    )
-
-
-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
-
-
-@flax.struct.dataclass
-class FlaxDataCollatorSpeechSeq2SeqWithPadding:
-    """
-    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 begin-of-sentence 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_input_length (:obj:`float`, `optional`):
-            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).
-        max_target_length (:obj:`int`, `optional`):
-            Maximum length of the ``labels`` of the returned list and optionally padding length (see above).
-        pad_input_to_multiple_of (:obj:`int`, `optional`):
-            If set will pad the input sequence to a multiple of the provided value.
-            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
-            7.5 (Volta).
-        pad_target_to_multiple_of (:obj:`int`, `optional`):
-            If set will pad the target sequence to a multiple of the provided value.
-            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
-            7.5 (Volta).
-    """
-
-    processor: Any
-    decoder_start_token_id: int
-    input_padding: Union[bool, str] = "longest"
-    target_padding: Union[bool, str] = "max_length"
-    max_input_length: Optional[float] = None
-    max_target_length: Optional[int] = None
-    pad_input_to_multiple_of: Optional[int] = None
-    pad_target_to_multiple_of: 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]
-        input_features = {model_input_name: features[model_input_name]}
-        label_features = {"input_ids": features["labels"]}
-
-        # reformat list to dict and set to pytorch format
-        batch = self.processor.feature_extractor.pad(
-            input_features,
-            max_length=self.max_input_length,
-            padding=self.input_padding,
-            pad_to_multiple_of=self.pad_input_to_multiple_of,
-            return_tensors="np",
-        )
-
-        labels_batch = self.processor.tokenizer.pad(
-            label_features,
-            max_length=self.max_target_length,
-            padding=self.target_padding,
-            pad_to_multiple_of=self.pad_target_to_multiple_of,
-            return_tensors="np",
-        )
-
-        # if bos token is appended in previous tokenization step,
-        # cut bos token here as it's append later anyways
-        labels = labels_batch["input_ids"]
-        if (labels[:, 0] == self.decoder_start_token_id).all().item():
-            labels = labels[:, 1:]
-            labels_batch.attention_mask = labels_batch.attention_mask[:, 1:]
-
-        decoder_input_ids = shift_tokens_right(labels, self.decoder_start_token_id)
-
-        # replace padding with -100 to ignore correctly when computing the loss
-        labels = np.ma.array(labels, mask=np.not_equal(labels_batch.attention_mask, 1))
-        labels = labels.filled(fill_value=-100)
-
-        batch["labels"] = labels
-        batch["decoder_input_ids"] = decoder_input_ids
-
-        return batch
-
-
-def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False, drop_last=True):
-    """
-    Returns batches of size `batch_size` from `dataset`. If `drop_last` is set to `False`, the final batch may be incomplete,
-    and range in size from 1 to `batch_size`. Shuffle batches if `shuffle` is `True`.
-    """
-    if shuffle:
-        batch_idx = jax.random.permutation(rng, len(dataset))
-        batch_idx = np.asarray(batch_idx)
-    else:
-        batch_idx = np.arange(len(dataset))
-
-    if drop_last:
-        steps_per_epoch = len(dataset) // batch_size
-        batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
-        batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
-    else:
-        steps_per_epoch = math.ceil(len(dataset) / batch_size)
-        batch_idx = np.array_split(batch_idx, steps_per_epoch)
-
-    for idx in batch_idx:
-        batch = dataset[idx]
-        yield batch
-
-
-class TrainState(train_state.TrainState):
-    dropout_rng: jnp.ndarray
-
-    def replicate(self):
-        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
-
-
-def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
-    summary_writer.scalar("train_time", train_time, step)
-    train_metrics = get_metrics(train_metrics)
-    for key, vals in train_metrics.items():
-        tag = f"train_{key}"
-        for i, val in enumerate(vals):
-            summary_writer.scalar(tag, val, step - len(vals) + i + 1)
-
-    for metric_name, value in eval_metrics.items():
-        summary_writer.scalar(f"eval_{metric_name}", value, step)
-
-
-def create_learning_rate_fn(
-    num_train_steps: int, num_warmup_steps: int, learning_rate: float
-) -> Callable[[int], jnp.array]:
-    """Returns a linear warmup, linear_decay learning rate function."""
-    warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
-    decay_fn = optax.linear_schedule(
-        init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
-    )
-    schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
-    return schedule_fn
-
-
-def main():
-    # 1. Parse input arguments
-    # See all possible arguments in src/transformers/training_args.py
-    # or by passing the --help flag to this script.
-    # We now keep distinct sets of args, for a cleaner separation of concerns.
-    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, 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()
-
-    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
-    # information sent is the one passed as arguments along with your JAX/Flax versions.
-    send_example_telemetry("run_speech_recognition_seq2seq", model_args, data_args, framework="flax")
-
-    # 2. Setup logging
-    # Make 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",
-        handlers=[logging.StreamHandler(sys.stdout)],
-    )
-    # Set the verbosity to info of the Transformers logger.
-    # We only want one process per machine to log things on the screen.
-    logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
-    if jax.process_index() == 0:
-        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)
-
-    # Check the output dir is valid
-    if (
-        os.path.exists(training_args.output_dir)
-        and os.listdir(training_args.output_dir)
-        and training_args.do_train
-        and not training_args.overwrite_output_dir
-    ):
-        raise ValueError(
-            f"Output directory ({training_args.output_dir}) already exists and is not empty."
-            "Use `--overwrite_output_dir` to overcome."
-        )
-
-    # Handle the repository creation
-    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)
-        repo = Repository(training_args.output_dir, clone_from=repo_name, token=training_args.hub_token)
-
-    # 3. Load dataset
-    raw_datasets = DatasetDict()
-
-    if training_args.do_train:
-        raw_datasets["train"] = load_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
-            split=data_args.train_split_name,
-            cache_dir=data_args.dataset_cache_dir,
-            use_auth_token=True if model_args.use_auth_token else None,
-        )
-
-    if training_args.do_eval:
-        raw_datasets["eval"] = load_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
-            split=data_args.eval_split_name,
-            cache_dir=data_args.dataset_cache_dir,
-            use_auth_token=True if model_args.use_auth_token else None,
-        )
-
-    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."
-        )
-
-    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 '{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 '{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)}."
-        )
-
-    # 5. Load pretrained model, tokenizer, and feature extractor
-    config = AutoConfig.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,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-    feature_extractor = AutoFeatureExtractor.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,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-    tokenizer = AutoTokenizer.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,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-
-    model = FlaxAutoModelForSpeechSeq2Seq.from_pretrained(
-        model_args.model_name_or_path,
-        config=config,
-        dtype=getattr(jnp, model_args.dtype),
-        cache_dir=model_args.cache_dir,
-        revision=model_args.model_revision,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-
-    if model.config.decoder_start_token_id is None:
-        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
-
-    # 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)
-    min_input_length = int(data_args.min_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
-    )
-    pad_input_to_multiple_of = data_args.pad_input_to_multiple_of
-    pad_target_to_multiple_of = data_args.pad_target_to_multiple_of
-    audio_column_name = data_args.audio_column_name
-    num_workers = data_args.preprocessing_num_workers
-    text_column_name = data_args.text_column_name
-    model_input_name = feature_extractor.model_input_names[0]
-    do_lower_case = data_args.do_lower_case
-
-    if training_args.do_train and data_args.max_train_samples is not None:
-        raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
-
-    if training_args.do_eval and data_args.max_eval_samples is not None:
-        raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))
-
-    if data_args.language is not None:
-        # We only need to set the task id when the language is specified (i.e. in a multilingual setting)
-        tokenizer.set_prefix_tokens(language=data_args.language, task=data_args.task)
-
-    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]
-        batch["input_length"] = len(sample["array"])
-
-        # process targets
-        input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
-        batch["labels"] = tokenizer(input_str).input_ids
-        return batch
-
-    vectorized_datasets = raw_datasets.map(
-        prepare_dataset,
-        remove_columns=next(iter(raw_datasets.values())).column_names,
-        num_proc=num_workers,
-        desc="preprocess train dataset",
-    )
-
-    # 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
-
-    vectorized_datasets = vectorized_datasets.filter(
-        is_audio_in_length_range,
-        num_proc=num_workers,
-        input_columns=["input_length"],
-    )
-
-    # 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
-
-    # 8. Load Metric
-    metric = load_metric("wer")
-
-    def compute_metrics(preds, labels):
-        for ind in range(len(labels)):
-            labels[ind][labels[ind] == -100] = tokenizer.pad_token_id
-
-        pred_str = tokenizer.batch_decode(preds, skip_special_tokens=True)
-        # we do not want to group tokens when computing the metrics
-        
-        try:
-            label_str = tokenizer.batch_decode(labels, skip_special_tokens=True)
-        except:
-            breakpoint()
-
-        wer = metric.compute(predictions=pred_str, references=label_str)
-        return {"wer": wer}
-
-    # 9. Save feature extractor, tokenizer and config
-    feature_extractor.save_pretrained(training_args.output_dir)
-    tokenizer.save_pretrained(training_args.output_dir)
-    config.save_pretrained(training_args.output_dir)
-
-    processor = AutoProcessor.from_pretrained(training_args.output_dir)
-
-    data_collator = FlaxDataCollatorSpeechSeq2SeqWithPadding(
-        processor=processor,
-        decoder_start_token_id=model.config.decoder_start_token_id,
-        input_padding="longest",
-        target_padding="longest",
-        max_target_length=max_label_length,
-        pad_input_to_multiple_of=pad_input_to_multiple_of,
-        pad_target_to_multiple_of=pad_target_to_multiple_of if pad_target_to_multiple_of else max_label_length,
-    )
-
-    # Enable tensorboard only on the master node
-    has_tensorboard = is_tensorboard_available()
-    if has_tensorboard and jax.process_index() == 0:
-        try:
-            from flax.metrics.tensorboard import SummaryWriter
-
-            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
-        except ImportError as ie:
-            has_tensorboard = False
-            logger.warning(
-                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
-            )
-    else:
-        logger.warning(
-            "Unable to display metrics through TensorBoard because the package is not installed: "
-            "Please run pip install tensorboard to enable."
-        )
-
-    # Initialize our training
-    rng = jax.random.PRNGKey(training_args.seed)
-    rng, dropout_rng = jax.random.split(rng)
-
-    # Store some constant
-    num_epochs = int(training_args.num_train_epochs)
-    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
-    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
-    eval_batch_size = per_device_eval_batch_size * jax.device_count()
-    steps_per_epoch = len(vectorized_datasets["train"]) // train_batch_size
-    total_train_steps = steps_per_epoch * num_epochs
-
-    # Create learning rate schedule
-    linear_decay_lr_schedule_fn = create_learning_rate_fn(
-        len(vectorized_datasets["train"]),
-        training_args.warmup_steps,
-        training_args.learning_rate,
-    )
-
-    # We use Optax's "masking" functionality to not apply weight decay
-    # to bias and LayerNorm scale parameters. decay_mask_fn returns a
-    # mask boolean with the same structure as the parameters.
-    # The mask is True for parameters that should be decayed.
-    def decay_mask_fn(params):
-        flat_params = traverse_util.flatten_dict(params)
-        # find out all LayerNorm parameters
-        layer_norm_candidates = ["layernorm", "layer_norm", "ln"]
-        layer_norm_named_params = set(
-            [
-                layer[-2:]
-                for layer_norm_name in layer_norm_candidates
-                for layer in flat_params.keys()
-                if layer_norm_name in "".join(layer).lower()
-            ]
-        )
-        flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_named_params) for path in flat_params}
-        return traverse_util.unflatten_dict(flat_mask)
-
-    # create adam optimizer
-    adamw = optax.adamw(
-        learning_rate=linear_decay_lr_schedule_fn,
-        b1=training_args.adam_beta1,
-        b2=training_args.adam_beta2,
-        eps=training_args.adam_epsilon,
-        weight_decay=training_args.weight_decay,
-        mask=decay_mask_fn,
-    )
-
-    # Setup train state
-    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
-
-    # label smoothed cross entropy
-    def loss_fn(logits, labels, label_smoothing_factor=0.0):
-        """
-        The label smoothing implementation is adapted from Flax's official example:
-        https://github.com/google/flax/blob/87a211135c6a377c8f29048a1cac3840e38b9da4/examples/wmt/train.py#L104
-        """
-        vocab_size = logits.shape[-1]
-        confidence = 1.0 - label_smoothing_factor
-        low_confidence = (1.0 - confidence) / (vocab_size - 1)
-        normalizing_constant = -(
-            confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
-        )
-        soft_labels = onehot(labels, vocab_size, on_value=confidence, off_value=low_confidence)
-
-        loss = optax.softmax_cross_entropy(logits, soft_labels)
-        loss = loss - normalizing_constant
-
-        # ignore padded tokens from loss, i.e. where labels are not set to -100
-        padding_mask = labels >= 0
-        loss = loss * padding_mask
-        loss = loss.sum()
-        num_labels = padding_mask.sum()
-        return loss, num_labels
-
-    # Define gradient update step fn
-    def train_step(state, batch, label_smoothing_factor=0.0):
-        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
-
-        def compute_loss(params):
-            labels = batch.pop("labels")
-            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
-            loss, num_labels = loss_fn(logits, labels, label_smoothing_factor)
-            return loss, num_labels
-
-        grad_fn = jax.value_and_grad(compute_loss, has_aux=True)
-        (loss, num_labels), grad = grad_fn(state.params)
-        num_labels = jax.lax.psum(num_labels, "batch")
-
-        # true loss = total loss / total samples
-        loss = jax.lax.psum(loss, "batch")
-        loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)
-
-        # true grad = total grad / total samples
-        grad = jax.lax.psum(grad, "batch")
-        grad = jax.tree_util.tree_map(lambda x: x / num_labels, grad)
-        new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
-
-        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
-        return new_state, metrics
-
-    # Define eval fn
-    def eval_step(params, batch, label_smoothing_factor=0.0):
-        labels = batch.pop("labels")
-        logits = model(**batch, params=params, train=False)[0]
-
-        loss, num_labels = loss_fn(logits, labels, label_smoothing_factor)
-        num_labels = jax.lax.psum(num_labels, "batch")
-
-        # true loss = total loss / total samples
-        loss = jax.lax.psum(loss, "batch")
-        loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)
-
-        metrics = {"loss": loss}
-        return metrics
-
-    # Define generation function
-    num_beams = model_args.num_beams if model_args.num_beams is not None else model.config.num_beams
-    gen_kwargs = {"max_length": max_label_length, "num_beams": num_beams}
-
-    def generate_step(params, batch):
-        model.params = params
-        output_ids = model.generate(batch[model_input_name], attention_mask=batch.get("attention_mask"), **gen_kwargs)
-        return output_ids.sequences
-
-    # Create parallel version of the train and eval step
-    p_train_step = jax.pmap(
-        partial(train_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch", donate_argnums=(0,)
-    )
-    p_eval_step = jax.pmap(partial(eval_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch")
-    p_generate_step = jax.pmap(generate_step, "batch")
-
-    # Replicate the train state on each device
-    state = state.replicate()
-
-    logger.info("***** Running training *****")
-    logger.info(f"  Num examples = {len(vectorized_datasets['train'])}")
-    logger.info(f"  Num Epochs = {num_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
-    logger.info(f"  Total optimization steps = {total_train_steps}")
-
-    train_time = 0
-    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
-    for epoch in epochs:
-        # ======================== Training ================================
-        train_start = time.time()
-
-        # Create sampling rng
-        rng, input_rng = jax.random.split(rng)
-        train_metrics = []
-
-        # Generate an epoch by shuffling sampling indices from the train dataset
-        train_loader = data_loader(input_rng, vectorized_datasets["train"], train_batch_size, shuffle=True)
-        # train
-        for _ in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
-            samples = next(train_loader)
-            batch = data_collator(samples)
-            batch = shard(batch.data)
-            state, train_metric = p_train_step(state, batch)
-            train_metrics.append(train_metric)
-
-        train_time += time.time() - train_start
-
-        train_metrics = unreplicate(train_metrics)
-        
-        if train_metrics:
-            epochs.write(
-                f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metrics['loss']}, Learning Rate:"
-                f" {train_metrics['learning_rate']})"
-            )
-
-        # ======================== Evaluating ==============================
-        eval_metrics = []
-        eval_preds = []
-        eval_labels = []
-
-        eval_loader = data_loader(input_rng, vectorized_datasets["eval"], eval_batch_size, drop_last=False)
-        eval_steps = math.ceil(len(vectorized_datasets["eval"]) / eval_batch_size)
-        for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
-            # Model forward
-            samples = next(eval_loader)
-            batch = data_collator(samples)
-            labels = batch["labels"]
-
-            metrics = pad_shard_unpad(p_eval_step, static_return=True)(
-                state.params, batch.data, min_device_batch=per_device_eval_batch_size
-            )
-            eval_metrics.append(metrics)
-
-            # generation
-            if training_args.predict_with_generate:
-                generated_ids = pad_shard_unpad(p_generate_step)(state.params, batch.data)
-                eval_preds.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
-                eval_labels.extend(labels)
-
-        # normalize eval metrics
-        eval_metrics = get_metrics(eval_metrics)
-        eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
-
-        # compute WER metric
-        wer_desc = ""
-        if training_args.predict_with_generate:
-            wer_metric = 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()])
-
-        # Print metrics and update progress bar
-        desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {wer_desc})"
-        epochs.write(desc)
-        epochs.desc = desc
-
-        # Save metrics
-        if has_tensorboard and jax.process_index() == 0:
-            cur_step = epoch * (len(vectorized_datasets["train"]) // train_batch_size)
-            if train_metrics:
-                write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)
-            else:
-                print(f"Train metrics not written because currently it is empty.")
-
-        # save checkpoint after each epoch and push checkpoint to the hub
-        if jax.process_index() == 0:
-            params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
-            model.save_pretrained(training_args.output_dir, params=params)
-            tokenizer.save_pretrained(training_args.output_dir)
-            if training_args.push_to_hub:
-                repo.push_to_hub(commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False)
-
-
-if __name__ == "__main__":
-    main()

backup_run_flax_speech_recognition_seq2seq_streaming.py

@@ -1,874 +0,0 @@
-#!/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.
-"""
-Fine-tuning the Flax library models for sequence to sequence speech recognition.
-"""
-# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
-
-import logging
-import math
-import os
-import sys
-import time
-from dataclasses import field
-from functools import partial
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import datasets
-import flax
-import jax
-import jax.numpy as jnp
-import numpy as np
-import optax
-from datasets import Dataset, DatasetDict, load_dataset, load_metric
-from flax import jax_utils, traverse_util
-from flax.jax_utils import pad_shard_unpad, unreplicate
-from flax.training import train_state
-from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
-from huggingface_hub import Repository, create_repo
-from tqdm import tqdm
-
-import transformers
-from transformers import (
-    AutoConfig,
-    AutoFeatureExtractor,
-    AutoProcessor,
-    AutoTokenizer,
-    FlaxAutoModelForSpeechSeq2Seq,
-    HfArgumentParser,
-    Seq2SeqTrainingArguments,
-    is_tensorboard_available,
-)
-from transformers.file_utils import get_full_repo_name
-from transformers.utils import check_min_version, send_example_telemetry
-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.27.0.dev0")
-
-require_version("datasets>=1.18.0", "To fix: pip install -r examples/flax/speech-recogintion/requirements.txt")
-
-logger = logging.getLogger(__name__)
-
-
-@flax.struct.dataclass
-class ModelArguments:
-    """
-    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
-    """
-
-    model_name_or_path: str = field(
-        metadata={"help": "Path to pretrained 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)."},
-    )
-    use_auth_token: bool = field(
-        default=False,
-        metadata={
-            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
-            "with private models)."
-        },
-    )
-    dtype: Optional[str] = field(
-        default="float32",
-        metadata={
-            "help": (
-                "Floating-point format in which the model weights should be initialized and trained. Choose one of"
-                " `[float32, float16, bfloat16]`."
-            )
-        },
-    )
-    num_beams: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": (
-                "Number of beams to use for evaluation. This argument will be passed to `model.generate`, "
-                "which is used during evaluation."
-            )
-        },
-    )
-
-
-@flax.struct.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)."}
-    )
-    text_column: Optional[str] = field(
-        default=None,
-        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
-    )
-    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."},
-    )
-    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="text",
-        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},
-    )
-    max_duration_in_seconds: float = field(
-        default=20.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: float = field(
-        default=128,
-        metadata={"help": "Truncate transcriptions that are longer `max_eval_length` tokens."},
-    )
-    pad_input_to_multiple_of: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "If set will pad the input sequence to a multiple of the provided value. "
-            "This is important to avoid triggering recompilations on TPU. If unspecified, will default to padding the inputs to max length."
-        },
-    )
-    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'"
-        },
-    )
-    do_lower_case: bool = field(
-        default=True,
-        metadata={"help": "Whether the target text should be lower cased."},
-    )
-    language: str = field(
-        default=None,
-        metadata={
-            "help": (
-                "Language for multilingual fine-tuning. This argument should be set for multilingual fine-tuning "
-                "only. For English speech recognition, it should be set to `None`."
-            )
-        },
-    )
-    task: str = field(
-        default="transcribe",
-        metadata={"help": "Task, either `transcribe` for speech recognition or `translate` for speech translation."},
-    )
-
-
-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
-
-
-@flax.struct.dataclass
-class FlaxDataCollatorSpeechSeq2SeqWithPadding:
-    """
-    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 begin-of-sentence 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_input_length (:obj:`float`, `optional`):
-            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).
-        max_target_length (:obj:`int`, `optional`):
-            Maximum length of the ``labels`` of the returned list and optionally padding length (see above).
-        pad_input_to_multiple_of (:obj:`int`, `optional`):
-            If set will pad the input sequence to a multiple of the provided value.
-            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
-            7.5 (Volta).
-        pad_target_to_multiple_of (:obj:`int`, `optional`):
-            If set will pad the target sequence to a multiple of the provided value.
-            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
-            7.5 (Volta).
-    """
-
-    processor: Any
-    decoder_start_token_id: int
-    input_padding: Union[bool, str] = "longest"
-    target_padding: Union[bool, str] = "max_length"
-    max_input_length: Optional[float] = None
-    max_target_length: Optional[int] = None
-    pad_input_to_multiple_of: Optional[int] = None
-    pad_target_to_multiple_of: 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]
-        input_features = {model_input_name: features[model_input_name]}
-        label_features = {"input_ids": features["labels"]}
-
-        # reformat list to dict and set to pytorch format
-        batch = self.processor.feature_extractor.pad(
-            input_features,
-            max_length=self.max_input_length,
-            padding=self.input_padding,
-            pad_to_multiple_of=self.pad_input_to_multiple_of,
-            return_tensors="np",
-        )
-
-        labels_batch = self.processor.tokenizer.pad(
-            label_features,
-            max_length=self.max_target_length,
-            padding=self.target_padding,
-            pad_to_multiple_of=self.pad_target_to_multiple_of,
-            return_tensors="np",
-        )
-
-        # if bos token is appended in previous tokenization step,
-        # cut bos token here as it's append later anyways
-        labels = labels_batch["input_ids"]
-        if (labels[:, 0] == self.decoder_start_token_id).all().item():
-            labels = labels[:, 1:]
-            labels_batch.attention_mask = labels_batch.attention_mask[:, 1:]
-
-        decoder_input_ids = shift_tokens_right(labels, self.decoder_start_token_id)
-
-        # replace padding with -100 to ignore correctly when computing the loss
-        labels = np.ma.array(labels, mask=np.not_equal(labels_batch.attention_mask, 1))
-        labels = labels.filled(fill_value=-100)
-
-        batch["labels"] = labels
-        batch["decoder_input_ids"] = decoder_input_ids
-
-        return batch
-
-
-def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False, drop_last=True):
-    """
-    Returns batches of size `batch_size` from `dataset`. If `drop_last` is set to `False`, the final batch may be incomplete,
-    and range in size from 1 to `batch_size`. Shuffle batches if `shuffle` is `True`.
-    """
-    if shuffle:
-        batch_idx = jax.random.permutation(rng, len(dataset))
-        batch_idx = np.asarray(batch_idx)
-    else:
-        batch_idx = np.arange(len(dataset))
-
-    if drop_last:
-        steps_per_epoch = len(dataset) // batch_size
-        batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
-        batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
-    else:
-        steps_per_epoch = math.ceil(len(dataset) / batch_size)
-        batch_idx = np.array_split(batch_idx, steps_per_epoch)
-
-    for idx in batch_idx:
-        batch = dataset[idx]
-        yield batch
-
-
-class TrainState(train_state.TrainState):
-    dropout_rng: jnp.ndarray
-
-    def replicate(self):
-        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
-
-
-def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
-    summary_writer.scalar("train_time", train_time, step)
-
-    train_metrics = get_metrics(train_metrics)
-    for key, vals in train_metrics.items():
-        tag = f"train_{key}"
-        for i, val in enumerate(vals):
-            summary_writer.scalar(tag, val, step - len(vals) + i + 1)
-
-    for metric_name, value in eval_metrics.items():
-        summary_writer.scalar(f"eval_{metric_name}", value, step)
-
-
-def create_learning_rate_fn(
-    num_train_steps: int, num_warmup_steps: int, learning_rate: float
-) -> Callable[[int], jnp.array]:
-    """Returns a linear warmup, linear_decay learning rate function."""
-    warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
-    decay_fn = optax.linear_schedule(
-        init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
-    )
-    schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
-    return schedule_fn
-
-
-def main():
-    # 1. Parse input arguments
-    # See all possible arguments in src/transformers/training_args.py
-    # or by passing the --help flag to this script.
-    # We now keep distinct sets of args, for a cleaner separation of concerns.
-    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, 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()
-
-    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
-    # information sent is the one passed as arguments along with your JAX/Flax versions.
-    send_example_telemetry("run_speech_recognition_seq2seq", model_args, data_args, framework="flax")
-
-    # 2. Setup logging
-    # Make 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",
-        handlers=[logging.StreamHandler(sys.stdout)],
-    )
-    # Set the verbosity to info of the Transformers logger.
-    # We only want one process per machine to log things on the screen.
-    logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
-    if jax.process_index() == 0:
-        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)
-
-    # Check the output dir is valid
-    if (
-        os.path.exists(training_args.output_dir)
-        and os.listdir(training_args.output_dir)
-        and training_args.do_train
-        and not training_args.overwrite_output_dir
-    ):
-        raise ValueError(
-            f"Output directory ({training_args.output_dir}) already exists and is not empty."
-            "Use `--overwrite_output_dir` to overcome."
-        )
-
-    # Handle the repository creation
-    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)
-        repo = Repository(training_args.output_dir, clone_from=repo_name, token=training_args.hub_token)
-
-    # 3. Load dataset
-    raw_datasets = DatasetDict()
-
-    if training_args.do_train:
-        raw_datasets["train"] = load_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
-            split=data_args.train_split_name,
-            cache_dir=data_args.dataset_cache_dir,
-            use_auth_token=True if model_args.use_auth_token else None,
-        )
-
-    if training_args.do_eval:
-        raw_datasets["eval"] = load_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
-            split=data_args.eval_split_name,
-            cache_dir=data_args.dataset_cache_dir,
-            use_auth_token=True if model_args.use_auth_token else None,
-        )
-
-    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."
-        )
-
-    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 '{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 '{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)}."
-        )
-
-    # 5. Load pretrained model, tokenizer, and feature extractor
-    config = AutoConfig.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,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-    feature_extractor = AutoFeatureExtractor.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,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-    tokenizer = AutoTokenizer.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,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-
-    model = FlaxAutoModelForSpeechSeq2Seq.from_pretrained(
-        model_args.model_name_or_path,
-        config=config,
-        dtype=getattr(jnp, model_args.dtype),
-        cache_dir=model_args.cache_dir,
-        revision=model_args.model_revision,
-        use_auth_token=True if model_args.use_auth_token else None,
-    )
-
-    if model.config.decoder_start_token_id is None:
-        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
-
-    # 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)
-    min_input_length = int(data_args.min_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
-    )
-    pad_input_to_multiple_of = data_args.pad_input_to_multiple_of
-    pad_target_to_multiple_of = data_args.pad_target_to_multiple_of
-    audio_column_name = data_args.audio_column_name
-    num_workers = data_args.preprocessing_num_workers
-    text_column_name = data_args.text_column_name
-    model_input_name = feature_extractor.model_input_names[0]
-    do_lower_case = data_args.do_lower_case
-
-    if training_args.do_train and data_args.max_train_samples is not None:
-        raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
-
-    if training_args.do_eval and data_args.max_eval_samples is not None:
-        raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))
-
-    if data_args.language is not None:
-        # We only need to set the task id when the language is specified (i.e. in a multilingual setting)
-        tokenizer.set_prefix_tokens(language=data_args.language, task=data_args.task)
-
-    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]
-        batch["input_length"] = len(sample["array"])
-
-        # process targets
-        input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
-        batch["labels"] = tokenizer(input_str).input_ids
-        return batch
-
-    vectorized_datasets = raw_datasets.map(
-        prepare_dataset,
-        remove_columns=next(iter(raw_datasets.values())).column_names,
-        num_proc=num_workers,
-        desc="preprocess train dataset",
-    )
-
-    # 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
-
-    vectorized_datasets = vectorized_datasets.filter(
-        is_audio_in_length_range,
-        num_proc=num_workers,
-        input_columns=["input_length"],
-    )
-
-    # 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
-
-    # 8. Load Metric
-    metric = load_metric("wer")
-
-    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)
-        # we do not want to group tokens when computing the metrics
-        label_str = tokenizer.batch_decode(labels, skip_special_tokens=True)
-
-        wer = metric.compute(predictions=pred_str, references=label_str)
-        return {"wer": wer}
-
-    # 9. Save feature extractor, tokenizer and config
-    feature_extractor.save_pretrained(training_args.output_dir)
-    tokenizer.save_pretrained(training_args.output_dir)
-    config.save_pretrained(training_args.output_dir)
-
-    processor = AutoProcessor.from_pretrained(training_args.output_dir)
-
-    data_collator = FlaxDataCollatorSpeechSeq2SeqWithPadding(
-        processor=processor,
-        decoder_start_token_id=model.config.decoder_start_token_id,
-        input_padding="longest",
-        target_padding="longest",
-        max_target_length=max_label_length,
-        pad_input_to_multiple_of=pad_input_to_multiple_of,
-        pad_target_to_multiple_of=pad_target_to_multiple_of if pad_target_to_multiple_of else max_label_length,
-    )
-
-    # Enable tensorboard only on the master node
-    has_tensorboard = is_tensorboard_available()
-    if has_tensorboard and jax.process_index() == 0:
-        try:
-            from flax.metrics.tensorboard import SummaryWriter
-
-            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
-        except ImportError as ie:
-            has_tensorboard = False
-            logger.warning(
-                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
-            )
-    else:
-        logger.warning(
-            "Unable to display metrics through TensorBoard because the package is not installed: "
-            "Please run pip install tensorboard to enable."
-        )
-
-    # Initialize our training
-    rng = jax.random.PRNGKey(training_args.seed)
-    rng, dropout_rng = jax.random.split(rng)
-
-    # Store some constant
-    num_epochs = int(training_args.num_train_epochs)
-    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
-    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
-    eval_batch_size = per_device_eval_batch_size * jax.device_count()
-    steps_per_epoch = len(vectorized_datasets["train"]) // train_batch_size
-    total_train_steps = steps_per_epoch * num_epochs
-
-    # Create learning rate schedule
-    linear_decay_lr_schedule_fn = create_learning_rate_fn(
-        len(vectorized_datasets["train"]),
-        training_args.warmup_steps,
-        training_args.learning_rate,
-    )
-
-    # We use Optax's "masking" functionality to not apply weight decay
-    # to bias and LayerNorm scale parameters. decay_mask_fn returns a
-    # mask boolean with the same structure as the parameters.
-    # The mask is True for parameters that should be decayed.
-    def decay_mask_fn(params):
-        flat_params = traverse_util.flatten_dict(params)
-        # find out all LayerNorm parameters
-        layer_norm_candidates = ["layernorm", "layer_norm", "ln"]
-        layer_norm_named_params = set(
-            [
-                layer[-2:]
-                for layer_norm_name in layer_norm_candidates
-                for layer in flat_params.keys()
-                if layer_norm_name in "".join(layer).lower()
-            ]
-        )
-        flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_named_params) for path in flat_params}
-        return traverse_util.unflatten_dict(flat_mask)
-
-    # create adam optimizer
-    adamw = optax.adamw(
-        learning_rate=linear_decay_lr_schedule_fn,
-        b1=training_args.adam_beta1,
-        b2=training_args.adam_beta2,
-        eps=training_args.adam_epsilon,
-        weight_decay=training_args.weight_decay,
-        mask=decay_mask_fn,
-    )
-
-    # Setup train state
-    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
-
-    # label smoothed cross entropy
-    def loss_fn(logits, labels, label_smoothing_factor=0.0):
-        """
-        The label smoothing implementation is adapted from Flax's official example:
-        https://github.com/google/flax/blob/87a211135c6a377c8f29048a1cac3840e38b9da4/examples/wmt/train.py#L104
-        """
-        vocab_size = logits.shape[-1]
-        confidence = 1.0 - label_smoothing_factor
-        low_confidence = (1.0 - confidence) / (vocab_size - 1)
-        normalizing_constant = -(
-            confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
-        )
-        soft_labels = onehot(labels, vocab_size, on_value=confidence, off_value=low_confidence)
-
-        loss = optax.softmax_cross_entropy(logits, soft_labels)
-        loss = loss - normalizing_constant
-
-        # ignore padded tokens from loss, i.e. where labels are not set to -100
-        padding_mask = labels >= 0
-        loss = loss * padding_mask
-        loss = loss.sum()
-        num_labels = padding_mask.sum()
-        return loss, num_labels
-
-    # Define gradient update step fn
-    def train_step(state, batch, label_smoothing_factor=0.0):
-        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
-
-        def compute_loss(params):
-            labels = batch.pop("labels")
-            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
-            loss, num_labels = loss_fn(logits, labels, label_smoothing_factor)
-            return loss, num_labels
-
-        grad_fn = jax.value_and_grad(compute_loss, has_aux=True)
-        (loss, num_labels), grad = grad_fn(state.params)
-        num_labels = jax.lax.psum(num_labels, "batch")
-
-        # true loss = total loss / total samples
-        loss = jax.lax.psum(loss, "batch")
-        loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)
-
-        # true grad = total grad / total samples
-        grad = jax.lax.psum(grad, "batch")
-        grad = jax.tree_util.tree_map(lambda x: x / num_labels, grad)
-        new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
-
-        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
-        return new_state, metrics
-
-    # Define eval fn
-    def eval_step(params, batch, label_smoothing_factor=0.0):
-        labels = batch.pop("labels")
-        logits = model(**batch, params=params, train=False)[0]
-
-        loss, num_labels = loss_fn(logits, labels, label_smoothing_factor)
-        num_labels = jax.lax.psum(num_labels, "batch")
-
-        # true loss = total loss / total samples
-        loss = jax.lax.psum(loss, "batch")
-        loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)
-
-        metrics = {"loss": loss}
-        return metrics
-
-    # Define generation function
-    num_beams = model_args.num_beams if model_args.num_beams is not None else model.config.num_beams
-    gen_kwargs = {"max_length": max_label_length, "num_beams": num_beams}
-
-    def generate_step(params, batch):
-        model.params = params
-        output_ids = model.generate(batch[model_input_name], attention_mask=batch.get("attention_mask"), **gen_kwargs)
-        return output_ids.sequences
-
-    # Create parallel version of the train and eval step
-    p_train_step = jax.pmap(
-        partial(train_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch", donate_argnums=(0,)
-    )
-    p_eval_step = jax.pmap(partial(eval_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch")
-    p_generate_step = jax.pmap(generate_step, "batch")
-
-    # Replicate the train state on each device
-    state = state.replicate()
-
-    logger.info("***** Running training *****")
-    logger.info(f"  Num examples = {len(vectorized_datasets['train'])}")
-    logger.info(f"  Num Epochs = {num_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
-    logger.info(f"  Total optimization steps = {total_train_steps}")
-
-    train_time = 0
-    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
-    for epoch in epochs:
-        # ======================== Training ================================
-        train_start = time.time()
-
-        # Create sampling rng
-        rng, input_rng = jax.random.split(rng)
-        train_metrics = []
-
-        # Generate an epoch by shuffling sampling indices from the train dataset
-        train_loader = data_loader(input_rng, vectorized_datasets["train"], train_batch_size, shuffle=True)
-        # train
-        for _ in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
-            samples = next(train_loader)
-            batch = data_collator(samples)
-            batch = shard(batch.data)
-            state, train_metric = p_train_step(state, batch)
-            train_metrics.append(train_metric)
-
-        train_time += time.time() - train_start
-
-        train_metric = unreplicate(train_metric)
-
-        epochs.write(
-            f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate:"
-            f" {train_metric['learning_rate']})"
-        )
-
-        # ======================== Evaluating ==============================
-        eval_metrics = []
-        eval_preds = []
-        eval_labels = []
-
-        eval_loader = data_loader(input_rng, vectorized_datasets["eval"], eval_batch_size, drop_last=False)
-        eval_steps = math.ceil(len(vectorized_datasets["eval"]) / eval_batch_size)
-        for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
-            # Model forward
-            samples = next(eval_loader)
-            batch = data_collator(samples)
-            labels = batch["labels"]
-
-            metrics = pad_shard_unpad(p_eval_step, static_return=True)(
-                state.params, batch.data, min_device_batch=per_device_eval_batch_size
-            )
-            eval_metrics.append(metrics)
-
-            # generation
-            if training_args.predict_with_generate:
-                generated_ids = pad_shard_unpad(p_generate_step)(state.params, batch.data)
-                eval_preds.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
-                eval_labels.extend(labels)
-
-        # normalize eval metrics
-        eval_metrics = get_metrics(eval_metrics)
-        eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
-
-        # compute WER metric
-        wer_desc = ""
-        if training_args.predict_with_generate:
-            wer_metric = 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()])
-
-        # Print metrics and update progress bar
-        desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {wer_desc})"
-        epochs.write(desc)
-        epochs.desc = desc
-
-        # Save metrics
-        if has_tensorboard and jax.process_index() == 0:
-            cur_step = epoch * (len(vectorized_datasets["train"]) // train_batch_size)
-            write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)
-
-        # save checkpoint after each epoch and push checkpoint to the hub
-        if jax.process_index() == 0:
-            params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
-            model.save_pretrained(training_args.output_dir, params=params)
-            tokenizer.save_pretrained(training_args.output_dir)
-            if training_args.push_to_hub:
-                repo.push_to_hub(commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False)
-
-
-if __name__ == "__main__":
-    main()

run.sh

@@ -1,9 +1,10 @@
 python run_flax_speech_recognition_seq2seq.py \
 	    --model_name_or_path openai/whisper-small.en \
-            --dataset_name hf-internal-testing/librispeech_asr_dummy \
-	    --dataset_config clean \
-            --train_split_name validation \
-            --eval_split_name validation \
+            --dataset_name mozilla-foundation/common_voice_11_0 \
+            --dataset_config es \
+	    --text_column_name sentence \
+            --train_split_name test\
+            --eval_split_name test\
             --output_dir whisper-small-flaxtest \
             --overwrite_output_dir \
             --num_train_epochs=2 \

run_flax_speech_recognition_seq2seq.py

@@ -1,5 +1,3 @@
-#!/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");

run_flax_speech_recognition_seq2seq_streaming.py

@@ -177,7 +177,7 @@ class DataTrainingArguments:
         metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},
     )
     max_duration_in_seconds: float = field(
-        default=20.0,
+        default=30.0,
         metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"},
     )
     min_duration_in_seconds: float = field(
@@ -501,7 +501,7 @@ def main():
     raw_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict()
     
     if training_args.do_train:
-        raw_datasets["train"] = load_maybe_streaming_dataset(
+        raw_datasets[data_args.train_split_name] = load_maybe_streaming_dataset(
             data_args.dataset_name,
             data_args.dataset_config_name,
             split=data_args.train_split_name,
@@ -511,7 +511,7 @@ def main():
         )
 
     if training_args.do_eval:
-        raw_datasets["eval"] = load_maybe_streaming_dataset(
+        raw_datasets[data_args.eval_split_name] = load_maybe_streaming_dataset(
             data_args.dataset_name,
             data_args.dataset_config_name,
             split=data_args.eval_split_name,
@@ -625,69 +625,50 @@ def main():
             remove_columns=raw_datasets_features,
         ).with_format("torch")
 
-        # Moving this to later. Better ways of doing this
-        #if training_args.do_train and data_args.streaming:
-            # manually shuffle if streaming (done by the trainer for non-streaming)
-            #vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(
-            #    buffer_size=data_args.shuffle_buffer_size,
-            #    seed=training_args.seed,
-            #)
-    
-    
+ 
     # 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
 
-    # For debugging
-    #def is_audio_in_length_range(length):
-    #    if min_input_length < length < max_input_length:
-    #        return True
-    #    else:
-    #        print(f"Warning: Input length {length} is not within the expected range [{min_input_length}, {max_input_length}].")
-    #        return False
-        
-    
+
     if training_args.do_train:
-        vectorized_datasets["train"] = vectorized_datasets["train"].filter(
+        vectorized_datasets[data_args.train_split_name] = vectorized_datasets[data_args.train_split_name].filter(
             is_audio_in_length_range,
             input_columns=["input_length"],
         )
-
-    # 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
-    # Not really needed for streaming
-    # 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 training_args.do_eval:
+        vectorized_datasets[data_args.eval_split_name] = vectorized_datasets[data_args.eval_split_name].filter(
+            is_audio_in_length_range,
+            input_columns=["input_length"],
+        )
+    
     # 8. Load Metric
     metric = evaluate.load("wer")
     do_normalize_eval = data_args.do_normalize_eval
 
-    def compute_metrics(pred):
-        pred_ids = pred.predictions
-
-        pred.label_ids[pred.label_ids == -100] = tokenizer.pad_token_id
+    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(pred_ids, skip_special_tokens=True)
+        pred_str = tokenizer.batch_decode(preds, skip_special_tokens=True)
         # we do not want to group tokens when computing the metrics
-        label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True)
-
-        if do_normalize_eval:
-            pred_str = [normalizer(pred) for pred in pred_str]
-            label_str = [normalizer(label) for label in label_str]
-            # filtering step to only evaluate the samples that correspond to non-zero references:
-            pred_str = [pred_str[i] for i in range(len(pred_str)) if len(label_str[i]) > 0]
-            label_str = [label_str[i] for i in range(len(label_str)) if len(label_str[i]) > 0]
-
-        wer = 100 * metric.compute(predictions=pred_str, references=label_str)
-
+        label_str = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        # TODO
+        # We should implement and test this as well
+        #if do_normalize_eval:
+        #    pred_str = [normalizer(pred) for pred in pred_str]
+        #    label_str = [normalizer(label) for label in label_str]
+        #    # filtering step to only evaluate the samples that correspond to non-zero references:
+        #    pred_str = [pred_str[i] for i in range(len(pred_str)) if len(label_str[i]) > 0]
+        #    label_str = [label_str[i] for i in range(len(label_str)) if len(label_str[i]) > 0]
+        
+        
+        wer = metric.compute(predictions=pred_str, references=label_str)
         return {"wer": wer}
-
+    
+    
     # 9. Save feature extractor, tokenizer and config
     feature_extractor.save_pretrained(training_args.output_dir)
     tokenizer.save_pretrained(training_args.output_dir)
@@ -730,8 +711,7 @@ def main():
     # Store some constant
     #num_epochs = int(training_args.num_train_epochs)
     train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
-    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
-    eval_batch_size = per_device_eval_batch_size * jax.device_count()
+    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
     
 
     # Create learning rate schedule
@@ -875,16 +855,16 @@ def main():
     # Create sampling rng
     #rng, input_rng = jax.random.split(rng)
     train_metrics = []
-    epoch = 0
     
+    # TODO
+    # Do the reset epoch stuff to shuffle
+    
+    epoch = 0
     
     # Create a batched data iterator
     num_workers = 0
-    # This is not working
-    # vectorized_datasets["train"] = vectorized_datasets["train"].shuffle()
-    batched_data_loader = torch.utils.data.DataLoader( batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=lambda x: x )
-    batched_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(batched_data_loader)
-    
+    batched_data_loader = torch.utils.data.DataLoader( batch_size=train_batch_size, dataset=vectorized_datasets[data_args.train_split_name],  num_workers=num_workers, collate_fn=lambda x: x )
+    batched_data_iterator = iter(batched_data_loader)
     
     # train
     for step in tqdm(range(data_args.num_train_steps), desc="Training...", position=1, leave=False):        
@@ -894,8 +874,9 @@ def main():
             
         except StopIteration:
             epoch += 1
-            batched_data_loader = torch.utils.data.DataLoader( batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=lambda x: x )
-            batched_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(batched_data_loader)            
+            # TODO - Not currently shuffled
+            batched_data_loader = torch.utils.data.DataLoader( batch_size=train_batch_size, dataset=vectorized_datasets[data_args.train_split_name],  num_workers=num_workers, collate_fn=lambda x: x )
+            batched_data_iterator = iter(batched_data_loader)
             samples = next(batched_data_iterator)
             
             logger.info(
@@ -908,18 +889,23 @@ def main():
         batch = shard(batch.data)
         state, train_metric = p_train_step(state, batch)
         
+        train_time += time.time() - train_start
+        train_metric = unreplicate(train_metric)
+        
         # ======================== Evaluating ==============================
         if step % training_args.eval_steps == 0 and step > 0:
-            eval_loader = data_loader(input_rng, vectorized_datasets["eval"], eval_batch_size, drop_last=False)
-
-            for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
+            batched_data_eval_loader = torch.utils.data.DataLoader( batch_size=eval_batch_size, dataset=vectorized_datasets[data_args.eval_split_name],  num_workers=num_workers, collate_fn=lambda x: x )
+            batched_data_eval_iterator = iter(batched_data_eval_loader)
+            
+            for _ in tqdm(range(data_args.max_eval_samples//eval_batch_size), desc="Evaluating...", position=2, leave=False):
                 # Model forward
-                samples = next(eval_loader)
-                batch = data_collator(samples)
+                samples = next(batched_data_eval_iterator)
+                reshaped_samples = {key: [feature[key] for feature in samples] for key in samples[0].keys()}
+                batch = data_collator(reshaped_samples)
                 labels = batch["labels"]
 
                 metrics = pad_shard_unpad(p_eval_step, static_return=True)(
-                    state.params, batch.data, min_device_batch=per_device_eval_batch_size
+                    state.params, batch.data, min_device_batch=training_args.per_device_eval_batch_size
                 )
                 eval_metrics.append(metrics)
 
@@ -944,22 +930,22 @@ def main():
             desc = f"Epoch... ({epoch} | Eval Loss: {eval_metrics['loss']} | {wer_desc})"
             logger.info(desc)
 
-    train_time += time.time() - train_start
-    train_metric = unreplicate(train_metric)
-
-
-    # Save metrics
-    if has_tensorboard and jax.process_index() == 0:
-        cur_step = epoch * (len(vectorized_datasets["train"]) // train_batch_size)
-        write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)
-
-    # save checkpoint after each epoch and push checkpoint to the hub
-    if jax.process_index() == 0:
-        params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
-        model.save_pretrained(training_args.output_dir, params=params)
-        tokenizer.save_pretrained(training_args.output_dir)
-        if training_args.push_to_hub:
-            repo.push_to_hub(commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False)
+            # Save metrics
+            if has_tensorboard and jax.process_index() == 0:
+                #TODO
+                breakpoint()
+                # cur_step = epoch * (len(vectorized_datasets[data_args.train_split_name]) // train_batch_size)
+                write_metric(summary_writer, train_metrics, eval_metrics, train_time, data_args.num_train_steps)
+
+            # TODO THis is not happening at every epoch!!!
+            breakpoint()
+            # save checkpoint after each epoch and push checkpoint to the hub
+            if jax.process_index() == 0:
+                params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
+                model.save_pretrained(training_args.output_dir, params=params)
+                tokenizer.save_pretrained(training_args.output_dir)
+                if training_args.push_to_hub:
+                    repo.push_to_hub(commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False)
 
 
 if __name__ == "__main__":

run_test.sh → run_streaming.sh

@@ -5,17 +5,18 @@ python run_flax_speech_recognition_seq2seq_streaming.py \
 	    --text_column_name sentence \
             --train_split_name test\
             --eval_split_name test\
-            --output_dir whisper-tiny-ft-dummy\
+            --output_dir whisper-small-flaxtest\
             --overwrite_output_dir\
             --num_train_epochs=1\
             --warmup_steps=8 \
             --do_train \
             --do_eval \
-            --num_train_steps 1000 \
-            --eval_steps 100 \
+            --num_train_steps 100 \
+	    --max_eval_samples 100 \
+            --eval_steps 50 \
             --learning_rate=2e-4 \
             --per_device_train_batch_size=2 \
-            --per_device_eval_batch_size=4 \
+            --per_device_eval_batch_size=2 \
             --predict_with_generate \
             --streaming=True \
             

whisper-small-flaxtest/added_tokens.json

@@ -0,0 +1,108 @@
+{
+  "<|af|>": 50326,
+  "<|am|>": 50333,
+  "<|ar|>": 50271,
+  "<|as|>": 50349,
+  "<|az|>": 50303,
+  "<|ba|>": 50354,
+  "<|be|>": 50329,
+  "<|bg|>": 50291,
+  "<|bn|>": 50301,
+  "<|bo|>": 50346,
+  "<|br|>": 50308,
+  "<|bs|>": 50314,
+  "<|ca|>": 50269,
+  "<|cs|>": 50282,
+  "<|cy|>": 50296,
+  "<|da|>": 50284,
+  "<|de|>": 50260,
+  "<|el|>": 50280,
+  "<|en|>": 50258,
+  "<|es|>": 50261,
+  "<|et|>": 50306,
+  "<|eu|>": 50309,
+  "<|fa|>": 50299,
+  "<|fi|>": 50276,
+  "<|fo|>": 50337,
+  "<|fr|>": 50264,
+  "<|gl|>": 50318,
+  "<|gu|>": 50332,
+  "<|haw|>": 50351,
+  "<|ha|>": 50353,
+  "<|hi|>": 50275,
+  "<|hr|>": 50290,
+  "<|ht|>": 50338,
+  "<|hu|>": 50285,
+  "<|hy|>": 50311,
+  "<|id|>": 50274,
+  "<|is|>": 50310,
+  "<|it|>": 50273,
+  "<|iw|>": 50278,
+  "<|ja|>": 50265,
+  "<|jw|>": 50355,
+  "<|ka|>": 50328,
+  "<|kk|>": 50315,
+  "<|km|>": 50322,
+  "<|kn|>": 50305,
+  "<|ko|>": 50263,
+  "<|la|>": 50293,
+  "<|lb|>": 50344,
+  "<|ln|>": 50352,
+  "<|lo|>": 50335,
+  "<|lt|>": 50292,
+  "<|lv|>": 50300,
+  "<|mg|>": 50348,
+  "<|mi|>": 50294,
+  "<|mk|>": 50307,
+  "<|ml|>": 50295,
+  "<|mn|>": 50313,
+  "<|mr|>": 50319,
+  "<|ms|>": 50281,
+  "<|mt|>": 50342,
+  "<|my|>": 50345,
+  "<|ne|>": 50312,
+  "<|nl|>": 50270,
+  "<|nn|>": 50341,
+  "<|nocaptions|>": 50361,
+  "<|notimestamps|>": 50362,
+  "<|no|>": 50287,
+  "<|oc|>": 50327,
+  "<|pa|>": 50320,
+  "<|pl|>": 50268,
+  "<|ps|>": 50339,
+  "<|pt|>": 50266,
+  "<|ro|>": 50283,
+  "<|ru|>": 50262,
+  "<|sa|>": 50343,
+  "<|sd|>": 50331,
+  "<|si|>": 50321,
+  "<|sk|>": 50297,
+  "<|sl|>": 50304,
+  "<|sn|>": 50323,
+  "<|so|>": 50325,
+  "<|sq|>": 50316,
+  "<|sr|>": 50302,
+  "<|startoflm|>": 50359,
+  "<|startofprev|>": 50360,
+  "<|startoftranscript|>": 50257,
+  "<|su|>": 50356,
+  "<|sv|>": 50272,
+  "<|sw|>": 50317,
+  "<|ta|>": 50286,
+  "<|te|>": 50298,
+  "<|tg|>": 50330,
+  "<|th|>": 50288,
+  "<|tk|>": 50340,
+  "<|tl|>": 50347,
+  "<|transcribe|>": 50358,
+  "<|translate|>": 50357,
+  "<|tr|>": 50267,
+  "<|tt|>": 50350,
+  "<|uk|>": 50279,
+  "<|ur|>": 50289,
+  "<|uz|>": 50336,
+  "<|vi|>": 50277,
+  "<|yi|>": 50334,
+  "<|yo|>": 50324,
+  "<|zh|>": 50259
+}

whisper-small-flaxtest/config.json

@@ -0,0 +1,143 @@
+{
+  "_name_or_path": "openai/whisper-tiny.en",
+  "activation_dropout": 0.0,
+  "activation_function": "gelu",
+  "apply_spec_augment": false,
+  "architectures": [
+    "WhisperForConditionalGeneration"
+  ],
+  "attention_dropout": 0.0,
+  "begin_suppress_tokens": [
+    220,
+    50256
+  ],
+  "bos_token_id": 50257,
+  "d_model": 384,
+  "decoder_attention_heads": 6,
+  "decoder_ffn_dim": 1536,
+  "decoder_layerdrop": 0.0,
+  "decoder_layers": 4,
+  "decoder_start_token_id": 50257,
+  "dropout": 0.0,
+  "encoder_attention_heads": 6,
+  "encoder_ffn_dim": 1536,
+  "encoder_layerdrop": 0.0,
+  "encoder_layers": 4,
+  "eos_token_id": 50256,
+  "forced_decoder_ids": [
+    [
+      1,
+      50362
+    ]
+  ],
+  "init_std": 0.02,
+  "is_encoder_decoder": true,
+  "mask_feature_length": 10,
+  "mask_feature_min_masks": 0,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_masks": 2,
+  "mask_time_prob": 0.05,
+  "max_length": 448,
+  "max_source_positions": 1500,
+  "max_target_positions": 448,
+  "model_type": "whisper",
+  "num_hidden_layers": 4,
+  "num_mel_bins": 80,
+  "pad_token_id": 50256,
+  "scale_embedding": false,
+  "suppress_tokens": [
+    1,
+    2,
+    7,
+    8,
+    9,
+    10,
+    14,
+    25,
+    26,
+    27,
+    28,
+    29,
+    31,
+    58,
+    59,
+    60,
+    61,
+    62,
+    63,
+    90,
+    91,
+    92,
+    93,
+    357,
+    366,
+    438,
+    532,
+    685,
+    705,
+    796,
+    930,
+    1058,
+    1220,
+    1267,
+    1279,
+    1303,
+    1343,
+    1377,
+    1391,
+    1635,
+    1782,
+    1875,
+    2162,
+    2361,
+    2488,
+    3467,
+    4008,
+    4211,
+    4600,
+    4808,
+    5299,
+    5855,
+    6329,
+    7203,
+    9609,
+    9959,
+    10563,
+    10786,
+    11420,
+    11709,
+    11907,
+    13163,
+    13697,
+    13700,
+    14808,
+    15306,
+    16410,
+    16791,
+    17992,
+    19203,
+    19510,
+    20724,
+    22305,
+    22935,
+    27007,
+    30109,
+    30420,
+    33409,
+    34949,
+    40283,
+    40493,
+    40549,
+    47282,
+    49146,
+    50257,
+    50359,
+    50360,
+    50361
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.27.0.dev0",
+  "use_cache": true,
+  "vocab_size": 51864
+}

whisper-small-flaxtest/events.out.tfevents.1677611724.t1v-n-d163ce9a-w-0.1583171.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9277daa3c525efff75929a791e43739542c8b06300cd7e0063ffc45416db3592
+size 40

whisper-small-flaxtest/events.out.tfevents.1677613551.t1v-n-d163ce9a-w-0.1702844.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:090e22643a37a0ccc28a8d5654c3d6c08569bfb99b5779934dad613905abe4bb
+size 40

whisper-small-flaxtest/events.out.tfevents.1677613844.t1v-n-d163ce9a-w-0.1706687.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e1c74cb53e0effd274e51aae5374d6a418525fb8eb2933e9108f568203ddd44
+size 40

whisper-small-flaxtest/events.out.tfevents.1677614511.t1v-n-d163ce9a-w-0.1904376.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bef180e271f9a1f610835e098dfe01c68fbe5b791708714f866eaeeae76c1730
+size 40

whisper-small-flaxtest/events.out.tfevents.1677615119.t1v-n-d163ce9a-w-0.2101561.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d63e63d1335b2c7c223774e073f1ffeada56385d68fdbd3fe717767749d153b9
+size 40

whisper-small-flaxtest/events.out.tfevents.1677615611.t1v-n-d163ce9a-w-0.2298739.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fe126e87c513b3ef6f7a5bd6333f15c9aeb960d710d91208127f0195a4d336e
+size 40

whisper-small-flaxtest/normalizer.json

@@ -0,0 +1,1742 @@
+{
+  "accessorise": "accessorize",
+  "accessorised": "accessorized",
+  "accessorises": "accessorizes",
+  "accessorising": "accessorizing",
+  "acclimatisation": "acclimatization",
+  "acclimatise": "acclimatize",
+  "acclimatised": "acclimatized",
+  "acclimatises": "acclimatizes",
+  "acclimatising": "acclimatizing",
+  "accoutrements": "accouterments",
+  "aeon": "eon",
+  "aeons": "eons",
+  "aerogramme": "aerogram",
+  "aerogrammes": "aerograms",
+  "aeroplane": "airplane",
+  "aeroplanes": "airplanes",
+  "aesthete": "esthete",
+  "aesthetes": "esthetes",
+  "aesthetic": "esthetic",
+  "aesthetically": "esthetically",
+  "aesthetics": "esthetics",
+  "aetiology": "etiology",
+  "ageing": "aging",
+  "aggrandisement": "aggrandizement",
+  "agonise": "agonize",
+  "agonised": "agonized",
+  "agonises": "agonizes",
+  "agonising": "agonizing",
+  "agonisingly": "agonizingly",
+  "almanack": "almanac",
+  "almanacks": "almanacs",
+  "aluminium": "aluminum",
+  "amortisable": "amortizable",
+  "amortisation": "amortization",
+  "amortisations": "amortizations",
+  "amortise": "amortize",
+  "amortised": "amortized",
+  "amortises": "amortizes",
+  "amortising": "amortizing",
+  "amphitheatre": "amphitheater",
+  "amphitheatres": "amphitheaters",
+  "anaemia": "anemia",
+  "anaemic": "anemic",
+  "anaesthesia": "anesthesia",
+  "anaesthetic": "anesthetic",
+  "anaesthetics": "anesthetics",
+  "anaesthetise": "anesthetize",
+  "anaesthetised": "anesthetized",
+  "anaesthetises": "anesthetizes",
+  "anaesthetising": "anesthetizing",
+  "anaesthetist": "anesthetist",
+  "anaesthetists": "anesthetists",
+  "anaesthetize": "anesthetize",
+  "anaesthetized": "anesthetized",
+  "anaesthetizes": "anesthetizes",
+  "anaesthetizing": "anesthetizing",
+  "analogue": "analog",
+  "analogues": "analogs",
+  "analyse": "analyze",
+  "analysed": "analyzed",
+  "analyses": "analyzes",
+  "analysing": "analyzing",
+  "anglicise": "anglicize",
+  "anglicised": "anglicized",
+  "anglicises": "anglicizes",
+  "anglicising": "anglicizing",
+  "annualised": "annualized",
+  "antagonise": "antagonize",
+  "antagonised": "antagonized",
+  "antagonises": "antagonizes",
+  "antagonising": "antagonizing",
+  "apologise": "apologize",
+  "apologised": "apologized",
+  "apologises": "apologizes",
+  "apologising": "apologizing",
+  "appal": "appall",
+  "appals": "appalls",
+  "appetiser": "appetizer",
+  "appetisers": "appetizers",
+  "appetising": "appetizing",
+  "appetisingly": "appetizingly",
+  "arbour": "arbor",
+  "arbours": "arbors",
+  "archaeologically": "archeologically",
+  "archaeologist": "archeologist",
+  "archaeologists": "archeologists",
+  "archaeology": "archeology</span>",
+  "archeological": "archaeological",
+  "ardour": "ardor",
+  "armour": "armor",
+  "armoured": "armored",
+  "armourer": "armorer",
+  "armourers": "armorers",
+  "armouries": "armories",
+  "armoury": "armory",
+  "artefact": "artifact",
+  "artefacts": "artifacts",
+  "authorise": "authorize",
+  "authorised": "authorized",
+  "authorises": "authorizes",
+  "authorising": "authorizing",
+  "axe": "ax",
+  "backpedalled": "backpedaled",
+  "backpedalling": "backpedaling",
+  "bannister": "banister",
+  "bannisters": "banisters",
+  "baptise": "baptize",
+  "baptised": "baptized",
+  "baptises": "baptizes",
+  "baptising": "baptizing",
+  "bastardise": "bastardize",
+  "bastardised": "bastardized",
+  "bastardises": "bastardizes",
+  "bastardising": "bastardizing",
+  "battleax": "battleaxe",
+  "baulk": "balk",
+  "baulked": "balked",
+  "baulking": "balking",
+  "baulks": "balks",
+  "bedevilled": "bedeviled",
+  "bedevilling": "bedeviling",
+  "behaviour": "behavior",
+  "behavioural": "behavioral",
+  "behaviourism": "behaviorism",
+  "behaviourist": "behaviorist",
+  "behaviourists": "behaviorists",
+  "behaviours": "behaviors",
+  "behove": "behoove",
+  "behoved": "behooved",
+  "behoves": "behooves",
+  "bejewelled": "bejeweled",
+  "belabour": "belabor",
+  "belaboured": "belabored",
+  "belabouring": "belaboring",
+  "belabours": "belabors",
+  "bevelled": "beveled",
+  "bevvies": "bevies",
+  "bevvy": "bevy",
+  "biassed": "biased",
+  "biassing": "biasing",
+  "bingeing": "binging",
+  "bougainvillaea": "bougainvillea",
+  "bougainvillaeas": "bougainvilleas",
+  "bowdlerise": "bowdlerize",
+  "bowdlerised": "bowdlerized",
+  "bowdlerises": "bowdlerizes",
+  "bowdlerising": "bowdlerizing",
+  "breathalyse": "breathalyze",
+  "breathalysed": "breathalyzed",
+  "breathalyser": "breathalyzer",
+  "breathalysers": "breathalyzers",
+  "breathalyses": "breathalyzes",
+  "breathalysing": "breathalyzing",
+  "brutalise": "brutalize",
+  "brutalised": "brutalized",
+  "brutalises": "brutalizes",
+  "brutalising": "brutalizing",
+  "busses": "buses",
+  "bussing": "busing",
+  "caesarean": "cesarean",
+  "caesareans": "cesareans",
+  "calibre": "caliber",
+  "calibres": "calibers",
+  "calliper": "caliper",
+  "callipers": "calipers",
+  "callisthenics": "calisthenics",
+  "canalise": "canalize",
+  "canalised": "canalized",
+  "canalises": "canalizes",
+  "canalising": "canalizing",
+  "cancelation": "cancellation",
+  "cancelations": "cancellations",
+  "cancelled": "canceled",
+  "cancelling": "canceling",
+  "candour": "candor",
+  "cannibalise": "cannibalize",
+  "cannibalised": "cannibalized",
+  "cannibalises": "cannibalizes",
+  "cannibalising": "cannibalizing",
+  "canonise": "canonize",
+  "canonised": "canonized",
+  "canonises": "canonizes",
+  "canonising": "canonizing",
+  "capitalise": "capitalize",
+  "capitalised": "capitalized",
+  "capitalises": "capitalizes",
+  "capitalising": "capitalizing",
+  "caramelise": "caramelize",
+  "caramelised": "caramelized",
+  "caramelises": "caramelizes",
+  "caramelising": "caramelizing",
+  "carbonise": "carbonize",
+  "carbonised": "carbonized",
+  "carbonises": "carbonizes",
+  "carbonising": "carbonizing",
+  "carolled": "caroled",
+  "carolling": "caroling",
+  "catalogue": "catalog",
+  "catalogued": "cataloged",
+  "catalogues": "catalogs",
+  "cataloguing": "cataloging",
+  "catalyse": "catalyze",
+  "catalysed": "catalyzed",
+  "catalyses": "catalyzes",
+  "catalysing": "catalyzing",
+  "categorise": "categorize",
+  "categorised": "categorized",
+  "categorises": "categorizes",
+  "categorising": "categorizing",
+  "cauterise": "cauterize",
+  "cauterised": "cauterized",
+  "cauterises": "cauterizes",
+  "cauterising": "cauterizing",
+  "cavilled": "caviled",
+  "cavilling": "caviling",
+  "centigramme": "centigram",
+  "centigrammes": "centigrams",
+  "centilitre": "centiliter",
+  "centilitres": "centiliters",
+  "centimetre": "centimeter",
+  "centimetres": "centimeters",
+  "centralise": "centralize",
+  "centralised": "centralized",
+  "centralises": "centralizes",
+  "centralising": "centralizing",
+  "centre": "center",
+  "centred": "centered",
+  "centrefold": "centerfold",
+  "centrefolds": "centerfolds",
+  "centrepiece": "centerpiece",
+  "centrepieces": "centerpieces",
+  "centres": "centers",
+  "channelled": "channeled",
+  "channelling": "channeling",
+  "characterise": "characterize",
+  "characterised": "characterized",
+  "characterises": "characterizes",
+  "characterising": "characterizing",
+  "cheque": "check",
+  "chequebook": "checkbook",
+  "chequebooks": "checkbooks",
+  "chequered": "checkered",
+  "cheques": "checks",
+  "chilli": "chili",
+  "chimaera": "chimera",
+  "chimaeras": "chimeras",
+  "chiselled": "chiseled",
+  "chiselling": "chiseling",
+  "circularise": "circularize",
+  "circularised": "circularized",
+  "circularises": "circularizes",
+  "circularising": "circularizing",
+  "civilise": "civilize",
+  "civilised": "civilized",
+  "civilises": "civilizes",
+  "civilising": "civilizing",
+  "clamour": "clamor",
+  "clamoured": "clamored",
+  "clamouring": "clamoring",
+  "clamours": "clamors",
+  "clangour": "clangor",
+  "clarinettist": "clarinetist",
+  "clarinettists": "clarinetists",
+  "collectivise": "collectivize",
+  "collectivised": "collectivized",
+  "collectivises": "collectivizes",
+  "collectivising": "collectivizing",
+  "colonisation": "colonization",
+  "colonise": "colonize",
+  "colonised": "colonized",
+  "coloniser": "colonizer",
+  "colonisers": "colonizers",
+  "colonises": "colonizes",
+  "colonising": "colonizing",
+  "colour": "color",
+  "colourant": "colorant",
+  "colourants": "colorants",
+  "coloured": "colored",
+  "coloureds": "coloreds",
+  "colourful": "colorful",
+  "colourfully": "colorfully",
+  "colouring": "coloring",
+  "colourize": "colorize",
+  "colourized": "colorized",
+  "colourizes": "colorizes",
+  "colourizing": "colorizing",
+  "colourless": "colorless",
+  "colours": "colors",
+  "commercialise": "commercialize",
+  "commercialised": "commercialized",
+  "commercialises": "commercializes",
+  "commercialising": "commercializing",
+  "compartmentalise": "compartmentalize",
+  "compartmentalised": "compartmentalized",
+  "compartmentalises": "compartmentalizes",
+  "compartmentalising": "compartmentalizing",
+  "computerise": "computerize",
+  "computerised": "computerized",
+  "computerises": "computerizes",
+  "computerising": "computerizing",
+  "conceptualise": "conceptualize",
+  "conceptualised": "conceptualized",
+  "conceptualises": "conceptualizes",
+  "conceptualising": "conceptualizing",
+  "connexion": "connection",
+  "connexions": "connections",
+  "contextualise": "contextualize",
+  "contextualised": "contextualized",
+  "contextualises": "contextualizes",
+  "contextualising": "contextualizing",
+  "cosier": "cozier",
+  "cosies": "cozies",
+  "cosiest": "coziest",
+  "cosily": "cozily",
+  "cosiness": "coziness",
+  "cosy": "cozy",
+  "councillor": "councilor",
+  "councillors": "councilors",
+  "counselled": "counseled",
+  "counselling": "counseling",
+  "counsellor": "counselor",
+  "counsellors": "counselors",
+  "crenelated": "crenellated",
+  "criminalise": "criminalize",
+  "criminalised": "criminalized",
+  "criminalises": "criminalizes",
+  "criminalising": "criminalizing",
+  "criticise": "criticize",
+  "criticised": "criticized",
+  "criticises": "criticizes",
+  "criticising": "criticizing",
+  "crueller": "crueler",
+  "cruellest": "cruelest",
+  "crystallisation": "crystallization",
+  "crystallise": "crystallize",
+  "crystallised": "crystallized",
+  "crystallises": "crystallizes",
+  "crystallising": "crystallizing",
+  "cudgelled": "cudgeled",
+  "cudgelling": "cudgeling",
+  "customise": "customize",
+  "customised": "customized",
+  "customises": "customizes",
+  "customising": "customizing",
+  "cypher": "cipher",
+  "cyphers": "ciphers",
+  "decentralisation": "decentralization",
+  "decentralise": "decentralize",
+  "decentralised": "decentralized",
+  "decentralises": "decentralizes",
+  "decentralising": "decentralizing",
+  "decriminalisation": "decriminalization",
+  "decriminalise": "decriminalize",
+  "decriminalised": "decriminalized",
+  "decriminalises": "decriminalizes",
+  "decriminalising": "decriminalizing",
+  "defence": "defense",
+  "defenceless": "defenseless",
+  "defences": "defenses",
+  "dehumanisation": "dehumanization",
+  "dehumanise": "dehumanize",
+  "dehumanised": "dehumanized",
+  "dehumanises": "dehumanizes",
+  "dehumanising": "dehumanizing",
+  "demeanour": "demeanor",
+  "demilitarisation": "demilitarization",
+  "demilitarise": "demilitarize",
+  "demilitarised": "demilitarized",
+  "demilitarises": "demilitarizes",
+  "demilitarising": "demilitarizing",
+  "demobilisation": "demobilization",
+  "demobilise": "demobilize",
+  "demobilised": "demobilized",
+  "demobilises": "demobilizes",
+  "demobilising": "demobilizing",
+  "democratisation": "democratization",
+  "democratise": "democratize",
+  "democratised": "democratized",
+  "democratises": "democratizes",
+  "democratising": "democratizing",
+  "demonise": "demonize",
+  "demonised": "demonized",
+  "demonises": "demonizes",
+  "demonising": "demonizing",
+  "demoralisation": "demoralization",
+  "demoralise": "demoralize",
+  "demoralised": "demoralized",
+  "demoralises": "demoralizes",
+  "demoralising": "demoralizing",
+  "denationalisation": "denationalization",
+  "denationalise": "denationalize",
+  "denationalised": "denationalized",
+  "denationalises": "denationalizes",
+  "denationalising": "denationalizing",
+  "deodorise": "deodorize",
+  "deodorised": "deodorized",
+  "deodorises": "deodorizes",
+  "deodorising": "deodorizing",
+  "depersonalise": "depersonalize",
+  "depersonalised": "depersonalized",
+  "depersonalises": "depersonalizes",
+  "depersonalising": "depersonalizing",
+  "deputise": "deputize",
+  "deputised": "deputized",
+  "deputises": "deputizes",
+  "deputising": "deputizing",
+  "desensitisation": "desensitization",
+  "desensitise": "desensitize",
+  "desensitised": "desensitized",
+  "desensitises": "desensitizes",
+  "desensitising": "desensitizing",
+  "destabilisation": "destabilization",
+  "destabilise": "destabilize",
+  "destabilised": "destabilized",
+  "destabilises": "destabilizes",
+  "destabilising": "destabilizing",
+  "dialled": "dialed",
+  "dialling": "dialing",
+  "dialogue": "dialog",
+  "dialogues": "dialogs",
+  "diarrhoea": "diarrhea",
+  "digitise": "digitize",
+  "digitised": "digitized",
+  "digitises": "digitizes",
+  "digitising": "digitizing",
+  "disc": "disk",
+  "discolour": "discolor",
+  "discoloured": "discolored",
+  "discolouring": "discoloring",
+  "discolours": "discolors",
+  "discs": "disks",
+  "disembowelled": "disemboweled",
+  "disembowelling": "disemboweling",
+  "disfavour": "disfavor",
+  "dishevelled": "disheveled",
+  "dishonour": "dishonor",
+  "dishonourable": "dishonorable",
+  "dishonourably": "dishonorably",
+  "dishonoured": "dishonored",
+  "dishonouring": "dishonoring",
+  "dishonours": "dishonors",
+  "disorganisation": "disorganization",
+  "disorganised": "disorganized",
+  "distil": "distill",
+  "distils": "distills",
+  "dramatisation": "dramatization",
+  "dramatisations": "dramatizations",
+  "dramatise": "dramatize",
+  "dramatised": "dramatized",
+  "dramatises": "dramatizes",
+  "dramatising": "dramatizing",
+  "draught": "draft",
+  "draughtboard": "draftboard",
+  "draughtboards": "draftboards",
+  "draughtier": "draftier",
+  "draughtiest": "draftiest",
+  "draughts": "drafts",
+  "draughtsman": "draftsman",
+  "draughtsmanship": "draftsmanship",
+  "draughtsmen": "draftsmen",
+  "draughtswoman": "draftswoman",
+  "draughtswomen": "draftswomen",
+  "draughty": "drafty",
+  "drivelled": "driveled",
+  "drivelling": "driveling",
+  "duelled": "dueled",
+  "duelling": "dueling",
+  "economise": "economize",
+  "economised": "economized",
+  "economises": "economizes",
+  "economising": "economizing",
+  "editorialise": "editorialize",
+  "editorialised": "editorialized",
+  "editorialises": "editorializes",
+  "editorialising": "editorializing",
+  "edoema": "edema",
+  "empathise": "empathize",
+  "empathised": "empathized",
+  "empathises": "empathizes",
+  "empathising": "empathizing",
+  "emphasise": "emphasize",
+  "emphasised": "emphasized",
+  "emphasises": "emphasizes",
+  "emphasising": "emphasizing",
+  "enamelled": "enameled",
+  "enamelling": "enameling",
+  "enamoured": "enamored",
+  "encyclopaedia": "encyclopedia",
+  "encyclopaedias": "encyclopedias",
+  "encyclopaedic": "encyclopedic",
+  "endeavour": "endeavor",
+  "endeavoured": "endeavored",
+  "endeavouring": "endeavoring",
+  "endeavours": "endeavors",
+  "energise": "energize",
+  "energised": "energized",
+  "energises": "energizes",
+  "energising": "energizing",
+  "enrol": "enroll",
+  "enrols": "enrolls",
+  "enthral": "enthrall",
+  "enthrals": "enthralls",
+  "epaulette": "epaulet",
+  "epaulettes": "epaulets",
+  "epicentre": "epicenter",
+  "epicentres": "epicenters",
+  "epilogue": "epilog",
+  "epilogues": "epilogs",
+  "epitomise": "epitomize",
+  "epitomised": "epitomized",
+  "epitomises": "epitomizes",
+  "epitomising": "epitomizing",
+  "equalisation": "equalization",
+  "equalise": "equalize",
+  "equalised": "equalized",
+  "equaliser": "equalizer",
+  "equalisers": "equalizers",
+  "equalises": "equalizes",
+  "equalising": "equalizing",
+  "eulogise": "eulogize",
+  "eulogised": "eulogized",
+  "eulogises": "eulogizes",
+  "eulogising": "eulogizing",
+  "evangelise": "evangelize",
+  "evangelised": "evangelized",
+  "evangelises": "evangelizes",
+  "evangelising": "evangelizing",
+  "exorcise": "exorcize",
+  "exorcised": "exorcized",
+  "exorcises": "exorcizes",
+  "exorcising": "exorcizing",
+  "extemporisation": "extemporization",
+  "extemporise": "extemporize",
+  "extemporised": "extemporized",
+  "extemporises": "extemporizes",
+  "extemporising": "extemporizing",
+  "externalisation": "externalization",
+  "externalisations": "externalizations",
+  "externalise": "externalize",
+  "externalised": "externalized",
+  "externalises": "externalizes",
+  "externalising": "externalizing",
+  "factorise": "factorize",
+  "factorised": "factorized",
+  "factorises": "factorizes",
+  "factorising": "factorizing",
+  "faecal": "fecal",
+  "faeces": "feces",
+  "familiarisation": "familiarization",
+  "familiarise": "familiarize",
+  "familiarised": "familiarized",
+  "familiarises": "familiarizes",
+  "familiarising": "familiarizing",
+  "fantasise": "fantasize",
+  "fantasised": "fantasized",
+  "fantasises": "fantasizes",
+  "fantasising": "fantasizing",
+  "favour": "favor",
+  "favourable": "favorable",
+  "favourably": "favorably",
+  "favoured": "favored",
+  "favouring": "favoring",
+  "favourite": "favorite",
+  "favourites": "favorites",
+  "favouritism": "favoritism",
+  "favours": "favors",
+  "feminise": "feminize",
+  "feminised": "feminized",
+  "feminises": "feminizes",
+  "feminising": "feminizing",
+  "fertilisation": "fertilization",
+  "fertilise": "fertilize",
+  "fertilised": "fertilized",
+  "fertiliser": "fertilizer",
+  "fertilisers": "fertilizers",
+  "fertilises": "fertilizes",
+  "fertilising": "fertilizing",
+  "fervour": "fervor",
+  "fibre": "fiber",
+  "fibreglass": "fiberglass",
+  "fibres": "fibers",
+  "fictionalisation": "fictionalization",
+  "fictionalisations": "fictionalizations",
+  "fictionalise": "fictionalize",
+  "fictionalised": "fictionalized",
+  "fictionalises": "fictionalizes",
+  "fictionalising": "fictionalizing",
+  "fillet": "filet",
+  "filleted": "fileted",
+  "filleting": "fileting",
+  "fillets": "filets",
+  "finalisation": "finalization",
+  "finalise": "finalize",
+  "finalised": "finalized",
+  "finalises": "finalizes",
+  "finalising": "finalizing",
+  "flautist": "flutist",
+  "flautists": "flutists",
+  "flavour": "flavor",
+  "flavoured": "flavored",
+  "flavouring": "flavoring",
+  "flavourings": "flavorings",
+  "flavourless": "flavorless",
+  "flavours": "flavors",
+  "flavoursome": "flavorsome",
+  "flyer / flier": "flier / flyer",
+  "foetal": "fetal",
+  "foetid": "fetid",
+  "foetus": "fetus",
+  "foetuses": "fetuses",
+  "formalisation": "formalization",
+  "formalise": "formalize",
+  "formalised": "formalized",
+  "formalises": "formalizes",
+  "formalising": "formalizing",
+  "fossilisation": "fossilization",
+  "fossilise": "fossilize",
+  "fossilised": "fossilized",
+  "fossilises": "fossilizes",
+  "fossilising": "fossilizing",
+  "fraternisation": "fraternization",
+  "fraternise": "fraternize",
+  "fraternised": "fraternized",
+  "fraternises": "fraternizes",
+  "fraternising": "fraternizing",
+  "fulfil": "fulfill",
+  "fulfilment": "fulfillment",
+  "fulfils": "fulfills",
+  "funnelled": "funneled",
+  "funnelling": "funneling",
+  "gage": "gauge",
+  "gaged": "gauged",
+  "gages": "gauges",
+  "gaging": "gauging",
+  "galvanise": "galvanize",
+  "galvanised": "galvanized",
+  "galvanises": "galvanizes",
+  "galvanising": "galvanizing",
+  "gambolled": "gamboled",
+  "gambolling": "gamboling",
+  "gaol": "jail",
+  "gaolbird": "jailbird",
+  "gaolbirds": "jailbirds",
+  "gaolbreak": "jailbreak",
+  "gaolbreaks": "jailbreaks",
+  "gaoled": "jailed",
+  "gaoler": "jailer",
+  "gaolers": "jailers",
+  "gaoling": "jailing",
+  "gaols": "jails",
+  "gasses": "gases",
+  "generalisation": "generalization",
+  "generalisations": "generalizations",
+  "generalise": "generalize",
+  "generalised": "generalized",
+  "generalises": "generalizes",
+  "generalising": "generalizing",
+  "ghettoise": "ghettoize",
+  "ghettoised": "ghettoized",
+  "ghettoises": "ghettoizes",
+  "ghettoising": "ghettoizing",
+  "gipsies": "gypsies",
+  "glamor": "glamour",
+  "glamorise": "glamorize",
+  "glamorised": "glamorized",
+  "glamorises": "glamorizes",
+  "glamorising": "glamorizing",
+  "globalisation": "globalization",
+  "globalise": "globalize",
+  "globalised": "globalized",
+  "globalises": "globalizes",
+  "globalising": "globalizing",
+  "glueing": "gluing",
+  "goitre": "goiter",
+  "goitres": "goiters",
+  "gonorrhoea": "gonorrhea",
+  "gramme": "gram",
+  "grammes": "grams",
+  "gravelled": "graveled",
+  "grey": "gray",
+  "greyed": "grayed",
+  "greying": "graying",
+  "greyish": "grayish",
+  "greyness": "grayness",
+  "greys": "grays",
+  "grovelled": "groveled",
+  "grovelling": "groveling",
+  "groyne": "groin",
+  "groynes": "groins",
+  "gruelling": "grueling",
+  "gruellingly": "gruelingly",
+  "gryphon": "griffin",
+  "gryphons": "griffins",
+  "gynaecological": "gynecological",
+  "gynaecologist": "gynecologist",
+  "gynaecologists": "gynecologists",
+  "gynaecology": "gynecology",
+  "haematological": "hematological",
+  "haematologist": "hematologist",
+  "haematologists": "hematologists",
+  "haematology": "hematology",
+  "haemoglobin": "hemoglobin",
+  "haemophilia": "hemophilia",
+  "haemophiliac": "hemophiliac",
+  "haemophiliacs": "hemophiliacs",
+  "haemorrhage": "hemorrhage",
+  "haemorrhaged": "hemorrhaged",
+  "haemorrhages": "hemorrhages",
+  "haemorrhaging": "hemorrhaging",
+  "haemorrhoids": "hemorrhoids",
+  "harbour": "harbor",
+  "harboured": "harbored",
+  "harbouring": "harboring",
+  "harbours": "harbors",
+  "harmonisation": "harmonization",
+  "harmonise": "harmonize",
+  "harmonised": "harmonized",
+  "harmonises": "harmonizes",
+  "harmonising": "harmonizing",
+  "homoeopath": "homeopath",
+  "homoeopathic": "homeopathic",
+  "homoeopaths": "homeopaths",
+  "homoeopathy": "homeopathy",
+  "homogenise": "homogenize",
+  "homogenised": "homogenized",
+  "homogenises": "homogenizes",
+  "homogenising": "homogenizing",
+  "honour": "honor",
+  "honourable": "honorable",
+  "honourably": "honorably",
+  "honoured": "honored",
+  "honouring": "honoring",
+  "honours": "honors",
+  "hospitalisation": "hospitalization",
+  "hospitalise": "hospitalize",
+  "hospitalised": "hospitalized",
+  "hospitalises": "hospitalizes",
+  "hospitalising": "hospitalizing",
+  "humanise": "humanize",
+  "humanised": "humanized",
+  "humanises": "humanizes",
+  "humanising": "humanizing",
+  "humour": "humor",
+  "humoured": "humored",
+  "humouring": "humoring",
+  "humourless": "humorless",
+  "humours": "humors",
+  "hybridise": "hybridize",
+  "hybridised": "hybridized",
+  "hybridises": "hybridizes",
+  "hybridising": "hybridizing",
+  "hypnotise": "hypnotize",
+  "hypnotised": "hypnotized",
+  "hypnotises": "hypnotizes",
+  "hypnotising": "hypnotizing",
+  "hypothesise": "hypothesize",
+  "hypothesised": "hypothesized",
+  "hypothesises": "hypothesizes",
+  "hypothesising": "hypothesizing",
+  "idealisation": "idealization",
+  "idealise": "idealize",
+  "idealised": "idealized",
+  "idealises": "idealizes",
+  "idealising": "idealizing",
+  "idolise": "idolize",
+  "idolised": "idolized",
+  "idolises": "idolizes",
+  "idolising": "idolizing",
+  "immobilisation": "immobilization",
+  "immobilise": "immobilize",
+  "immobilised": "immobilized",
+  "immobiliser": "immobilizer",
+  "immobilisers": "immobilizers",
+  "immobilises": "immobilizes",
+  "immobilising": "immobilizing",
+  "immortalise": "immortalize",
+  "immortalised": "immortalized",
+  "immortalises": "immortalizes",
+  "immortalising": "immortalizing",
+  "immunisation": "immunization",
+  "immunise": "immunize",
+  "immunised": "immunized",
+  "immunises": "immunizes",
+  "immunising": "immunizing",
+  "impanelled": "impaneled",
+  "impanelling": "impaneling",
+  "imperilled": "imperiled",
+  "imperilling": "imperiling",
+  "individualise": "individualize",
+  "individualised": "individualized",
+  "individualises": "individualizes",
+  "individualising": "individualizing",
+  "industrialise": "industrialize",
+  "industrialised": "industrialized",
+  "industrialises": "industrializes",
+  "industrialising": "industrializing",
+  "inflexion": "inflection",
+  "inflexions": "inflections",
+  "initialise": "initialize",
+  "initialised": "initialized",
+  "initialises": "initializes",
+  "initialising": "initializing",
+  "initialled": "initialed",
+  "initialling": "initialing",
+  "instal": "install",
+  "instalment": "installment",
+  "instalments": "installments",
+  "instals": "installs",
+  "instil": "instill",
+  "instils": "instills",
+  "institutionalisation": "institutionalization",
+  "institutionalise": "institutionalize",
+  "institutionalised": "institutionalized",
+  "institutionalises": "institutionalizes",
+  "institutionalising": "institutionalizing",
+  "intellectualise": "intellectualize",
+  "intellectualised": "intellectualized",
+  "intellectualises": "intellectualizes",
+  "intellectualising": "intellectualizing",
+  "internalisation": "internalization",
+  "internalise": "internalize",
+  "internalised": "internalized",
+  "internalises": "internalizes",
+  "internalising": "internalizing",
+  "internationalisation": "internationalization",
+  "internationalise": "internationalize",
+  "internationalised": "internationalized",
+  "internationalises": "internationalizes",
+  "internationalising": "internationalizing",
+  "ionisation": "ionization",
+  "ionise": "ionize",
+  "ionised": "ionized",
+  "ioniser": "ionizer",
+  "ionisers": "ionizers",
+  "ionises": "ionizes",
+  "ionising": "ionizing",
+  "italicise": "italicize",
+  "italicised": "italicized",
+  "italicises": "italicizes",
+  "italicising": "italicizing",
+  "itemise": "itemize",
+  "itemised": "itemized",
+  "itemises": "itemizes",
+  "itemising": "itemizing",
+  "jeopardise": "jeopardize",
+  "jeopardised": "jeopardized",
+  "jeopardises": "jeopardizes",
+  "jeopardising": "jeopardizing",
+  "jewelled": "jeweled",
+  "jeweller": "jeweler",
+  "jewellers": "jewelers",
+  "jewellery": "jewelry",
+  "judgement": "judgment",
+  "kilogramme": "kilogram",
+  "kilogrammes": "kilograms",
+  "kilometre": "kilometer",
+  "kilometres": "kilometers",
+  "labelled": "labeled",
+  "labelling": "labeling",
+  "labour": "labor",
+  "laboured": "labored",
+  "labourer": "laborer",
+  "labourers": "laborers",
+  "labouring": "laboring",
+  "labours": "labors",
+  "lacklustre": "lackluster",
+  "legalisation": "legalization",
+  "legalise": "legalize",
+  "legalised": "legalized",
+  "legalises": "legalizes",
+  "legalising": "legalizing",
+  "legitimise": "legitimize",
+  "legitimised": "legitimized",
+  "legitimises": "legitimizes",
+  "legitimising": "legitimizing",
+  "leukaemia": "leukemia",
+  "levelled": "leveled",
+  "leveller": "leveler",
+  "levellers": "levelers",
+  "levelling": "leveling",
+  "libelled": "libeled",
+  "libelling": "libeling",
+  "libellous": "libelous",
+  "liberalisation": "liberalization",
+  "liberalise": "liberalize",
+  "liberalised": "liberalized",
+  "liberalises": "liberalizes",
+  "liberalising": "liberalizing",
+  "licence": "license",
+  "licenced": "licensed",
+  "licences": "licenses",
+  "licencing": "licensing",
+  "likeable": "likable",
+  "lionisation": "lionization",
+  "lionise": "lionize",
+  "lionised": "lionized",
+  "lionises": "lionizes",
+  "lionising": "lionizing",
+  "liquidise": "liquidize",
+  "liquidised": "liquidized",
+  "liquidiser": "liquidizer",
+  "liquidisers": "liquidizers",
+  "liquidises": "liquidizes",
+  "liquidising": "liquidizing",
+  "litre": "liter",
+  "litres": "liters",
+  "localise": "localize",
+  "localised": "localized",
+  "localises": "localizes",
+  "localising": "localizing",
+  "louvre": "louver",
+  "louvred": "louvered",
+  "louvres": "louvers",
+  "lustre": "luster",
+  "magnetise": "magnetize",
+  "magnetised": "magnetized",
+  "magnetises": "magnetizes",
+  "magnetising": "magnetizing",
+  "manoeuvrability": "maneuverability",
+  "manoeuvrable": "maneuverable",
+  "manoeuvre": "maneuver",
+  "manoeuvred": "maneuvered",
+  "manoeuvres": "maneuvers",
+  "manoeuvring": "maneuvering",
+  "manoeuvrings": "maneuverings",
+  "marginalisation": "marginalization",
+  "marginalise": "marginalize",
+  "marginalised": "marginalized",
+  "marginalises": "marginalizes",
+  "marginalising": "marginalizing",
+  "marshalled": "marshaled",
+  "marshalling": "marshaling",
+  "marvelled": "marveled",
+  "marvelling": "marveling",
+  "marvellous": "marvelous",
+  "marvellously": "marvelously",
+  "materialisation": "materialization",
+  "materialise": "materialize",
+  "materialised": "materialized",
+  "materialises": "materializes",
+  "materialising": "materializing",
+  "maximisation": "maximization",
+  "maximise": "maximize",
+  "maximised": "maximized",
+  "maximises": "maximizes",
+  "maximising": "maximizing",
+  "meagre": "meager",
+  "mechanisation": "mechanization",
+  "mechanise": "mechanize",
+  "mechanised": "mechanized",
+  "mechanises": "mechanizes",
+  "mechanising": "mechanizing",
+  "mediaeval": "medieval",
+  "memorialise": "memorialize",
+  "memorialised": "memorialized",
+  "memorialises": "memorializes",
+  "memorialising": "memorializing",
+  "memorise": "memorize",
+  "memorised": "memorized",
+  "memorises": "memorizes",
+  "memorising": "memorizing",
+  "mesmerise": "mesmerize",
+  "mesmerised": "mesmerized",
+  "mesmerises": "mesmerizes",
+  "mesmerising": "mesmerizing",
+  "metabolise": "metabolize",
+  "metabolised": "metabolized",
+  "metabolises": "metabolizes",
+  "metabolising": "metabolizing",
+  "metre": "meter",
+  "metres": "meters",
+  "mhm": "hmm",
+  "micrometre": "micrometer",
+  "micrometres": "micrometers",
+  "militarise": "militarize",
+  "militarised": "militarized",
+  "militarises": "militarizes",
+  "militarising": "militarizing",
+  "milligramme": "milligram",
+  "milligrammes": "milligrams",
+  "millilitre": "milliliter",
+  "millilitres": "milliliters",
+  "millimetre": "millimeter",
+  "millimetres": "millimeters",
+  "miniaturisation": "miniaturization",
+  "miniaturise": "miniaturize",
+  "miniaturised": "miniaturized",
+  "miniaturises": "miniaturizes",
+  "miniaturising": "miniaturizing",
+  "minibusses": "minibuses",
+  "minimise": "minimize",
+  "minimised": "minimized",
+  "minimises": "minimizes",
+  "minimising": "minimizing",
+  "misbehaviour": "misbehavior",
+  "misdemeanour": "misdemeanor",
+  "misdemeanours": "misdemeanors",
+  "misspelt": "misspelled",
+  "mitre": "miter",
+  "mitres": "miters",
+  "mm": "hmm",
+  "mmm": "hmm",
+  "mobilisation": "mobilization",
+  "mobilise": "mobilize",
+  "mobilised": "mobilized",
+  "mobilises": "mobilizes",
+  "mobilising": "mobilizing",
+  "modelled": "modeled",
+  "modeller": "modeler",
+  "modellers": "modelers",
+  "modelling": "modeling",
+  "modernise": "modernize",
+  "modernised": "modernized",
+  "modernises": "modernizes",
+  "modernising": "modernizing",
+  "moisturise": "moisturize",
+  "moisturised": "moisturized",
+  "moisturiser": "moisturizer",
+  "moisturisers": "moisturizers",
+  "moisturises": "moisturizes",
+  "moisturising": "moisturizing",
+  "monologue": "monolog",
+  "monologues": "monologs",
+  "monopolisation": "monopolization",
+  "monopolise": "monopolize",
+  "monopolised": "monopolized",
+  "monopolises": "monopolizes",
+  "monopolising": "monopolizing",
+  "moralise": "moralize",
+  "moralised": "moralized",
+  "moralises": "moralizes",
+  "moralising": "moralizing",
+  "motorised": "motorized",
+  "mould": "mold",
+  "moulded": "molded",
+  "moulder": "molder",
+  "mouldered": "moldered",
+  "mouldering": "moldering",
+  "moulders": "molders",
+  "mouldier": "moldier",
+  "mouldiest": "moldiest",
+  "moulding": "molding",
+  "mouldings": "moldings",
+  "moulds": "molds",
+  "mouldy": "moldy",
+  "moult": "molt",
+  "moulted": "molted",
+  "moulting": "molting",
+  "moults": "molts",
+  "moustache": "mustache",
+  "moustached": "mustached",
+  "moustaches": "mustaches",
+  "moustachioed": "mustachioed",
+  "multicoloured": "multicolored",
+  "nationalisation": "nationalization",
+  "nationalisations": "nationalizations",
+  "nationalise": "nationalize",
+  "nationalised": "nationalized",
+  "nationalises": "nationalizes",
+  "nationalising": "nationalizing",
+  "naturalisation": "naturalization",
+  "naturalise": "naturalize",
+  "naturalised": "naturalized",
+  "naturalises": "naturalizes",
+  "naturalising": "naturalizing",
+  "neighbour": "neighbor",
+  "neighbourhood": "neighborhood",
+  "neighbourhoods": "neighborhoods",
+  "neighbouring": "neighboring",
+  "neighbourliness": "neighborliness",
+  "neighbourly": "neighborly",
+  "neighbours": "neighbors",
+  "neutralisation": "neutralization",
+  "neutralise": "neutralize",
+  "neutralised": "neutralized",
+  "neutralises": "neutralizes",
+  "neutralising": "neutralizing",
+  "normalisation": "normalization",
+  "normalise": "normalize",
+  "normalised": "normalized",
+  "normalises": "normalizes",
+  "normalising": "normalizing",
+  "odour": "odor",
+  "odourless": "odorless",
+  "odours": "odors",
+  "oesophagus": "esophagus",
+  "oesophaguses": "esophaguses",
+  "oestrogen": "estrogen",
+  "offence": "offense",
+  "offences": "offenses",
+  "omelette": "omelet",
+  "omelettes": "omelets",
+  "optimise": "optimize",
+  "optimised": "optimized",
+  "optimises": "optimizes",
+  "optimising": "optimizing",
+  "organisation": "organization",
+  "organisational": "organizational",
+  "organisations": "organizations",
+  "organise": "organize",
+  "organised": "organized",
+  "organiser": "organizer",
+  "organisers": "organizers",
+  "organises": "organizes",
+  "organising": "organizing",
+  "orthopaedic": "orthopedic",
+  "orthopaedics": "orthopedics",
+  "ostracise": "ostracize",
+  "ostracised": "ostracized",
+  "ostracises": "ostracizes",
+  "ostracising": "ostracizing",
+  "outmanoeuvre": "outmaneuver",
+  "outmanoeuvred": "outmaneuvered",
+  "outmanoeuvres": "outmaneuvers",
+  "outmanoeuvring": "outmaneuvering",
+  "overemphasise": "overemphasize",
+  "overemphasised": "overemphasized",
+  "overemphasises": "overemphasizes",
+  "overemphasising": "overemphasizing",
+  "oxidisation": "oxidization",
+  "oxidise": "oxidize",
+  "oxidised": "oxidized",
+  "oxidises": "oxidizes",
+  "oxidising": "oxidizing",
+  "paederast": "pederast",
+  "paederasts": "pederasts",
+  "paediatric": "pediatric",
+  "paediatrician": "pediatrician",
+  "paediatricians": "pediatricians",
+  "paediatrics": "pediatrics",
+  "paedophile": "pedophile",
+  "paedophiles": "pedophiles",
+  "paedophilia": "pedophilia",
+  "palaeolithic": "paleolithic",
+  "palaeontologist": "paleontologist",
+  "palaeontologists": "paleontologists",
+  "palaeontology": "paleontology",
+  "panelled": "paneled",
+  "panelling": "paneling",
+  "panellist": "panelist",
+  "panellists": "panelists",
+  "paralyse": "paralyze",
+  "paralysed": "paralyzed",
+  "paralyses": "paralyzes",
+  "paralysing": "paralyzing",
+  "parcelled": "parceled",
+  "parcelling": "parceling",
+  "parlour": "parlor",
+  "parlours": "parlors",
+  "particularise": "particularize",
+  "particularised": "particularized",
+  "particularises": "particularizes",
+  "particularising": "particularizing",
+  "passivisation": "passivization",
+  "passivise": "passivize",
+  "passivised": "passivized",
+  "passivises": "passivizes",
+  "passivising": "passivizing",
+  "pasteurisation": "pasteurization",
+  "pasteurise": "pasteurize",
+  "pasteurised": "pasteurized",
+  "pasteurises": "pasteurizes",
+  "pasteurising": "pasteurizing",
+  "patronise": "patronize",
+  "patronised": "patronized",
+  "patronises": "patronizes",
+  "patronising": "patronizing",
+  "patronisingly": "patronizingly",
+  "pedalled": "pedaled",
+  "pedalling": "pedaling",
+  "pedestrianisation": "pedestrianization",
+  "pedestrianise": "pedestrianize",
+  "pedestrianised": "pedestrianized",
+  "pedestrianises": "pedestrianizes",
+  "pedestrianising": "pedestrianizing",
+  "penalise": "penalize",
+  "penalised": "penalized",
+  "penalises": "penalizes",
+  "penalising": "penalizing",
+  "pencilled": "penciled",
+  "pencilling": "penciling",
+  "personalise": "personalize",
+  "personalised": "personalized",
+  "personalises": "personalizes",
+  "personalising": "personalizing",
+  "pharmacopoeia": "pharmacopeia",
+  "pharmacopoeias": "pharmacopeias",
+  "philosophise": "philosophize",
+  "philosophised": "philosophized",
+  "philosophises": "philosophizes",
+  "philosophising": "philosophizing",
+  "philtre": "filter",
+  "philtres": "filters",
+  "phoney": "phony",
+  "plagiarise": "plagiarize",
+  "plagiarised": "plagiarized",
+  "plagiarises": "plagiarizes",
+  "plagiarising": "plagiarizing",
+  "plough": "plow",
+  "ploughed": "plowed",
+  "ploughing": "plowing",
+  "ploughman": "plowman",
+  "ploughmen": "plowmen",
+  "ploughs": "plows",
+  "ploughshare": "plowshare",
+  "ploughshares": "plowshares",
+  "polarisation": "polarization",
+  "polarise": "polarize",
+  "polarised": "polarized",
+  "polarises": "polarizes",
+  "polarising": "polarizing",
+  "politicisation": "politicization",
+  "politicise": "politicize",
+  "politicised": "politicized",
+  "politicises": "politicizes",
+  "politicising": "politicizing",
+  "popularisation": "popularization",
+  "popularise": "popularize",
+  "popularised": "popularized",
+  "popularises": "popularizes",
+  "popularising": "popularizing",
+  "pouffe": "pouf",
+  "pouffes": "poufs",
+  "practise": "practice",
+  "practised": "practiced",
+  "practises": "practices",
+  "practising": "practicing",
+  "praesidium": "presidium",
+  "praesidiums": "presidiums",
+  "pressurisation": "pressurization",
+  "pressurise": "pressurize",
+  "pressurised": "pressurized",
+  "pressurises": "pressurizes",
+  "pressurising": "pressurizing",
+  "pretence": "pretense",
+  "pretences": "pretenses",
+  "primaeval": "primeval",
+  "prioritisation": "prioritization",
+  "prioritise": "prioritize",
+  "prioritised": "prioritized",
+  "prioritises": "prioritizes",
+  "prioritising": "prioritizing",
+  "privatisation": "privatization",
+  "privatisations": "privatizations",
+  "privatise": "privatize",
+  "privatised": "privatized",
+  "privatises": "privatizes",
+  "privatising": "privatizing",
+  "professionalisation": "professionalization",
+  "professionalise": "professionalize",
+  "professionalised": "professionalized",
+  "professionalises": "professionalizes",
+  "professionalising": "professionalizing",
+  "programme": "program",
+  "programmes": "programs",
+  "prologue": "prolog",
+  "prologues": "prologs",
+  "propagandise": "propagandize",
+  "propagandised": "propagandized",
+  "propagandises": "propagandizes",
+  "propagandising": "propagandizing",
+  "proselytise": "proselytize",
+  "proselytised": "proselytized",
+  "proselytiser": "proselytizer",
+  "proselytisers": "proselytizers",
+  "proselytises": "proselytizes",
+  "proselytising": "proselytizing",
+  "psychoanalyse": "psychoanalyze",
+  "psychoanalysed": "psychoanalyzed",
+  "psychoanalyses": "psychoanalyzes",
+  "psychoanalysing": "psychoanalyzing",
+  "publicise": "publicize",
+  "publicised": "publicized",
+  "publicises": "publicizes",
+  "publicising": "publicizing",
+  "pulverisation": "pulverization",
+  "pulverise": "pulverize",
+  "pulverised": "pulverized",
+  "pulverises": "pulverizes",
+  "pulverising": "pulverizing",
+  "pummelled": "pummel",
+  "pummelling": "pummeled",
+  "pyjama": "pajama",
+  "pyjamas": "pajamas",
+  "pzazz": "pizzazz",
+  "quarrelled": "quarreled",
+  "quarrelling": "quarreling",
+  "radicalise": "radicalize",
+  "radicalised": "radicalized",
+  "radicalises": "radicalizes",
+  "radicalising": "radicalizing",
+  "rancour": "rancor",
+  "randomise": "randomize",
+  "randomised": "randomized",
+  "randomises": "randomizes",
+  "randomising": "randomizing",
+  "rationalisation": "rationalization",
+  "rationalisations": "rationalizations",
+  "rationalise": "rationalize",
+  "rationalised": "rationalized",
+  "rationalises": "rationalizes",
+  "rationalising": "rationalizing",
+  "ravelled": "raveled",
+  "ravelling": "raveling",
+  "realisable": "realizable",
+  "realisation": "realization",
+  "realisations": "realizations",
+  "realise": "realize",
+  "realised": "realized",
+  "realises": "realizes",
+  "realising": "realizing",
+  "recognisable": "recognizable",
+  "recognisably": "recognizably",
+  "recognisance": "recognizance",
+  "recognise": "recognize",
+  "recognised": "recognized",
+  "recognises": "recognizes",
+  "recognising": "recognizing",
+  "reconnoitre": "reconnoiter",
+  "reconnoitred": "reconnoitered",
+  "reconnoitres": "reconnoiters",
+  "reconnoitring": "reconnoitering",
+  "refuelled": "refueled",
+  "refuelling": "refueling",
+  "regularisation": "regularization",
+  "regularise": "regularize",
+  "regularised": "regularized",
+  "regularises": "regularizes",
+  "regularising": "regularizing",
+  "remodelled": "remodeled",
+  "remodelling": "remodeling",
+  "remould": "remold",
+  "remoulded": "remolded",
+  "remoulding": "remolding",
+  "remoulds": "remolds",
+  "reorganisation": "reorganization",
+  "reorganisations": "reorganizations",
+  "reorganise": "reorganize",
+  "reorganised": "reorganized",
+  "reorganises": "reorganizes",
+  "reorganising": "reorganizing",
+  "revelled": "reveled",
+  "reveller": "reveler",
+  "revellers": "revelers",
+  "revelling": "reveling",
+  "revitalise": "revitalize",
+  "revitalised": "revitalized",
+  "revitalises": "revitalizes",
+  "revitalising": "revitalizing",
+  "revolutionise": "revolutionize",
+  "revolutionised": "revolutionized",
+  "revolutionises": "revolutionizes",
+  "revolutionising": "revolutionizing",
+  "rhapsodise": "rhapsodize",
+  "rhapsodised": "rhapsodized",
+  "rhapsodises": "rhapsodizes",
+  "rhapsodising": "rhapsodizing",
+  "rigour": "rigor",
+  "rigours": "rigors",
+  "ritualised": "ritualized",
+  "rivalled": "rivaled",
+  "rivalling": "rivaling",
+  "romanticise": "romanticize",
+  "romanticised": "romanticized",
+  "romanticises": "romanticizes",
+  "romanticising": "romanticizing",
+  "rumour": "rumor",
+  "rumoured": "rumored",
+  "rumours": "rumors",
+  "sabre": "saber",
+  "sabres": "sabers",
+  "saltpetre": "saltpeter",
+  "sanitise": "sanitize",
+  "sanitised": "sanitized",
+  "sanitises": "sanitizes",
+  "sanitising": "sanitizing",
+  "satirise": "satirize",
+  "satirised": "satirized",
+  "satirises": "satirizes",
+  "satirising": "satirizing",
+  "saviour": "savior",
+  "saviours": "saviors",
+  "savour": "savor",
+  "savoured": "savored",
+  "savouries": "savories",
+  "savouring": "savoring",
+  "savours": "savors",
+  "savoury": "savory",
+  "scandalise": "scandalize",
+  "scandalised": "scandalized",
+  "scandalises": "scandalizes",
+  "scandalising": "scandalizing",
+  "sceptic": "skeptic",
+  "sceptical": "skeptical",
+  "sceptically": "skeptically",
+  "scepticism": "skepticism",
+  "sceptics": "skeptics",
+  "sceptre": "scepter",
+  "sceptres": "scepters",
+  "scrutinise": "scrutinize",
+  "scrutinised": "scrutinized",
+  "scrutinises": "scrutinizes",
+  "scrutinising": "scrutinizing",
+  "secularisation": "secularization",
+  "secularise": "secularize",
+  "secularised": "secularized",
+  "secularises": "secularizes",
+  "secularising": "secularizing",
+  "sensationalise": "sensationalize",
+  "sensationalised": "sensationalized",
+  "sensationalises": "sensationalizes",
+  "sensationalising": "sensationalizing",
+  "sensitise": "sensitize",
+  "sensitised": "sensitized",
+  "sensitises": "sensitizes",
+  "sensitising": "sensitizing",
+  "sentimentalise": "sentimentalize",
+  "sentimentalised": "sentimentalized",
+  "sentimentalises": "sentimentalizes",
+  "sentimentalising": "sentimentalizing",
+  "sepulchre": "sepulcher",
+  "sepulchres": "sepulchers",
+  "serialisation": "serialization",
+  "serialisations": "serializations",
+  "serialise": "serialize",
+  "serialised": "serialized",
+  "serialises": "serializes",
+  "serialising": "serializing",
+  "sermonise": "sermonize",
+  "sermonised": "sermonized",
+  "sermonises": "sermonizes",
+  "sermonising": "sermonizing",
+  "sheikh": "sheik",
+  "shovelled": "shoveled",
+  "shovelling": "shoveling",
+  "shrivelled": "shriveled",
+  "shrivelling": "shriveling",
+  "signalise": "signalize",
+  "signalised": "signalized",
+  "signalises": "signalizes",
+  "signalising": "signalizing",
+  "signalled": "signaled",
+  "signalling": "signaling",
+  "smoulder": "smolder",
+  "smouldered": "smoldered",
+  "smouldering": "smoldering",
+  "smoulders": "smolders",
+  "snivelled": "sniveled",
+  "snivelling": "sniveling",
+  "snorkelled": "snorkeled",
+  "snorkelling": "snorkeling",
+  "snowplough": "snowplow",
+  "snowploughs": "snowplow",
+  "socialisation": "socialization",
+  "socialise": "socialize",
+  "socialised": "socialized",
+  "socialises": "socializes",
+  "socialising": "socializing",
+  "sodomise": "sodomize",
+  "sodomised": "sodomized",
+  "sodomises": "sodomizes",
+  "sodomising": "sodomizing",
+  "solemnise": "solemnize",
+  "solemnised": "solemnized",
+  "solemnises": "solemnizes",
+  "solemnising": "solemnizing",
+  "sombre": "somber",
+  "specialisation": "specialization",
+  "specialisations": "specializations",
+  "specialise": "specialize",
+  "specialised": "specialized",
+  "specialises": "specializes",
+  "specialising": "specializing",
+  "spectre": "specter",
+  "spectres": "specters",
+  "spiralled": "spiraled",
+  "spiralling": "spiraling",
+  "splendour": "splendor",
+  "splendours": "splendors",
+  "squirrelled": "squirreled",
+  "squirrelling": "squirreling",
+  "stabilisation": "stabilization",
+  "stabilise": "stabilize",
+  "stabilised": "stabilized",
+  "stabiliser": "stabilizer",
+  "stabilisers": "stabilizers",
+  "stabilises": "stabilizes",
+  "stabilising": "stabilizing",
+  "standardisation": "standardization",
+  "standardise": "standardize",
+  "standardised": "standardized",
+  "standardises": "standardizes",
+  "standardising": "standardizing",
+  "stencilled": "stenciled",
+  "stencilling": "stenciling",
+  "sterilisation": "sterilization",
+  "sterilisations": "sterilizations",
+  "sterilise": "sterilize",
+  "sterilised": "sterilized",
+  "steriliser": "sterilizer",
+  "sterilisers": "sterilizers",
+  "sterilises": "sterilizes",
+  "sterilising": "sterilizing",
+  "stigmatisation": "stigmatization",
+  "stigmatise": "stigmatize",
+  "stigmatised": "stigmatized",
+  "stigmatises": "stigmatizes",
+  "stigmatising": "stigmatizing",
+  "storey": "story",
+  "storeys": "stories",
+  "subsidisation": "subsidization",
+  "subsidise": "subsidize",
+  "subsidised": "subsidized",
+  "subsidiser": "subsidizer",
+  "subsidisers": "subsidizers",
+  "subsidises": "subsidizes",
+  "subsidising": "subsidizing",
+  "succour": "succor",
+  "succoured": "succored",
+  "succouring": "succoring",
+  "succours": "succors",
+  "sulphate": "sulfate",
+  "sulphates": "sulfates",
+  "sulphide": "sulfide",
+  "sulphides": "sulfides",
+  "sulphur": "sulfur",
+  "sulphurous": "sulfurous",
+  "summarise": "summarize",
+  "summarised": "summarized",
+  "summarises": "summarizes",
+  "summarising": "summarizing",
+  "swivelled": "swiveled",
+  "swivelling": "swiveling",
+  "symbolise": "symbolize",
+  "symbolised": "symbolized",
+  "symbolises": "symbolizes",
+  "symbolising": "symbolizing",
+  "sympathise": "sympathize",
+  "sympathised": "sympathized",
+  "sympathiser": "sympathizer",
+  "sympathisers": "sympathizers",
+  "sympathises": "sympathizes",
+  "sympathising": "sympathizing",
+  "synchronisation": "synchronization",
+  "synchronise": "synchronize",
+  "synchronised": "synchronized",
+  "synchronises": "synchronizes",
+  "synchronising": "synchronizing",
+  "synthesise": "synthesize",
+  "synthesised": "synthesized",
+  "synthesiser": "synthesizer",
+  "synthesisers": "synthesizers",
+  "synthesises": "synthesizes",
+  "synthesising": "synthesizing",
+  "syphon": "siphon",
+  "syphoned": "siphoned",
+  "syphoning": "siphoning",
+  "syphons": "siphons",
+  "systematisation": "systematization",
+  "systematise": "systematize",
+  "systematised": "systematized",
+  "systematises": "systematizes",
+  "systematising": "systematizing",
+  "tantalise": "tantalize",
+  "tantalised": "tantalized",
+  "tantalises": "tantalizes",
+  "tantalising": "tantalizing",
+  "tantalisingly": "tantalizingly",
+  "tasselled": "tasseled",
+  "technicolour": "technicolor",
+  "temporise": "temporize",
+  "temporised": "temporized",
+  "temporises": "temporizes",
+  "temporising": "temporizing",
+  "tenderise": "tenderize",
+  "tenderised": "tenderized",
+  "tenderises": "tenderizes",
+  "tenderising": "tenderizing",
+  "terrorise": "terrorize",
+  "terrorised": "terrorized",
+  "terrorises": "terrorizes",
+  "terrorising": "terrorizing",
+  "theatre": "theater",
+  "theatregoer": "theatergoer",
+  "theatregoers": "theatergoers",
+  "theatres": "theaters",
+  "theorise": "theorize",
+  "theorised": "theorized",
+  "theorises": "theorizes",
+  "theorising": "theorizing",
+  "tonne": "ton",
+  "tonnes": "tons",
+  "towelled": "toweled",
+  "towelling": "toweling",
+  "toxaemia": "toxemia",
+  "tranquillise": "tranquilize",
+  "tranquillised": "tranquilized",
+  "tranquilliser": "tranquilizer",
+  "tranquillisers": "tranquilizers",
+  "tranquillises": "tranquilizes",
+  "tranquillising": "tranquilizing",
+  "tranquillity": "tranquility",
+  "tranquillize": "tranquilize",
+  "tranquillized": "tranquilized",
+  "tranquillizer": "tranquilizer",
+  "tranquillizers": "tranquilizers",
+  "tranquillizes": "tranquilizes",
+  "tranquillizing": "tranquilizing",
+  "tranquilly": "tranquility",
+  "transistorised": "transistorized",
+  "traumatise": "traumatize",
+  "traumatised": "traumatized",
+  "traumatises": "traumatizes",
+  "traumatising": "traumatizing",
+  "travelled": "traveled",
+  "traveller": "traveler",
+  "travellers": "travelers",
+  "travelling": "traveling",
+  "travelog": "travelogue",
+  "travelogs": "travelogues",
+  "trialled": "trialed",
+  "trialling": "trialing",
+  "tricolour": "tricolor",
+  "tricolours": "tricolors",
+  "trivialise": "trivialize",
+  "trivialised": "trivialized",
+  "trivialises": "trivializes",
+  "trivialising": "trivializing",
+  "tumour": "tumor",
+  "tumours": "tumors",
+  "tunnelled": "tunneled",
+  "tunnelling": "tunneling",
+  "tyrannise": "tyrannize",
+  "tyrannised": "tyrannized",
+  "tyrannises": "tyrannizes",
+  "tyrannising": "tyrannizing",
+  "tyre": "tire",
+  "tyres": "tires",
+  "unauthorised": "unauthorized",
+  "uncivilised": "uncivilized",
+  "underutilised": "underutilized",
+  "unequalled": "unequaled",
+  "unfavourable": "unfavorable",
+  "unfavourably": "unfavorably",
+  "unionisation": "unionization",
+  "unionise": "unionize",
+  "unionised": "unionized",
+  "unionises": "unionizes",
+  "unionising": "unionizing",
+  "unorganised": "unorganized",
+  "unravelled": "unraveled",
+  "unravelling": "unraveling",
+  "unrecognisable": "unrecognizable",
+  "unrecognised": "unrecognized",
+  "unrivalled": "unrivaled",
+  "unsavoury": "unsavory",
+  "untrammelled": "untrammeled",
+  "urbanisation": "urbanization",
+  "urbanise": "urbanize",
+  "urbanised": "urbanized",
+  "urbanises": "urbanizes",
+  "urbanising": "urbanizing",
+  "utilisable": "utilizable",
+  "utilisation": "utilization",
+  "utilise": "utilize",
+  "utilised": "utilized",
+  "utilises": "utilizes",
+  "utilising": "utilizing",
+  "valour": "valor",
+  "vandalise": "vandalize",
+  "vandalised": "vandalized",
+  "vandalises": "vandalizes",
+  "vandalising": "vandalizing",
+  "vaporisation": "vaporization",
+  "vaporise": "vaporize",
+  "vaporised": "vaporized",
+  "vaporises": "vaporizes",
+  "vaporising": "vaporizing",
+  "vapour": "vapor",
+  "vapours": "vapors",
+  "verbalise": "verbalize",
+  "verbalised": "verbalized",
+  "verbalises": "verbalizes",
+  "verbalising": "verbalizing",
+  "victimisation": "victimization",
+  "victimise": "victimize",
+  "victimised": "victimized",
+  "victimises": "victimizes",
+  "victimising": "victimizing",
+  "videodisc": "videodisk",
+  "videodiscs": "videodisks",
+  "vigour": "vigor",
+  "visualisation": "visualization",
+  "visualisations": "visualizations",
+  "visualise": "visualize",
+  "visualised": "visualized",
+  "visualises": "visualizes",
+  "visualising": "visualizing",
+  "vocalisation": "vocalization",
+  "vocalisations": "vocalizations",
+  "vocalise": "vocalize",
+  "vocalised": "vocalized",
+  "vocalises": "vocalizes",
+  "vocalising": "vocalizing",
+  "vulcanised": "vulcanized",
+  "vulgarisation": "vulgarization",
+  "vulgarise": "vulgarize",
+  "vulgarised": "vulgarized",
+  "vulgarises": "vulgarizes",
+  "vulgarising": "vulgarizing",
+  "waggon": "wagon",
+  "waggons": "wagons",
+  "watercolour": "watercolor",
+  "watercolours": "watercolors",
+  "weaselled": "weaseled",
+  "weaselling": "weaseling",
+  "westernisation": "westernization",
+  "westernise": "westernize",
+  "westernised": "westernized",
+  "westernises": "westernizes",
+  "westernising": "westernizing",
+  "womanise": "womanize",
+  "womanised": "womanized",
+  "womaniser": "womanizer",
+  "womanisers": "womanizers",
+  "womanises": "womanizes",
+  "womanising": "womanizing",
+  "woollen": "woolen",
+  "woollens": "woolens",
+  "woollies": "woolies",
+  "woolly": "wooly",
+  "worshipped": "worshiped",
+  "worshipper": "worshiper",
+  "worshipping": "worshiping",
+  "yodelled": "yodeled",
+  "yodelling": "yodeling",
+  "yoghourt": "yogurt",
+  "yoghourts": "yogurts",
+  "yoghurt": "yogurt",
+  "yoghurts": "yogurts"
+}

whisper-small-flaxtest/preprocessor_config.json

@@ -0,0 +1,14 @@
+{
+  "chunk_length": 30,
+  "feature_extractor_type": "WhisperFeatureExtractor",
+  "feature_size": 80,
+  "hop_length": 160,
+  "n_fft": 400,
+  "n_samples": 480000,
+  "nb_max_frames": 3000,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "processor_class": "WhisperProcessor",
+  "return_attention_mask": false,
+  "sampling_rate": 16000
+}

whisper-small-flaxtest/special_tokens_map.json

@@ -0,0 +1,114 @@
+{
+  "additional_special_tokens": [
+    "<|startoftranscript|>",
+    "<|en|>",
+    "<|zh|>",
+    "<|de|>",
+    "<|es|>",
+    "<|ru|>",
+    "<|ko|>",
+    "<|fr|>",
+    "<|ja|>",
+    "<|pt|>",
+    "<|tr|>",
+    "<|pl|>",
+    "<|ca|>",
+    "<|nl|>",
+    "<|ar|>",
+    "<|sv|>",
+    "<|it|>",
+    "<|id|>",
+    "<|hi|>",
+    "<|fi|>",
+    "<|vi|>",
+    "<|iw|>",
+    "<|uk|>",
+    "<|el|>",
+    "<|ms|>",
+    "<|cs|>",
+    "<|ro|>",
+    "<|da|>",
+    "<|hu|>",
+    "<|ta|>",
+    "<|no|>",
+    "<|th|>",
+    "<|ur|>",
+    "<|hr|>",
+    "<|bg|>",
+    "<|lt|>",
+    "<|la|>",
+    "<|mi|>",
+    "<|ml|>",
+    "<|cy|>",
+    "<|sk|>",
+    "<|te|>",
+    "<|fa|>",
+    "<|lv|>",
+    "<|bn|>",
+    "<|sr|>",
+    "<|az|>",
+    "<|sl|>",
+    "<|kn|>",
+    "<|et|>",
+    "<|mk|>",
+    "<|br|>",
+    "<|eu|>",
+    "<|is|>",
+    "<|hy|>",
+    "<|ne|>",
+    "<|mn|>",
+    "<|bs|>",
+    "<|kk|>",
+    "<|sq|>",
+    "<|sw|>",
+    "<|gl|>",
+    "<|mr|>",
+    "<|pa|>",
+    "<|si|>",
+    "<|km|>",
+    "<|sn|>",
+    "<|yo|>",
+    "<|so|>",
+    "<|af|>",
+    "<|oc|>",
+    "<|ka|>",
+    "<|be|>",
+    "<|tg|>",
+    "<|sd|>",
+    "<|gu|>",
+    "<|am|>",
+    "<|yi|>",
+    "<|lo|>",
+    "<|uz|>",
+    "<|fo|>",
+    "<|ht|>",
+    "<|ps|>",
+    "<|tk|>",
+    "<|nn|>",
+    "<|mt|>",
+    "<|sa|>",
+    "<|lb|>",
+    "<|my|>",
+    "<|bo|>",
+    "<|tl|>",
+    "<|mg|>",
+    "<|as|>",
+    "<|tt|>",
+    "<|haw|>",
+    "<|ln|>",
+    "<|ha|>",
+    "<|ba|>",
+    "<|jw|>",
+    "<|su|>",
+    "<|translate|>",
+    "<|transcribe|>",
+    "<|startoflm|>",
+    "<|startofprev|>",
+    "<|nocaptions|>",
+    "<|notimestamps|>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

whisper-small-flaxtest/tokenizer_config.json

@@ -0,0 +1,35 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "errors": "replace",
+  "model_max_length": 1024,
+  "pad_token": null,
+  "processor_class": "WhisperProcessor",
+  "return_attention_mask": false,
+  "special_tokens_map_file": null,
+  "tokenizer_class": "WhisperTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

whisper-tiny-ft-dummy/events.out.tfevents.1677582573.t1v-n-d163ce9a-w-0.1276805.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28df6cfd95c9539437e6d7457e72914758a48bc8c81c04d68d45c9beec6b5bdc
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677585137.t1v-n-d163ce9a-w-0.1284051.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0facb62340abd1a1ec68e3ddb84c1e26464d970168d3c5b8c3c52ec5d52de1d6
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677587059.t1v-n-d163ce9a-w-0.1287692.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e14a45dd8b6f3c888c4779ee8bcd6b20049ea8d45e1b0f6c34b7455dddfff433
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677587350.t1v-n-d163ce9a-w-0.1292303.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6378de6a369cc141269aaff4436e172c11d71a3b1dcc2d11cdae62da20d488c7
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677588068.t1v-n-d163ce9a-w-0.1297330.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92bbb8a3a7af5963003ed46c79ce553a9f98f9897c228c98e662440c90f1c8c5
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677588142.t1v-n-d163ce9a-w-0.1301760.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27b4453568339824f34c48944b0bf00ca2a4e603301b5d10dbafa527b557b776
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677588581.t1v-n-d163ce9a-w-0.1306471.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04d88ab6f9b46f027b704a597a8a2136d6260304a365e1c303e9269e2c9e5b2d
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677590425.t1v-n-d163ce9a-w-0.1318486.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3e2d7e93cfc0d7e463bbe781f6a502d64db4721d515e2ca070dd1673ecf39fc
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677590675.t1v-n-d163ce9a-w-0.1323104.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc61c0bcaa754613477811c72780874609f1047b22afc279f3ec8458d68cd1eb
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677591418.t1v-n-d163ce9a-w-0.1328351.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c6b869dc6ce4f75deb0966514f50eb606d7e808bcd6d51a9bd40f4b1f5902da
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677591678.t1v-n-d163ce9a-w-0.1333009.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0f75dd974d807c09a331f5b745d02ccf57c483cd88e094f12633f5a5acf2fa4
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677591869.t1v-n-d163ce9a-w-0.1337579.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76910deccce21a87d5b2b12eb9b71a6ffbf998d4dfb4f19f56f71e31b8067404
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677592021.t1v-n-d163ce9a-w-0.1344023.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8414ba6c82d77bad684b451c11194d51834d0917c9fd2ecf72b26e198d7b5408
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677592180.t1v-n-d163ce9a-w-0.1350466.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83b7ebb7a355c38ce543ddace67c17fcda3b39efad8637520028b7548de964ba
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677592712.t1v-n-d163ce9a-w-0.1355445.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7675eefbe2395c77c95b12e778c549db892e3a6849aac7022bb535c19fda1f17
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677596196.t1v-n-d163ce9a-w-0.1363328.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5280463efd0ffadd26520dc5af0cf8bc7272ce109d9c7cb5d130cbffaae90d76
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677596396.t1v-n-d163ce9a-w-0.1367877.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03b9ae6d9966b67a2ecb933764153693d2c2e8815ccd08203af297694d498819
+size 40

whisper-tiny-ft-dummy/events.out.tfevents.1677596532.t1v-n-d163ce9a-w-0.1372356.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:851cba504d6263302d36e66de1f9ce7521b2bf5df056c0f22e320b2e9247c2a2
+size 40