mzjvp6ho: saving weights and logs of step 5k

events.out.tfevents.1660117225.t1v-n-eedfb410-w-0.55204.0.v2

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

events.out.tfevents.1660130897.t1v-n-eedfb410-w-0.8420.0.v2

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

events.out.tfevents.1660143983.t1v-n-eedfb410-w-0.3332902.0.v2

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

events.out.tfevents.1660145355.t1v-n-eedfb410-w-0.2349240.0.v2

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

events.out.tfevents.1660206880.t1v-n-eedfb410-w-0.1479163.0.v2

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

events.out.tfevents.1660208728.t1v-n-eedfb410-w-0.503538.0.v2

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

events.out.tfevents.1660218137.t1v-n-eedfb410-w-0.2916397.0.v2

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

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d6e7fe76ddde6be27c0129735dc3bb50191fff23d8f87075f1335970abf06211
+oid sha256:2824056b9c2f157ff862c17877b5aa4a77f0f6107345973495c02df3828b7469
 size 3850218852

run.recover.sh

@@ -11,9 +11,8 @@ WANDB_ENTITY=NbAiLab WANDB_PROJECT=wav2vec2 python run_flax_speech_recognition_c
         --per_device_train_batch_size="2" \
         --per_device_eval_batch_size="2" \
         --gradient_accumulation_steps="1" \
-        --precision="full_mixed" \
+        --precision="half_mixed" \
         --matmul_precision="bfloat16" \
-        --multisteps \
         --learning_rate="6.394633237505332e-05" \
         --skip_steps="275000" \
         --warmup_steps="2000" \

run.sh

@@ -1,3 +1,6 @@
+# See https://github.com/sanchit-gandhi/seq2seq-speech/issues/23#issuecomment-1122183173: do_lower_case should only be set to True for the tokenizer if the tokenizer has upper case letters in the vocab
+# Let's also not add extra remove_punctuation
+# And limit max duration to 25 seconds
 WANDB_ENTITY=NbAiLab WANDB_PROJECT=wav2vec2 python run_flax_speech_recognition_ctc.py \
         --model_name_or_path="facebook/wav2vec2-xls-r-1b" \
         --hub_model_id="NbAiLab/wav2vec2-1b-npsc-nst-tpu" \
@@ -11,7 +14,7 @@ WANDB_ENTITY=NbAiLab WANDB_PROJECT=wav2vec2 python run_flax_speech_recognition_c
         --precision="full_mixed" \
         --matmul_precision="bfloat16" \
         --multisteps \
-        --learning_rate="1e-4" \
+        --learning_rate="2e-5" \
         --warmup_steps="2000" \
         --length_column_name="input_length" \
         --evaluation_strategy="steps" \
@@ -32,7 +35,7 @@ WANDB_ENTITY=NbAiLab WANDB_PROJECT=wav2vec2 python run_flax_speech_recognition_c
         --mask_feature_length="64" \
         --gradient_checkpointing \
         --min_duration_in_seconds="0.5" \
-        --max_duration_in_seconds="30.0" \
+        --max_duration_in_seconds="25.0" \
         --use_auth_token \
         --seed="42" \
         --group_by_length \
@@ -40,10 +43,5 @@ WANDB_ENTITY=NbAiLab WANDB_PROJECT=wav2vec2 python run_flax_speech_recognition_c
         --push_to_hub \
         --preprocessing_num_workers="32" \
         --ctc_zero_infinity \
-        --do_lower_case \
         --wandb_project="wav2vec2" \
         --wandb_name="wav2vec2-1b-npsc-nst-tpu" \
-        --remove_punctuation
-
-
-# --fp16

run_flax_speech_recognition_ctc.py

@@ -415,12 +415,12 @@ class MixedPrecisionTrainState(struct.PyTreeNode):
         )
 
     @classmethod
-    def create(cls, *, apply_fn, params, tx, to_dtype, **kwargs):
+    def create(cls, *, apply_fn, params, tx, to_dtype, step=0, **kwargs):
         """Creates a new instance with `step=0` and initialized `opt_state`."""
         # downcast optimizer state to bf16 if mixed-precision training
         opt_state = tx.init(to_dtype(params)) if tx is not None else None
         return cls(
-            step=0,
+            step=step,
             apply_fn=apply_fn,
             params=params,
             tx=tx,
@@ -1339,6 +1339,7 @@ def main():
 
     # Setup train state
     state = MixedPrecisionTrainState.create(
+        step=data_args.skip_steps,
         apply_fn=model.__call__,
         get_attention_mask_fn=model._get_feature_vector_attention_mask,
         params=model.params,
@@ -1517,14 +1518,13 @@ def main():
         if training_args.do_train:
             # ======================== Training ================================
             train_start = time.time()
+            # Create sampling rng
+            rng, input_rng = jax.random.split(rng)
 
             if epoch < skip_epochs:
                 logger.info(f"Skipping epoch {epoch + 1}")
                 continue
 
-            # Create sampling rng
-            rng, input_rng = jax.random.split(rng)
-
             # Generate an epoch by randomly shuffling sampling indices from the train dataset and grouping by length
             train_samples_idx = get_grouped_indices(vectorized_datasets[data_args.train_split_name], batch_size_per_update, input_rng)
             train_batch_idx = generate_batch_splits(train_samples_idx, batch_size_per_update)

special_tokens_map.json

@@ -399,6 +399,104 @@
       "rstrip": false,
       "single_word": false
     },
+    {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
     {
       "content": "</s>",
       "lstrip": false,

wandb/debug-internal.log

@@ -1 +1 @@
-run-20220805_230151-2y71vcu4/logs/debug-internal.log
+run-20220811_101752-mzjvp6ho/logs/debug-internal.log

wandb/debug.log

@@ -1 +1 @@
-run-20220805_230151-2y71vcu4/logs/debug.log
+run-20220811_101752-mzjvp6ho/logs/debug.log

wandb/latest-run

@@ -1 +1 @@
-run-20220805_230151-2y71vcu4
+run-20220811_101752-mzjvp6ho

wandb/run-20220810_073735-23avj35z/files/code/run_flax_speech_recognition_ctc.py

@@ -0,0 +1,1631 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace 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 connectionist temporal classification (CTC) 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 re
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import datasets
+import numpy as np
+from datasets import DatasetDict, load_dataset, load_metric
+from tqdm import tqdm
+
+import flax
+import jax
+import jax.numpy as jnp
+import optax
+import transformers
+import wandb as wandb
+from flax import core, jax_utils, struct, traverse_util
+from flax.jax_utils import unreplicate, pad_shard_unpad
+from flax.training.common_utils import get_metrics, shard, shard_prng_key
+from huggingface_hub import Repository
+from models import Wav2Vec2Config, FlaxWav2Vec2ForCTC
+from optax._src import linear_algebra
+from transformers import (
+    AutoFeatureExtractor,
+    AutoProcessor,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainingArguments,
+    is_tensorboard_available,
+    set_seed,
+)
+from transformers.file_utils import get_full_repo_name
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.17.0.dev0")
+
+require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/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)."
+        },
+    )
+    freeze_feature_encoder: bool = field(
+        default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}
+    )
+    attention_dropout: float = field(
+        default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}
+    )
+    activation_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The hidden activation dropout probability in the embeddings, encoder, and pooler."
+        },
+    )
+    hidden_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The dropout probability for all fully connected layers in the embeddings, encoder, and pooler."
+        },
+    )
+    feat_proj_dropout: float = field(
+        default=0.0,
+        metadata={
+            "help": "The feat proj dropout probability for feature encoder representations."
+        },
+    )
+    final_dropout: float = field(
+        default=0.0,
+        metadata={"help": "The dropout probability for the final projection layer."},
+    )
+    mask_time_prob: float = field(
+        default=0.1,
+        metadata={
+            "help": "The spec aug dropout probability for feature encoder representations."
+        },
+    )
+    mask_time_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the time axis."},
+    )
+    mask_feature_prob: float = field(
+        default=0.0,
+        metadata={
+            "help": "Probability of each feature vector along the feature axis to be chosen as the start of the vector"
+                    "span to be masked. Approximately ``mask_feature_prob * sequence_length // mask_feature_length`` feature bins will be masked along the time axis."
+        },
+    )
+    mask_feature_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the feature axis."},
+    )
+    layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})
+    ctc_loss_reduction: Optional[str] = field(
+        default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}
+    )
+    ctc_zero_infinity: Optional[bool] = field(
+        default=False, metadata={"help": "If True, will try yo aboud the CTC loss goinf to infinity."}
+    )
+
+
+@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."
+        },
+    )
+    max_test_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of test 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": "Truncate audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_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: Optional[int] = field(
+        default=512,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    min_label_length: Optional[int] = field(
+        default=2,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    pad_input_to_multiple_of: Optional[int] = field(
+        default=32000,
+        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."
+        },
+    )
+    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."
+        },
+    )
+    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 training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    do_lower_case: bool = field(
+        default=True,
+        metadata={"help": "Whether the target text should be lower cased."},
+    )
+    wandb_project: str = field(
+        default="flax-speech-recognition-ctc",
+        metadata={"help": "The name of the wandb project."},
+    )
+    wandb_name: str = field(
+        default=None,
+        metadata={"help": "The name of the wandb run."},
+    )
+    wandb_job_type: str = field(
+        default="CTC",
+        metadata={"help": "The name of the wandb job type."},
+    )
+    test_split_name: str = field(
+        default="test",
+        metadata={"help": "The name of the test data set split to use (via the datasets library). Defaults to 'test'"},
+    )
+    remove_punctuation: bool = field(
+        default=False, metadata={"help": "Whether or not to remove punctuation during training."}
+    )
+    skip_steps: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": "Skip this number of steps. Useful to continue training"
+        },
+    )
+
+
+# @flax.struct.dataclass
+@dataclass
+class FlaxTrainingArguments(TrainingArguments):
+    precision: str = field(
+        default="full",
+        metadata={
+            "help": "Whether to enable mixed-precision training. If true, the optimizer is stored in half-precision (bfloat16) and computations are executed in half-precision"
+            "**Note that this only specifies the dtype of the computation and optimizer state. It does not influence the dtype of model parameters.**"
+        },
+    )
+    matmul_precision: str = field(
+        default="default",
+        metadata={
+            "help": "Default floating-point precision of internal computations used in TPU matrix multiplications and convolutions. "
+            "This configuration option controls the default precision for JAX operations that take an optional precision argument (e.g. `lax.conv_general_dilated` and `lax.dot`). "
+            "This configuration option does not change the behaviours of such calls with explicit precision arguments; "
+            "it only changes the behaviors of calls with no such argument provided. "
+            "One of `['highest', 'float32', 'high', 'bfloat16_3x', 'default', 'bfloat16', 'fastest', None]`."
+        },
+    )
+    multisteps: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to use Optax MultiSteps for gradient accumulation. If `False` (default) and `gradient_accumulation_steps > 1`, "
+            "a custom gradient accumulation implementation will be employed."
+        },
+    )
+
+
+def to_fp32(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+
+
+def to_bf16(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t)
+
+
+class MixedPrecisionTrainState(struct.PyTreeNode):
+    """Train state for use with a single Optax optimizer.
+    Adapted from flax train_state https://github.com/google/flax/blob/main/flax/training/train_state.py
+
+    Synopsis::
+
+        state = TrainState.create(
+            apply_fn=model.apply,
+            params=variables['params'],
+            tx=tx)
+        grad_fn = jax.grad(make_loss_fn(state.apply_fn))
+        for batch in data:
+          grads = grad_fn(state.params, batch)
+          state = state.apply_gradients(grads=grads)
+
+    Args:
+      step: Counter starts at 0 and is incremented by every call to
+        `.apply_gradients()`.
+      apply_fn: Usually set to `model.apply()`. Kept in this dataclass for
+        convenience to have a shorter params list for the `train_step()` function
+        in your training loop.
+      params: The parameters to be updated by `tx` and used by `apply_fn`.
+      tx: An Optax gradient transformation.
+      opt_state: The state for `tx`.
+      dropout_rng: PRNG key for stochastic operations.
+      bf16: Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training.
+    """
+
+    step: int
+    apply_fn: Callable = struct.field(pytree_node=False)
+    get_attention_mask_fn: Callable = struct.field(pytree_node=False)
+    params: core.FrozenDict[str, Any]
+    tx: optax.GradientTransformation = struct.field(pytree_node=False)
+    opt_state: optax.OptState
+    dropout_rng: jnp.ndarray
+    max_grad_norm: Optional[float] = 1.0
+
+    def apply_gradients(self, *, grads, to_dtype, **kwargs):
+        """Updates `step`, `params`, `opt_state` and `**kwargs` in return value.
+
+        Note that internally this function calls `.tx.update()` followed by a call
+        to `optax.apply_updates()` to update `params` and `opt_state`.
+
+        Args:
+          grads: Gradients that have the same pytree structure as `.params`.
+          **kwargs: Additional dataclass attributes that should be `.replace()`-ed.
+
+        Returns:
+          An updated instance of `self` with `step` incremented by one, `params`
+          and `opt_state` updated by applying `grads`, and additional attributes
+          replaced as specified by `kwargs`.
+        """
+
+        # clip gradients by global l2 norm
+        casted_max_grad_norm = to_dtype(self.max_grad_norm)
+        g_norm = linear_algebra.global_norm(grads)
+        g_norm = jnp.maximum(casted_max_grad_norm, g_norm)
+        grads = jax.tree_util.tree_map(lambda t: (t / g_norm) * casted_max_grad_norm, grads)
+
+        # perform update step in fp32 and subsequently downcast optimizer states if mixed precision training
+        # grads and opt_state in bf16 (need to upcast), params in fp32 (leave as is)
+        updates, new_opt_state = self.tx.update(to_fp32(grads), to_fp32(self.opt_state), self.params)
+
+        new_params = optax.apply_updates(self.params, updates)
+        return self.replace(
+            step=self.step + 1,
+            params=new_params,
+            opt_state=to_dtype(new_opt_state),
+            **kwargs,
+        )
+
+    @classmethod
+    def create(cls, *, apply_fn, params, tx, to_dtype, **kwargs):
+        """Creates a new instance with `step=0` and initialized `opt_state`."""
+        # downcast optimizer state to bf16 if mixed-precision training
+        opt_state = tx.init(to_dtype(params)) if tx is not None else None
+        return cls(
+            step=0,
+            apply_fn=apply_fn,
+            params=params,
+            tx=tx,
+            opt_state=opt_state,
+            **kwargs,
+        )
+
+    def replicate(self):
+        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
+
+
+@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).
+        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
+    input_padding: Union[bool, str] = "longest"
+    label_padding: Union[bool, str] = "max_length"
+    pad_input_to_multiple_of: Optional[int] = None
+    pad_to_multiple_of_label: Optional[int] = None
+    max_input_length: Optional[float] = None
+    max_label_length: Optional[float] = 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
+        input_features = [{"input_values": feature["input_values"]} for feature in features]
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+
+        # reformat list to dict and set to pytorch format
+        batch = self.processor.feature_extractor.pad(
+            input_features,
+            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_label_length,
+            padding=self.label_padding,
+            pad_to_multiple_of=self.pad_to_multiple_of_label,
+            return_tensors="np",
+        )
+
+        labels = labels_batch["input_ids"]
+        labels = np.ma.array(labels, mask=np.not_equal(labels_batch.attention_mask, 1))
+        labels = labels.filled(fill_value=-100)
+
+        batch["labels"] = labels
+
+        return batch
+
+
+def get_grouped_indices(
+    dataset, batch_size: int, rng: Optional[List[int]] = None, mega_batch_mult: Optional[int] = None
+) -> np.array:
+    """
+    Adapted from the `get_length_grouped_indices` function in the PyTorch Trainer utils file (https://github.com/huggingface/transformers/blob/main/src/transformers/trainer_pt_utils.py#L486)
+    Function that returns a list of indices in which each slice of `batch_size` consecutive indices correspond to elements of similar
+    lengths. To do this, the indices are:
+
+    - randomly permuted (if a JAX rng is specified)
+    - grouped in mega-batches of size `mega_batch_mult * batch_size`
+    - sorted by length in each mega-batch
+
+    The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of
+    maximum length placed first, so that an OOM happens sooner rather than later.
+    """
+    lengths = dataset["input_length"]
+
+    # Default for mega_batch_mult: 50 or the number to get 4 megabatches, whichever is smaller.
+    if mega_batch_mult is None:
+        mega_batch_mult = min(len(lengths) // (batch_size * 4), 50)
+        # Just in case, for tiny datasets
+        if mega_batch_mult == 0:
+            mega_batch_mult = 1
+
+    # We need to use JAX for the random permutation as the PRNG key will be set based on the seed outside of the sampler.
+    num_samples = len(lengths)
+    indices = jax.random.permutation(rng, np.arange(num_samples)) if rng is not None else np.arange(num_samples)
+
+    megabatch_size = mega_batch_mult * batch_size
+    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
+    megabatches = [list(sorted(megabatch, key=lambda i: lengths[i], reverse=True)) for megabatch in megabatches]
+
+    # The rest is to get the biggest batch first.
+    # Since each megabatch is sorted by descending length, the longest element is the first
+    megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches]
+    max_idx = np.argmax(megabatch_maximums).item()
+    # Switch to put the longest batch in first position
+    # (note that this is different to the PT grouped sampler in which we only put the longest element in the first position, and not its batch)
+    megabatches[0], megabatches[max_idx] = megabatches[max_idx], megabatches[0]
+
+    megabatches = np.array([i for megabatch in megabatches for i in megabatch])
+
+    return megabatches
+
+
+def generate_batch_splits(samples_idx: np.ndarray, batch_size: int, drop_last=True) -> np.ndarray:
+    """Generate batches of data for a specified batch size from sample indices. If the dataset size is not divisible by
+    the batch size and `drop_last` is `True`, the last incomplete batch is dropped. Else, it is returned."""
+    num_samples = len(samples_idx)
+    if drop_last:
+        samples_to_remove = num_samples % batch_size
+        if samples_to_remove != 0:
+            samples_idx = samples_idx[:-samples_to_remove]
+        sections_split = num_samples // batch_size
+        samples_idx = samples_idx.reshape((sections_split, batch_size))
+    else:
+        sections_split = math.ceil(num_samples / batch_size)
+        samples_idx = np.array_split(samples_idx, sections_split)
+    return samples_idx
+
+
+def write_train_metric(summary_writer, train_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)
+
+
+def write_eval_metric(summary_writer, eval_metrics, step, pred_str=None):
+    for metric_name, value in eval_metrics.items():
+        summary_writer.scalar(f"eval_{metric_name}", value, step)
+
+    if pred_str is not None:
+        # write output actual predictions for debugging
+        summary_writer.text("eval_predictions", "\n".join(pred_str), step)
+
+
+def write_wandb_log(metrics, step, prefix=None):
+    if jax.process_index() == 0:
+        log_metrics = {}
+        for k, v in metrics.items():
+            if "layer" in k:
+                log_metrics[f"{k}/"] = v
+            elif prefix is not None:
+                log_metrics[f"{prefix}/{k}"] = v
+            else:
+                log_metrics[k] = v
+        wandb.log(log_metrics, step)
+
+
+def write_wandb_pred(pred_str, label_str, step, num_log=50, prefix="eval"):
+    if jax.process_index() == 0:
+        # convert str data to a wandb compatible format
+        str_data = [[label_str[i], pred_str[i]] for i in range(len(pred_str))]
+        # we'll log the first 50 predictions for each epoch
+        wandb.log(
+            {
+                f"{prefix}/step_{int(step / 1000)}k": wandb.Table(
+                    columns=["label_str", "pred_str"], data=str_data[:num_log]
+                )
+            },
+            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 ctc_loss(
+    logits,
+    logits_attention_mask,
+    labels,
+    blank_id,
+    loss_reduction="mean",
+    output_emission_dict=False,
+    log_epsilon=-100000.0,
+):
+    """Computes CTC loss.
+    This function performs forward computation over an FSA with `N * 2` states
+    where `N` is the max number of labels. The states are split into two groups:
+    Phi states and emission states. a phi-state accepts repetition of
+    phi (blank)-symbols and transits to emission state when the correct label is
+    observed. An emission state accepts repetition of the label and transits to
+    the next phi states at any time (so called epsilon-transition).
+    Below, `B` denotes the batch size, `T` denotes the time steps in `logits`,
+    and `N` denotes the time steps in `labels`.
+    Args:
+      logits: (B, T, K)-array containing log-probabilities of each class.
+      logitpaddings: (B, T)-array. Padding indicators for `logits`.
+      labels: (B, N)-array containing reference integer labels.
+      labelpaddings: (B, N)-array. Padding indicators for `labels`. Currently,
+        `labels` must be right-padded, i.e. each row of `labelpaddings` must be
+        repetition of zeroes, followed by repetition of ones.
+      blank_id: Id for blank token.
+      loss_reduction: one of "mean", "sum", "default"
+        - "none": no reduction is applied.
+        - "mean": output loss will be divided by target lengths and then the
+          mean over the batch is taken.
+        - "sum": output loss are summed over batch
+      output_emission_dict: whether to output additional information about the emission probs
+    Returns:
+      A pair of `(per_seq_loss, aux)`.
+      per_seq_loss:
+        (B,)-array containing loss values for each sequence in the batch.
+      aux: Dictionary containing interim variables used for computing losses.
+        aux['logalpha_phi']: (T, B, N+1)-array. Log-forward-probabilities of each
+          phi-state corresponding to the n-th label.
+        aux['logalpha_emit']: (T, B, N)-array. Log-forward-probabilities of each
+          emission-state corresponding to the n-th label.
+        aux['logprobs_phi']: (T, B, 1)-array. Probability of the phi-symbol
+          corresponding to each time frame.
+        aux['logprobs_emit']: (T, B, N)-array. Probability of the n-th label
+          corresponding to each time frame.
+    """
+    # label paddings are indicated by -100
+    labelpaddings = labels < 0
+    # logit paddings are the inverse of attention_mask
+    logitpaddings = ~logits_attention_mask
+
+    # Copied from https://github.com/tensorflow/lingvo/blob/master/lingvo/jax/layers/ctc_objectives.py
+    batchsize, unused_maxinputlen, num_classes = logits.shape
+    batchsize_, maxlabellen = labels.shape
+
+    logprobs = jax.nn.log_softmax(logits)
+    labellens = maxlabellen - jnp.sum(labelpaddings, axis=1).astype(jnp.int32)
+
+    # repeat[b, n] == 1.0 when label[b, n] == label[b, n+1].
+    repeat = (labels[:, :-1] == labels[:, 1:]).astype(jnp.float32)
+    repeat = jnp.pad(repeat, ((0, 0), (0, 1)))
+
+    logprobs_phi = logprobs[:, :, blank_id : blank_id + 1]  # [B, T, 1]
+    logprobs_phi = jnp.transpose(logprobs_phi, (1, 0, 2))  # [T, B, 1]
+
+    one_hot = jax.nn.one_hot(labels, num_classes=num_classes)  # [B, N, K]
+    logprobs_emit = jnp.einsum("btk,bnk->btn", logprobs, one_hot)
+    logprobs_emit = jnp.transpose(logprobs_emit, (1, 0, 2))  # [T, B, N]
+
+    logalpha_phi_init = jnp.ones((batchsize, maxlabellen + 1)) * log_epsilon  # [B, N]
+    logalpha_phi_init = logalpha_phi_init.at[:, 0].set(0.0)
+    logalpha_emit_init = jnp.ones((batchsize, maxlabellen)) * log_epsilon  # [B, N]
+
+    def loop_body(prev, x):
+        prev_phi, prev_emit = prev
+        # emit-to-phi epsilon transition, except if the next label is repetition
+        prev_phi_orig = prev_phi
+        prev_phi = prev_phi.at[:, 1:].set(jnp.logaddexp(prev_phi[:, 1:], prev_emit + log_epsilon * repeat))
+
+        logprob_emit, logprob_phi, pad = x
+
+        # phi-to-emit transition
+        next_emit = jnp.logaddexp(prev_phi[:, :-1] + logprob_emit, prev_emit + logprob_emit)
+        # self-loop transition
+        next_phi = prev_phi + logprob_phi
+        # emit-to-phi blank transition only when the next label is repetition
+        next_phi = next_phi.at[:, 1:].set(
+            jnp.logaddexp(next_phi[:, 1:], prev_emit + logprob_phi + log_epsilon * (1.0 - repeat))
+        )
+
+        pad = pad.reshape((batchsize, 1))
+        next_emit = pad * prev_emit + (1.0 - pad) * next_emit
+        next_phi = pad * prev_phi_orig + (1.0 - pad) * next_phi
+
+        return (next_phi, next_emit), (next_phi, next_emit)
+
+    xs = (logprobs_emit, logprobs_phi, logitpaddings.transpose((1, 0)))
+    _, (logalpha_phi, logalpha_emit) = jax.lax.scan(loop_body, (logalpha_phi_init, logalpha_emit_init), xs)
+
+    # last row needs to be updated with the last epsilon transition
+    logalpha_phi_last = logalpha_phi[-1].at[:, 1:].set(jnp.logaddexp(logalpha_phi[-1, :, 1:], logalpha_emit[-1]))
+    logalpha_phi = logalpha_phi.at[-1].set(logalpha_phi_last)
+
+    # extract per_seq_loss
+    one_hot = jax.nn.one_hot(labellens, num_classes=maxlabellen + 1)  # [B, N+1]
+    per_seq_loss = -jnp.einsum("bn,bn->b", logalpha_phi_last, one_hot)
+
+    if loss_reduction == "mean":
+        target_lengths = labelpaddings.shape[-1] - labelpaddings.sum(axis=-1)
+        loss = (per_seq_loss / target_lengths).mean()
+    elif loss_reduction == "sum":
+        loss = per_seq_loss.sum()
+    else:
+        loss = per_seq_loss
+
+    if not output_emission_dict:
+        return loss
+
+    return loss, {
+        "logalpha_phi": logalpha_phi,
+        "logalpha_emit": logalpha_emit,
+        "logprobs_phi": logprobs_phi,
+        "logprobs_emit": logprobs_emit,
+    }
+
+
+def make_dataset(data_args, seed=42):
+    # Pre-processing dataset
+    import re
+
+    def map_nst(entry):
+        text = entry["text"].lower()
+        text = text.replace("(...vær stille under dette opptaket...)", "")
+        text = re.sub('[áàâ]', 'a', text)
+        text = re.sub('[ä]', 'æ', text)
+        text = re.sub('[éèëê]', 'e', text)
+        text = re.sub('[íìïî]', 'i', text)
+        text = re.sub('[óòöô]', 'o', text)
+        text = re.sub('[ö]', 'ø', text)
+        text = re.sub('[ç]', 'c', text)
+        text = re.sub('[úùüû]', 'u', text)
+        # text = re.sub('\\(?=(Punktum|Komma|Utropstegn|Spørsmålstegn))', ' ', text)
+        text = re.sub('\s+', ' ', text)
+        return {"text": text}
+
+    def filter_nst(entry):
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.match(entry["type"], "pIW|CA"):
+            return False  # Spelling out words
+        return True
+
+    def filter_npsc(entry):
+        # False if there are digits in the text
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.search("\d", entry["text"]):
+            return False
+        return True
+
+    def map_npsc(entry):
+        batch = {"text": entry["text"].lower()}
+        batch["text"] = re.sub('[áàâ]', 'a', batch["text"])
+        batch["text"] = re.sub('[ä]', 'æ', batch["text"])
+        batch["text"] = re.sub('[éèëê]', 'e', batch["text"])
+        batch["text"] = re.sub('[íìïî]', 'i', batch["text"])
+        batch["text"] = re.sub('[óòöô]', 'o', batch["text"])
+        batch["text"] = re.sub('[ö]', 'ø', batch["text"])
+        batch["text"] = re.sub('[ç]', 'c', batch["text"])
+        batch["text"] = re.sub('[úùüû]', 'u', batch["text"])
+        batch["text"] = re.sub('\s', ' ', batch["text"])
+        batch["text"] = re.sub('<ee>', 'eee', batch["text"])
+        batch["text"] = re.sub('<qq>', 'qqq', batch["text"])
+        batch["text"] = re.sub('<mm>', 'mmm', batch["text"])
+        batch["text"] = re.sub('<inaudible>', 'xxx', batch["text"])
+        # batch["text"] = re.sub('<inaudible>', '?', batch["text"])
+        if "<" in batch["text"]:
+            raise ValueError(batch["text"])
+        return batch
+
+    nst = datasets.load_dataset("NbAiLab/NST", "no-close")
+    npsc = datasets.load_dataset("NbAiLab/NPSC", "16K_mp3")
+    # TODO NST_hesitate
+
+    split = len(npsc[data_args.train_split_name]) / (len(npsc[data_args.train_split_name]) + len(npsc[data_args.eval_split_name]))  # Use same train/val ratio as NPSC
+    nst_train = nst[data_args.train_split_name].train_test_split(train_size=split, seed=seed)
+    nst[data_args.train_split_name] = nst_train["train"]
+    nst[data_args.eval_split_name] = nst_train["test"]
+
+    nst = nst.filter(filter_nst).map(
+        map_nst,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NST",
+    ).shuffle(seed=seed)
+    npsc = npsc.filter(filter_npsc).map(
+        map_npsc,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NPSC",
+    ).shuffle(seed=seed)
+
+    npsc_base = npsc.remove_columns([col for col in npsc[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+    nst_base = nst.remove_columns([col for col in nst[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+
+    combined = {}
+    for split in data_args.train_split_name, data_args.eval_split_name, data_args.test_split_name:
+        probs = np.array([len(nst_base[split]), len(npsc_base[split])])  # Weight by number of examples
+        probs = (probs / probs.sum()).tolist()
+        comb = datasets.interleave_datasets([nst_base[split], npsc_base[split]], probabilities=probs, seed=seed)
+        combined[split] = comb
+
+    return datasets.DatasetDict(**combined)
+
+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, FlaxTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # 2. 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()
+
+    # Set up wandb run
+    if jax.process_index() == 0:
+        wandb.init(project=data_args.wandb_project, name=data_args.wandb_name, job_type=data_args.wandb_job_type)
+
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set the default TPU matmul precision and display the number of devices
+    jax.config.update("jax_default_matmul_precision", training_args.matmul_precision)
+    logger.info(f"JAX devices: {jax.device_count()}, matmul precision: {training_args.matmul_precision}")
+
+    # 4. Load dataset
+
+    set_seed(training_args.seed)
+    raw_datasets = make_dataset(data_args, seed=training_args.seed)
+   
+    # raw_datasets = DatasetDict()
+
+    # if training_args.do_train:
+    #     raw_datasets[data_args.train_split_name] = 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[data_args.eval_split_name] = 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 training_args.do_predict:
+    #     test_split = data_args.test_split_name.split("+")
+    #     for split in test_split:
+    #         raw_datasets[split] = load_dataset(
+    #             data_args.dataset_name,
+    #             data_args.dataset_config_name,
+    #             split=split,
+    #             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 and not training_args.do_predict:
+        raise ValueError(
+            "Cannot not train, not do evaluation and not do prediction. At least one of "
+            "training, evaluation or prediction has to be done."
+        )
+
+    # if not training, there is no need to run multiple epochs
+    if not training_args.do_train:
+        training_args.num_train_epochs = 1
+
+    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
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    config = Wav2Vec2Config.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,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    # update config according to training args, model args, and tokenizer attributes
+    config.update(
+        {
+            "feat_proj_dropout": model_args.feat_proj_dropout,
+            "attention_dropout": model_args.attention_dropout,
+            "hidden_dropout": model_args.hidden_dropout,
+            "final_dropout": model_args.final_dropout,
+            "mask_time_prob": model_args.mask_time_prob,
+            "mask_time_length": model_args.mask_time_length,
+            "mask_feature_prob": model_args.mask_feature_prob,
+            "mask_feature_length": model_args.mask_feature_length,
+            "gradient_checkpointing": training_args.gradient_checkpointing,
+            "layerdrop": model_args.layerdrop,
+            "ctc_loss_reduction": model_args.ctc_loss_reduction,
+            "ctc_zero_infinity": model_args.ctc_zero_infinity,
+            "pad_token_id": tokenizer.pad_token_id,
+            "vocab_size": tokenizer.vocab_size,  # len(tokenizer),
+            "activation_dropout": model_args.activation_dropout,
+        }
+    )
+
+    if tokenizer.do_lower_case and data_args.dataset_name != "librispeech_asr":
+        raise ValueError(
+            "Setting the tokenizer attribute `do_lower_case` to `True` converts all input strings to "
+            "uppercase prior to tokenization. This should only be done when the tokenizer is built on an uppercased corpus,"
+            "i.e. for the dataset `librispeech_asr` only. If your dataset is not `librispeech_asr`, the tokenizer is mostly likely "
+            "built on an lowercased corpus. In this case, set `tokenizer.do_lower_case` to ``False`."
+        )
+
+    if training_args.precision == "full_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = True
+    elif training_args.precision == "half_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = False
+    else:
+        dtype = jnp.float32
+        training_args.mixed_precision = False
+
+    try:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+        )
+    except:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+            from_pt=True,
+        )
+
+    # 6. Resample speech dataset ALWAYS
+    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_target_length = data_args.max_label_length
+    min_target_length = data_args.min_label_length
+    pad_input_to_multiple_of = data_args.pad_input_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
+    dataset_name = data_args.dataset_name
+    chars_to_ignore = ', ? . ! - ; : " “ % ‘ ” ?'.split(" ")
+    chars_to_ignore_regex = f'[{"".join(chars_to_ignore)}]'
+    # gigaspeech_punctuation = {" <comma>": ",", " <period>": ".", " <questionmark>": "?", " <exclamationpoint>": "!"}
+    # gigaspeech_disfluencies = ["<other>", "<sil>"]
+    # swb_disfluencies = ["[noise]", "[laughter]", "[silence]", "<a_aside>", "<b_aside>", "<e_aside>", "[laughter-",
+    #                 "[vocalized-noise]", "_1"]
+    # swb_punctuations = ["{", "}", "[", "]-", "]"]
+    # earnings_disfluencies = ["<crosstalk>", "<affirmative>", "<inaudible>", "inaudible", "<laugh>", "<unk>"]
+    ignore_segments = ["ignore_time_segment_in_scoring", "<noise>", "<music>", "[noise]", "[laughter]", "[silence]",
+                       "[vocalized-noise]", "<crosstalk>", "<affirmative>", "<inaudible>", "<laugh>", "<other>", "<sil>", ""]
+
+    if training_args.do_train and data_args.max_train_samples is not None:
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].select(range(data_args.max_train_samples))
+
+    if training_args.do_eval and data_args.max_eval_samples is not None:
+        raw_datasets[data_args.eval_split_name] = raw_datasets[data_args.eval_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_predict and data_args.max_test_samples is not None:
+        raw_datasets[data_args.test_split_name] = raw_datasets[data_args.test_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_train and data_args.remove_punctuation:
+
+        def remove_punctuation(batch):
+            batch[text_column_name] = (
+                re.sub(chars_to_ignore_regex, "", batch[text_column_name]).replace("'", "").replace('"', "")
+            )
+
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].map(
+            remove_punctuation,
+            num_proc=data_args.preprocessing_num_workers,
+            desc="removing punctuation from train split",
+        )
+
+    # filter data where the targets are ignored in scoring
+    def is_target_labels(input_str):
+        return input_str.lower() not in ignore_segments
+
+    raw_datasets = raw_datasets.filter(
+            is_target_labels,
+            num_proc=num_workers,
+            input_columns=[text_column_name],
+            desc="filtering data where the targets are ignored in scoring",
+        )
+
+    def prepare_dataset(batch):
+        # process audio
+        try:
+            sample = batch[audio_column_name]
+        except ValueError:
+            sample = {"array": np.array([0.]), "sampling_rate": feature_extractor.sampling_rate}
+        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
+        # process audio length
+        batch[model_input_name] = inputs.input_values[0]
+        batch["input_length"] = len(batch["input_values"])
+
+        # process targets
+        input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
+
+        # if dataset_name == "google/xtreme_s":
+        #     # Finally, we tokenize the processed text
+        #     batch["labels"] = tokenizer(input_str).input_ids
+        #     batch["labels_length"] = len(batch["labels"])
+        #     return batch
+
+        # # Common Voice 9
+        # if input_str.startswith('"') and input_str.endswith('"'):
+        #     # we can remove trailing quotation marks as they do not affect the transcription
+        #     input_str = input_str[1:-1]
+        # # normalize quotation marks
+        # input_str = re.sub(r'["“”]', '"', input_str)
+        # # normalize apostrophes
+        # input_str = re.sub(r"[’']", "'", input_str)
+        # # normalize hyphens
+        # input_str = re.sub(r"[—–]", "-", input_str)
+        # # replace double quotation marks with single
+        # input_str = input_str.replace('""', '"')
+        # if dataset_name == "mozilla-foundation/common_voice_9_0" and len(input_str):
+        #     # for CV9, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # TEDLIUM-3
+        # # delete the <unk> token from the text and replace spaced apostrophes with un-spaced
+        # input_str = input_str.replace("<unk>", "").replace(" '", "'")
+
+        # # GigaSpeech
+        # for disfluency in gigaspeech_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # convert spelled out punctuation to symbolic form
+        # for punctuation, replacement in gigaspeech_punctuation.items():
+        #     input_str = input_str.replace(punctuation, replacement)
+        # if dataset_name == "speechcolab/gigaspeech" and len(input_str):
+        #     # for GS, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # SWB
+        # for disfluency in swb_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # remove parenthesised text (test data only)
+        # input_str = re.sub("[\(].*?[\)]", "", input_str)
+        # for punctuation in swb_punctuations:
+        #     input_str = input_str.replace(punctuation, "")
+        # # replace anomalous words with their correct transcriptions
+        # split_str = input_str.split("/")
+        # if len(split_str) > 1:
+        #     input_str = " ".join(
+        #         [" ".join([" ".join(i.split(" ")[:-1]) for i in split_str])] + [split_str[-1].split(" ")[-1]])
+
+        # # Earnings 22
+        # for disfluency in earnings_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # replace mal-formatted ellipsis
+        # input_str = input_str.replace("…", ".")
+
+        # JIWER compliance
+        # remove multiple spaces
+        input_str = re.sub(r"\s\s+", " ", input_str)
+        # strip trailing spaces
+        input_str = input_str.strip()
+
+        # Finally, we tokenize the processed text
+        batch["labels"] = tokenizer(input_str).input_ids
+        batch["labels_length"] = len(batch["labels"])
+        return batch
+
+    vectorized_datasets = raw_datasets.map(
+        prepare_dataset,
+        remove_columns=next(iter(raw_datasets.values())).column_names,
+        num_proc=num_workers,
+        desc="preprocess dataset",
+    )
+
+    # filter data with inputs shorter than min_input_length or longer than max_input_length
+    def is_audio_in_length_range(length):
+        return length > min_input_length and length < max_input_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_audio_in_length_range,
+        num_proc=num_workers,
+        input_columns=["input_length"],
+    )
+
+    # filter data with targets shorter than min_target_length or longer than max_target_length
+    def is_labels_in_length_range(length):
+        return length > min_target_length  # and length < max_target_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_labels_in_length_range,
+        num_proc=num_workers,
+        input_columns=["labels_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 Metrics
+    wer_metric = load_metric("wer")
+    cer_metric = load_metric("cer")
+
+    def compute_metrics(pred_ids: List[List[int]], label_ids: List[List[int]]):
+        padded_ids = np.where(np.asarray(label_ids) == -100, tokenizer.pad_token_id, np.asarray(label_ids))
+
+        pred_str = tokenizer.batch_decode(pred_ids)
+        # we do not want to group tokens when computing the metrics
+        label_str = tokenizer.batch_decode(padded_ids, group_tokens=False)
+
+        wer = wer_metric.compute(predictions=pred_str, references=label_str)
+        cer = cer_metric.compute(predictions=pred_str, references=label_str)
+
+        return {"wer": wer, "cer": cer}, pred_str, label_str
+
+    # 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,
+        input_padding="longest",
+        pad_input_to_multiple_of=pad_input_to_multiple_of,
+        max_label_length=data_args.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."
+        )
+
+    # 10. Handle the repository creation
+    if training_args.push_to_hub:
+        with open(os.path.join(training_args.output_dir, ".gitattributes"), "r+") as f:
+            git_lfs_extensions = f.read()
+            if "*.wandb" not in git_lfs_extensions:
+                f.write("*.wandb filter=lfs diff=lfs merge=lfs -text")
+        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
+        repo = Repository(training_args.output_dir, clone_from=repo_name)
+
+    # 11. Initialize our training
+    rng = jax.random.PRNGKey(training_args.seed)
+    rng, dropout_rng = jax.random.split(rng)
+
+    # Store some constants
+    max_steps = int(training_args.max_steps)
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
+    batch_size_per_update = train_batch_size * gradient_accumulation_steps
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
+    to_dtype = to_bf16 if training_args.mixed_precision else to_fp32
+
+    if training_args.do_train:
+        num_train_samples = len(vectorized_datasets[data_args.train_split_name])
+        steps_per_epoch = num_train_samples // batch_size_per_update
+        if max_steps > 0:
+            num_epochs = -(training_args.max_steps // -steps_per_epoch)
+            total_train_steps = max_steps
+        else:
+            num_epochs = int(training_args.num_train_epochs)
+            total_train_steps = steps_per_epoch * num_epochs
+
+        # Create learning rate schedule
+        # Create learning rate schedule
+        linear_decay_lr_schedule_fn = create_learning_rate_fn(
+            total_train_steps,
+            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.
+        # Note that this mask is specifically adapted for FlaxWav2Vec2 and FlaxBart.
+        # For FlaxT5, one should correct the layer norm parameter naming
+        # accordingly - see `run_t5_mlm_flax.py` e.g.
+        def decay_mask_fn(params):
+            flat_params = traverse_util.flatten_dict(params)
+            layer_norm_params = [
+                (name, "scale")
+                for name in ["layer_norm", "self_attn_layer_norm", "layernorm_embedding", "final_layer_norm"]
+            ]
+            flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_params) for path in flat_params}
+            return traverse_util.unflatten_dict(flat_mask)
+
+        if training_args.adafactor:
+            # Create Adafactor optimizer
+            optim = optax.adafactor(
+                learning_rate=linear_decay_lr_schedule_fn,
+                dtype_momentum=jnp.bfloat16 if training_args.mixed_precision else jnp.float32,
+                weight_decay_rate=training_args.weight_decay,
+                weight_decay_mask=decay_mask_fn,
+            )
+        else:
+            # Create AdamW optimizer
+            optim = 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,
+            )
+
+        # Optax MultiSteps for gradient accumulation. We'll only call this optimizer transformation if gradient accumulation is required (i.e. gradient accumulation steps > 1)
+        if training_args.multisteps and gradient_accumulation_steps > 1:
+            optim = optax.MultiSteps(optim, gradient_accumulation_steps, use_grad_mean=False)
+    else:
+        num_epochs = 0
+        total_train_steps = 0
+        num_train_samples = 0
+        optim = None
+
+    # Setup train state
+    state = MixedPrecisionTrainState.create(
+        apply_fn=model.__call__,
+        get_attention_mask_fn=model._get_feature_vector_attention_mask,
+        params=model.params,
+        tx=optim,
+        to_dtype=to_dtype,
+        dropout_rng=dropout_rng,
+        max_grad_norm=training_args.max_grad_norm,
+    )
+
+    # Replicate the train state on each device
+    state = state.replicate()
+    blank_id = model.config.pad_token_id
+
+    # Define gradient update step fn
+    def train_step(state, batch):
+        # only one single rng per grad step, with or without accumulation, as the graph should be identical over one effective training batch
+        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
+
+        def compute_loss(params, minibatch):
+            labels = minibatch.pop("labels")
+            logits = state.apply_fn(
+                **minibatch,
+                params=params,
+                dropout_rng=dropout_rng,
+                freeze_feature_encoder=model_args.freeze_feature_encoder,
+                train=True,
+            )[0]
+            logits_mask = state.get_attention_mask_fn(logits.shape[1], batch["attention_mask"])
+            loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+            return loss
+
+        grad_fn = jax.value_and_grad(compute_loss)
+
+        if gradient_accumulation_steps == 1 or training_args.multisteps:
+            loss, grad = grad_fn(to_dtype(state.params), batch)
+
+        # Custom gradient accumulation
+        else:
+            # add a first dimension over gradient_accumulation_steps for minibatch slices
+            batch = jax.tree_util.tree_map(
+                lambda x: x.reshape(
+                    gradient_accumulation_steps, training_args.per_device_train_batch_size, *x.shape[1::]
+                ),
+                batch,
+            )
+
+            def accum_minibatch_step(accum_grad, minibatch):
+                # compute loss, num labels and grad over minibatch and accumulate
+                loss, grad = grad_fn(to_dtype(state.params), minibatch)
+                return jax.tree_util.tree_map(jnp.add, accum_grad, grad), loss
+
+            # create an initial state for accumulating losses, num labels and gradients
+            init_grad = jax.tree_util.tree_map(jnp.zeros_like, to_dtype(state.params))
+            # loop accum minibatch step over the number of gradient accumulation steps
+            grad, loss = jax.lax.scan(accum_minibatch_step, init_grad, batch)
+
+        # update state
+        new_state = state.apply_gradients(
+            grads=grad,
+            dropout_rng=new_dropout_rng,
+            to_dtype=to_dtype,
+        )
+
+        # compute gradient norms over all layers and globally for detailed monitoring
+        layer_grad_norm = jax.tree_util.tree_map(jnp.linalg.norm, grad)
+        logs = {
+            "layer_grad_norm": layer_grad_norm,
+            "grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm)),
+        }
+
+        # compute parameter norms over all layers and globally for detailed monitoring
+        layer_param_norm = jax.tree_util.tree_map(jnp.linalg.norm, new_state.params)
+        logs["layer_param_norm"] = layer_param_norm
+        logs["param_norm"] = jnp.linalg.norm(jax.tree_util.tree_leaves(layer_param_norm))
+
+        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+        metrics.update(logs)
+
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+
+        return new_state, metrics
+
+    # Define eval fn
+    def eval_step(params, batch):
+        labels = batch.pop("labels")
+        logits = model(**batch, params=params, train=False)[0]
+
+        logits_mask = model._get_feature_vector_attention_mask(logits.shape[1], batch["attention_mask"])
+        loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+        pred_ids = jnp.argmax(logits, axis=-1)
+
+        # summarize metrics
+        metrics = {"loss": loss}
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+        return metrics, pred_ids
+
+    # Create parallel version of the train and eval step
+    if training_args.do_train:
+        p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
+
+    if training_args.do_eval:
+        p_eval_step = jax.pmap(eval_step, "batch")
+
+    def run_evaluation(step):
+        if training_args.do_eval:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the eval dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[data_args.eval_split_name], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
+                samples = [vectorized_datasets[data_args.eval_split_name][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, step, prefix="eval")
+            write_wandb_pred(pred_str, label_str, step)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, step, pred_str=pred_str)
+
+    def save_checkpoint(step):
+        # save 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"{wandb.run.id}: saving weights and logs of step {int(step / 1000)}k", blocking=False)
+
+    skip_epochs = data_args.skip_steps // (num_train_samples // batch_size_per_update)
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_train_samples}")
+    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"  Num gradient accumulation steps = {gradient_accumulation_steps}")
+    logger.info(f"  Total train batch size (w. parallel & distributed) = {batch_size_per_update}")
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+    logger.info(f"  Gradient checkpointing: {config.gradient_checkpointing}")
+    logger.info(f"  Use scan: {config.use_scan}")
+    logger.info(f"  Fuse matmuls: {config.fuse_matmuls}")
+    logger.info(f"  Skipping: {data_args.skip_steps} steps ({skip_epochs} epochs)")
+
+    train_time = cur_step = 0
+    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
+    for epoch in epochs:
+        if training_args.do_train:
+            # ======================== Training ================================
+            train_start = time.time()
+
+            if epoch < skip_epochs:
+                logger.info(f"Skipping epoch {epoch + 1}")
+                continue
+
+            # Create sampling rng
+            rng, input_rng = jax.random.split(rng)
+
+            # Generate an epoch by randomly shuffling sampling indices from the train dataset and grouping by length
+            train_samples_idx = get_grouped_indices(vectorized_datasets[data_args.train_split_name], batch_size_per_update, input_rng)
+            train_batch_idx = generate_batch_splits(train_samples_idx, batch_size_per_update)
+
+            if data_args.skip_steps > cur_step:
+                logger.info(f"Skipping {data_args.skip_steps - (epoch * (num_train_samples // batch_size_per_update))} steps...")
+            # Gather the indices for creating the batch and do a training step
+            for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1), 1):
+                cur_step = epoch * (num_train_samples // batch_size_per_update) + step
+                if cur_step <= data_args.skip_steps:
+                    continue
+
+                samples = [vectorized_datasets[data_args.train_split_name][int(idx)] for idx in batch_idx]              
+                batch = data_collator(samples)
+                batch = shard(batch.data)
+                try:
+                    state, train_metric = p_train_step(state, batch)
+                except TypeError as e:
+                    logger.warning("Encountered following error: \n", e)
+
+
+                if cur_step % training_args.logging_steps == 0:
+                    # Save metrics
+                    train_metric = unreplicate(train_metric)
+                    train_time += time.time() - train_start
+                    # need to upcast all device arrays to fp32 for wandb logging (jnp.bfloat16 not supported) -> do this here OR in train_step
+                    write_wandb_log(to_fp32(train_metric), cur_step, prefix=data_args.train_split_name)
+                    # we won't log to tensorboard for now (it is fiddly logging param and grad norms on a layer-by-layer basis)
+                    # if has_tensorboard and jax.process_index() == 0:
+                    # write_train_metric(summary_writer, train_metrics, train_time, cur_step)
+
+                    epochs.write(
+                        f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']}, Gradient Norm: {train_metric['grad_norm']})"
+                    )
+
+                if cur_step % total_train_steps == 0:
+                    break
+
+                if training_args.eval_steps and cur_step % training_args.eval_steps == 0:
+                    run_evaluation(cur_step)
+
+                if cur_step % training_args.save_steps == 0:
+                    save_checkpoint(cur_step)
+
+            if training_args.eval_steps == 0 and (epoch + 1) != num_epochs:
+                # run evaluation at the end of the epoch if eval steps are not specified
+                run_evaluation(cur_step)
+                save_checkpoint(cur_step)
+
+    if training_args.do_train:
+        save_checkpoint(cur_step)
+
+    cur_step = max_steps if max_steps > 0 else cur_step  # set step to max steps so that eval happens in alignment with training
+
+    if training_args.do_eval:
+        run_evaluation(cur_step)
+
+    # TODO: collapse 'do_predict' into the run_evaluation function
+    if training_args.do_predict:
+        for split in [data_args.test_split_name]:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the test dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[split], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc=f"Predicting {split}...", position=2)):
+                samples = [vectorized_datasets[split][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({cur_step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, cur_step, prefix=split)
+            write_wandb_pred(pred_str, label_str, cur_step, prefix=split)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, cur_step, pred_str=pred_str)
+
+
+if __name__ == "__main__":
+    main()

wandb/run-20220810_073735-23avj35z/files/config.yaml

@@ -0,0 +1,33 @@
+wandb_version: 1
+
+_wandb:
+  desc: null
+  value:
+    cli_version: 0.12.9
+    code_path: code/run_flax_speech_recognition_ctc.py
+    framework: huggingface
+    huggingface_version: 4.21.0
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    python_version: 3.8.10
+    start_time: 1660117055
+    t:
+      1:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      2:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      3:
+      - 13
+      4: 3.8.10
+      5: 0.12.9
+      6: 4.21.0
+      8:
+      - 5

wandb/run-20220810_073735-23avj35z/files/diff.patch

@@ -0,0 +1,27 @@
+diff --git a/wandb/debug-internal.log b/wandb/debug-internal.log
+index 23926ef..9213b33 120000
+--- a/wandb/debug-internal.log
++++ b/wandb/debug-internal.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug-internal.log
+
++run-20220810_073735-23avj35z/logs/debug-internal.log
+
+diff --git a/wandb/debug.log b/wandb/debug.log
+index 279853d..bcac724 120000
+--- a/wandb/debug.log
++++ b/wandb/debug.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug.log
+
++run-20220810_073735-23avj35z/logs/debug.log
+
+diff --git a/wandb/latest-run b/wandb/latest-run
+index f069a7a..1406fac 120000
+--- a/wandb/latest-run
++++ b/wandb/latest-run
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4
+
++run-20220810_073735-23avj35z
+

wandb/run-20220810_073735-23avj35z/files/output.log

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

wandb/run-20220810_073735-23avj35z/files/requirements.txt

@@ -0,0 +1,158 @@
+absl-py==1.0.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+appdirs==1.4.4
+astunparse==1.6.3
+async-timeout==4.0.2
+attrs==21.4.0
+audioread==2.1.9
+backcall==0.2.0
+cachetools==4.2.4
+certifi==2021.10.8
+cffi==1.15.1
+charset-normalizer==2.0.10
+chex==0.1.3
+click==8.0.3
+cloud-tpu-client==0.10
+cloud-tpu-profiler==2.4.0
+clu==0.0.6
+colorama==0.4.5
+commonmark==0.9.1
+configparser==5.2.0
+contextlib2==21.6.0
+cycler==0.11.0
+datasets==2.4.0
+decorator==5.1.0
+dill==0.3.4
+dm-tree==0.1.6
+docker-pycreds==0.4.0
+etils==0.6.0
+exceptiongroup==1.0.0rc8
+filelock==3.4.2
+flatbuffers==2.0
+flax==0.5.3
+fonttools==4.28.5
+frozenlist==1.2.0
+fsspec==2021.11.1
+future==0.18.2
+gast==0.4.0
+gitdb==4.0.9
+gitpython==3.1.26
+google-api-core==1.31.5
+google-api-python-client==1.8.0
+google-auth-httplib2==0.1.0
+google-auth-oauthlib==0.4.6
+google-auth==2.3.3
+google-pasta==0.2.0
+googleapis-common-protos==1.54.0
+grpcio==1.43.0
+h5py==3.6.0
+httplib2==0.20.2
+huggingface-hub==0.2.1
+hypothesis==6.53.0
+idna==3.3
+importlib-metadata==4.10.0
+importlib-resources==5.4.0
+ipython==7.31.0
+jax==0.3.15
+jaxlib==0.3.15
+jedi==0.18.1
+jiwer==2.3.0
+joblib==1.1.0
+keras-preprocessing==1.1.2
+keras==2.7.0
+kiwisolver==1.3.2
+libclang==12.0.0
+librosa==0.9.2
+libtpu-nightly==0.1.dev20220722
+llvmlite==0.39.0
+markdown==3.3.6
+matplotlib-inline==0.1.3
+matplotlib==3.5.1
+ml-collections==0.1.0
+msgpack==1.0.3
+multidict==5.2.0
+multiprocess==0.70.12.2
+numba==0.56.0
+numpy==1.22.0
+oauth2client==4.1.3
+oauthlib==3.1.1
+opt-einsum==3.3.0
+optax==0.1.3
+packaging==21.3
+pandas==1.3.5
+parso==0.8.3
+pathtools==0.1.2
+pexpect==4.8.0
+pickleshare==0.7.5
+pillow==9.0.0
+pip==22.2.1
+pkg-resources==0.0.0
+pooch==1.6.0
+promise==2.3
+prompt-toolkit==3.0.24
+protobuf==3.19.1
+psutil==5.9.0
+ptyprocess==0.7.0
+pyarrow==6.0.1
+pyasn1-modules==0.2.8
+pyasn1==0.4.8
+pycparser==2.21
+pyctcdecode==0.4.0
+pygments==2.11.1
+pygtrie==2.5.0
+pyparsing==3.0.6
+python-dateutil==2.8.2
+python-levenshtein==0.12.2
+pytz==2021.3
+pyyaml==6.0
+regex==2021.11.10
+requests-oauthlib==1.3.0
+requests==2.27.0
+resampy==0.3.1
+responses==0.18.0
+rich==11.2.0
+rsa==4.8
+sacremoses==0.0.46
+scikit-learn==1.1.1
+scipy==1.7.3
+sentry-sdk==1.5.2
+setuptools==44.0.0
+shortuuid==1.0.8
+six==1.16.0
+smmap==5.0.0
+sortedcontainers==2.4.0
+soundfile==0.10.3.post1
+sox==1.4.1
+subprocess32==3.5.4
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.0
+tensorboard==2.7.0
+tensorflow-cpu==2.7.0
+tensorflow-datasets==4.4.0
+tensorflow-estimator==2.7.0
+tensorflow-io-gcs-filesystem==0.23.1
+tensorflow-metadata==1.5.0
+tensorflow==2.7.0
+tensorstore==0.1.21
+termcolor==1.1.0
+threadpoolctl==3.1.0
+tokenizers==0.11.2
+toolz==0.11.2
+torch==1.12.0
+torchaudio==0.12.0+cpu
+tqdm==4.62.3
+traitlets==5.1.1
+transformers==4.21.0
+typing-extensions==4.3.0
+uritemplate==3.0.1
+urllib3==1.26.7
+wandb==0.12.9
+wcwidth==0.2.5
+werkzeug==2.0.2
+wheel==0.37.1
+wrapt==1.13.3
+xxhash==2.0.2
+yarl==1.7.2
+yaspin==2.1.0
+zipp==3.7.0

wandb/run-20220810_073735-23avj35z/files/wandb-metadata.json

@@ -0,0 +1,70 @@
+{
+    "os": "Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29",
+    "python": "3.8.10",
+    "heartbeatAt": "2022-08-10T07:37:39.012020",
+    "startedAt": "2022-08-10T07:37:35.560272",
+    "docker": null,
+    "cpu_count": 96,
+    "cuda": null,
+    "args": [
+        "--model_name_or_path=./",
+        "--hub_model_id=NbAiLab/wav2vec2-1b-npsc-nst-tpu",
+        "--tokenizer_name=./",
+        "--output_dir=./",
+        "--overwrite_output_dir",
+        "--num_train_epochs=40",
+        "--per_device_train_batch_size=2",
+        "--per_device_eval_batch_size=2",
+        "--gradient_accumulation_steps=1",
+        "--precision=full_mixed",
+        "--matmul_precision=bfloat16",
+        "--multisteps",
+        "--learning_rate=6.394633237505332e-05",
+        "--skip_steps=275000",
+        "--warmup_steps=2000",
+        "--length_column_name=input_length",
+        "--evaluation_strategy=steps",
+        "--text_column_name=text",
+        "--save_steps=5000",
+        "--eval_steps=5000",
+        "--logging_steps=100",
+        "--layerdrop=0.041",
+        "--attention_dropout=0.094",
+        "--activation_dropout=0.055",
+        "--hidden_dropout=0.047",
+        "--save_total_limit=5",
+        "--freeze_feature_encoder",
+        "--feat_proj_dropout=0.04",
+        "--mask_time_prob=0.082",
+        "--mask_time_length=10",
+        "--mask_feature_prob=0.25",
+        "--mask_feature_length=64",
+        "--gradient_checkpointing",
+        "--min_duration_in_seconds=0.5",
+        "--max_duration_in_seconds=30.0",
+        "--use_auth_token",
+        "--seed=42",
+        "--group_by_length",
+        "--do_train",
+        "--do_eval",
+        "--push_to_hub",
+        "--preprocessing_num_workers=32",
+        "--ctc_zero_infinity",
+        "--do_lower_case",
+        "--wandb_project=wav2vec2",
+        "--wandb_name=wav2vec2-1b-npsc-nst-tpu (cont.)",
+        "--remove_punctuation"
+    ],
+    "state": "running",
+    "program": "run_flax_speech_recognition_ctc.py",
+    "codePath": "run_flax_speech_recognition_ctc.py",
+    "git": {
+        "remote": "https://huggingface.co/NbAiLab/wav2vec2-1b-npsc-nst-tpu",
+        "commit": "f624ac4bfedfbf56891676d7c5f2e37b4c8e0745"
+    },
+    "email": "versae@gmail.com",
+    "root": "/data/wav2vec2-1b-npsc-nst-tpu",
+    "host": "t1v-n-eedfb410-w-0",
+    "username": "javierr",
+    "executable": "/data/flax/bin/python"
+}

wandb/run-20220810_073735-23avj35z/files/wandb-summary.json

@@ -0,0 +1 @@
+{"train/grad_norm": 6.5625, "layer_grad_norm/": {"lm_head": {"bias": 0.031982421875, "kernel": 4.625}, "wav2vec2": {"encoder": {"layer_norm": {"bias": 0.0556640625, "scale": 0.06103515625}, "layers": {"0": {"attention": {"k_proj": {"bias": 5.9604644775390625e-05, "kernel": 0.023681640625}, "out_proj": {"bias": 0.04150390625, "kernel": 0.2431640625}, "q_proj": {"bias": 0.002899169921875, "kernel": 0.031005859375}, "v_proj": {"bias": 0.037109375, "kernel": 0.265625}}, "feed_forward": {"intermediate_dense": {"bias": 0.04443359375, "kernel": 0.515625}, "output_dense": {"bias": 0.0159912109375, "kernel": 0.439453125}}, "final_layer_norm": {"bias": 0.146484375, "scale": 0.322265625}, "layer_norm": {"bias": 0.0703125, "scale": 0.07080078125}}, "1": {"attention": {"k_proj": {"bias": 3.4332275390625e-05, "kernel": 0.03955078125}, "out_proj": {"bias": 0.0177001953125, "kernel": 0.134765625}, "q_proj": {"bias": 0.0035247802734375, "kernel": 0.0439453125}, "v_proj": {"bias": 0.02880859375, "kernel": 0.111328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.022216796875, "kernel": 0.3359375}, "output_dense": {"bias": 0.0157470703125, "kernel": 0.259765625}}, "final_layer_norm": {"bias": 0.046142578125, "scale": 0.05712890625}, "layer_norm": {"bias": 0.05712890625, "scale": 0.039794921875}}, "10": {"attention": {"k_proj": {"bias": 3.600120544433594e-05, "kernel": 0.123046875}, "out_proj": {"bias": 0.01416015625, "kernel": 0.2001953125}, "q_proj": {"bias": 0.0078125, "kernel": 0.12255859375}, "v_proj": {"bias": 0.022705078125, "kernel": 0.2001953125}}, "feed_forward": {"intermediate_dense": {"bias": 0.021484375, "kernel": 0.328125}, "output_dense": {"bias": 0.013671875, "kernel": 0.2734375}}, "final_layer_norm": {"bias": 0.0390625, "scale": 0.0322265625}, "layer_norm": {"bias": 0.04931640625, "scale": 0.0341796875}}, "11": {"attention": {"k_proj": {"bias": 8.344650268554688e-05, "kernel": 0.158203125}, "out_proj": {"bias": 0.0142822265625, "kernel": 0.28125}, "q_proj": {"bias": 0.0087890625, "kernel": 0.130859375}, "v_proj": {"bias": 0.024658203125, "kernel": 0.28515625}}, "feed_forward": {"intermediate_dense": {"bias": 0.01953125, "kernel": 0.310546875}, "output_dense": {"bias": 0.013916015625, "kernel": 0.244140625}}, "final_layer_norm": {"bias": 0.03271484375, "scale": 0.0308837890625}, "layer_norm": {"bias": 0.05029296875, "scale": 0.0439453125}}, "12": {"attention": {"k_proj": {"bias": 4.982948303222656e-05, "kernel": 0.123046875}, "out_proj": {"bias": 0.0147705078125, "kernel": 0.244140625}, "q_proj": {"bias": 0.0081787109375, "kernel": 0.1162109375}, "v_proj": {"bias": 0.023681640625, "kernel": 0.2294921875}}, "feed_forward": {"intermediate_dense": {"bias": 0.021484375, "kernel": 0.32421875}, "output_dense": {"bias": 0.014404296875, "kernel": 0.255859375}}, "final_layer_norm": {"bias": 0.0390625, "scale": 0.04248046875}, "layer_norm": {"bias": 0.046630859375, "scale": 0.0546875}}, "13": {"attention": {"k_proj": {"bias": 0.00012493133544921875, "kernel": 0.15625}, "out_proj": {"bias": 0.01519775390625, "kernel": 0.330078125}, "q_proj": {"bias": 0.0111083984375, "kernel": 0.158203125}, "v_proj": {"bias": 0.026611328125, "kernel": 0.3046875}}, "feed_forward": {"intermediate_dense": {"bias": 0.022216796875, "kernel": 0.34375}, "output_dense": {"bias": 0.01513671875, "kernel": 0.3125}}, "final_layer_norm": {"bias": 0.040283203125, "scale": 0.032958984375}, "layer_norm": {"bias": 0.051513671875, "scale": 0.091796875}}, "14": {"attention": {"k_proj": {"bias": 9.441375732421875e-05, "kernel": 0.1005859375}, "out_proj": {"bias": 0.015625, "kernel": 0.2412109375}, "q_proj": {"bias": 0.006256103515625, "kernel": 0.099609375}, "v_proj": {"bias": 0.0235595703125, "kernel": 0.2275390625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0257568359375, "kernel": 0.39453125}, "output_dense": {"bias": 0.015380859375, "kernel": 0.33984375}}, "final_layer_norm": {"bias": 0.05126953125, "scale": 0.05517578125}, "layer_norm": {"bias": 0.041748046875, "scale": 0.03076171875}}, "15": {"attention": {"k_proj": {"bias": 0.0003070831298828125, "kernel": 0.1806640625}, "out_proj": {"bias": 0.015625, "kernel": 0.5078125}, "q_proj": {"bias": 0.0106201171875, "kernel": 0.173828125}, "v_proj": {"bias": 0.026611328125, "kernel": 0.361328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.024169921875, "kernel": 0.376953125}, "output_dense": {"bias": 0.01556396484375, "kernel": 0.349609375}}, "final_layer_norm": {"bias": 0.041015625, "scale": 0.033447265625}, "layer_norm": {"bias": 0.048095703125, "scale": 0.072265625}}, "16": {"attention": {"k_proj": {"bias": 6.389617919921875e-05, "kernel": 0.1025390625}, "out_proj": {"bias": 0.016357421875, "kernel": 0.267578125}, "q_proj": {"bias": 0.0057373046875, "kernel": 0.1005859375}, "v_proj": {"bias": 0.0238037109375, "kernel": 0.220703125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0223388671875, "kernel": 0.359375}, "output_dense": {"bias": 0.0159912109375, "kernel": 0.341796875}}, "final_layer_norm": {"bias": 0.0380859375, "scale": 0.033935546875}, "layer_norm": {"bias": 0.043212890625, "scale": 0.034912109375}}, "17": {"attention": {"k_proj": {"bias": 4.57763671875e-05, "kernel": 0.0927734375}, "out_proj": {"bias": 0.0172119140625, "kernel": 0.23046875}, "q_proj": {"bias": 0.005889892578125, "kernel": 0.087890625}, "v_proj": {"bias": 0.0244140625, "kernel": 0.2177734375}}, "feed_forward": {"intermediate_dense": {"bias": 0.024169921875, "kernel": 0.390625}, "output_dense": {"bias": 0.01708984375, "kernel": 0.353515625}}, "final_layer_norm": {"bias": 0.041259765625, "scale": 0.036376953125}, "layer_norm": {"bias": 0.0439453125, "scale": 0.0341796875}}, "18": {"attention": {"k_proj": {"bias": 0.000247955322265625, "kernel": 0.126953125}, "out_proj": {"bias": 0.017578125, "kernel": 0.369140625}, "q_proj": {"bias": 0.0076904296875, "kernel": 0.1337890625}, "v_proj": {"bias": 0.027587890625, "kernel": 0.298828125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02587890625, "kernel": 0.44921875}, "output_dense": {"bias": 0.0174560546875, "kernel": 0.41015625}}, "final_layer_norm": {"bias": 0.04443359375, "scale": 0.03857421875}, "layer_norm": {"bias": 0.048583984375, "scale": 0.039794921875}}, "19": {"attention": {"k_proj": {"bias": 8.678436279296875e-05, "kernel": 0.140625}, "out_proj": {"bias": 0.017822265625, "kernel": 0.28125}, "q_proj": {"bias": 0.009033203125, "kernel": 0.140625}, "v_proj": {"bias": 0.0286865234375, "kernel": 0.283203125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02587890625, "kernel": 0.474609375}, "output_dense": {"bias": 0.0174560546875, "kernel": 0.421875}}, "final_layer_norm": {"bias": 0.041748046875, "scale": 0.0380859375}, "layer_norm": {"bias": 0.052734375, "scale": 0.04052734375}}, "2": {"attention": {"k_proj": {"bias": 4.982948303222656e-05, "kernel": 0.07421875}, "out_proj": {"bias": 0.0177001953125, "kernel": 0.2060546875}, "q_proj": {"bias": 0.006195068359375, "kernel": 0.06982421875}, "v_proj": {"bias": 0.03173828125, "kernel": 0.181640625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0234375, "kernel": 0.390625}, "output_dense": {"bias": 0.01556396484375, "kernel": 0.275390625}}, "final_layer_norm": {"bias": 0.047119140625, "scale": 0.03173828125}, "layer_norm": {"bias": 0.0556640625, "scale": 0.07275390625}}, "20": {"attention": {"k_proj": {"bias": 2.110004425048828e-05, "kernel": 0.095703125}, "out_proj": {"bias": 0.0185546875, "kernel": 0.142578125}, "q_proj": {"bias": 0.005157470703125, "kernel": 0.0947265625}, "v_proj": {"bias": 0.0263671875, "kernel": 0.140625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0250244140625, "kernel": 0.4765625}, "output_dense": {"bias": 0.018310546875, "kernel": 0.390625}}, "final_layer_norm": {"bias": 0.041015625, "scale": 0.04443359375}, "layer_norm": {"bias": 0.046142578125, "scale": 0.038330078125}}, "21": {"attention": {"k_proj": {"bias": 4.00543212890625e-05, "kernel": 0.1259765625}, "out_proj": {"bias": 0.0189208984375, "kernel": 0.2216796875}, "q_proj": {"bias": 0.006927490234375, "kernel": 0.12890625}, "v_proj": {"bias": 0.02734375, "kernel": 0.203125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0267333984375, "kernel": 0.51953125}, "output_dense": {"bias": 0.0185546875, "kernel": 0.41796875}}, "final_layer_norm": {"bias": 0.04541015625, "scale": 0.04736328125}, "layer_norm": {"bias": 0.044189453125, "scale": 0.054443359375}}, "22": {"attention": {"k_proj": {"bias": 3.3855438232421875e-05, "kernel": 0.1181640625}, "out_proj": {"bias": 0.019775390625, "kernel": 0.240234375}, "q_proj": {"bias": 0.006011962890625, "kernel": 0.11279296875}, "v_proj": {"bias": 0.028076171875, "kernel": 0.21875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0269775390625, "kernel": 0.515625}, "output_dense": {"bias": 0.0194091796875, "kernel": 0.400390625}}, "final_layer_norm": {"bias": 0.046142578125, "scale": 0.047119140625}, "layer_norm": {"bias": 0.049560546875, "scale": 0.0458984375}}, "23": {"attention": {"k_proj": {"bias": 0.0001087188720703125, "kernel": 0.16015625}, "out_proj": {"bias": 0.0198974609375, "kernel": 0.443359375}, "q_proj": {"bias": 0.008544921875, "kernel": 0.1630859375}, "v_proj": {"bias": 0.03173828125, "kernel": 0.35546875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0263671875, "kernel": 0.53125}, "output_dense": {"bias": 0.01953125, "kernel": 0.400390625}}, "final_layer_norm": {"bias": 0.044677734375, "scale": 0.04638671875}, "layer_norm": {"bias": 0.05615234375, "scale": 0.056396484375}}, "24": {"attention": {"k_proj": {"bias": 6.246566772460938e-05, "kernel": 0.169921875}, "out_proj": {"bias": 0.0191650390625, "kernel": 0.36328125}, "q_proj": {"bias": 0.00933837890625, "kernel": 0.18359375}, "v_proj": {"bias": 0.03271484375, "kernel": 0.328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02685546875, "kernel": 0.5390625}, "output_dense": {"bias": 0.01904296875, "kernel": 0.37890625}}, "final_layer_norm": {"bias": 0.04736328125, "scale": 0.04345703125}, "layer_norm": {"bias": 0.0625, "scale": 0.041015625}}, "25": {"attention": {"k_proj": {"bias": 6.079673767089844e-05, "kernel": 0.169921875}, "out_proj": {"bias": 0.0196533203125, "kernel": 0.3125}, "q_proj": {"bias": 0.00860595703125, "kernel": 0.16015625}, "v_proj": {"bias": 0.03271484375, "kernel": 0.32421875}}, "feed_forward": {"intermediate_dense": {"bias": 0.02734375, "kernel": 0.55859375}, "output_dense": {"bias": 0.01953125, "kernel": 0.375}}, "final_layer_norm": {"bias": 0.050537109375, "scale": 0.0478515625}, "layer_norm": {"bias": 0.06005859375, "scale": 0.06298828125}}, "26": {"attention": {"k_proj": {"bias": 5.9604644775390625e-05, "kernel": 0.169921875}, "out_proj": {"bias": 0.01953125, "kernel": 0.29296875}, "q_proj": {"bias": 0.01025390625, "kernel": 0.177734375}, "v_proj": {"bias": 0.0341796875, "kernel": 0.296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.026611328125, "kernel": 0.51171875}, "output_dense": {"bias": 0.01904296875, "kernel": 0.353515625}}, "final_layer_norm": {"bias": 0.0478515625, "scale": 0.04443359375}, "layer_norm": {"bias": 0.060791015625, "scale": 0.06396484375}}, "27": {"attention": {"k_proj": {"bias": 0.00011396408081054688, "kernel": 0.2021484375}, "out_proj": {"bias": 0.01806640625, "kernel": 0.44921875}, "q_proj": {"bias": 0.01068115234375, "kernel": 0.2138671875}, "v_proj": {"bias": 0.03466796875, "kernel": 0.435546875}}, "feed_forward": {"intermediate_dense": {"bias": 0.02734375, "kernel": 0.515625}, "output_dense": {"bias": 0.0181884765625, "kernel": 0.36328125}}, "final_layer_norm": {"bias": 0.05078125, "scale": 0.045654296875}, "layer_norm": {"bias": 0.06640625, "scale": 0.04931640625}}, "28": {"attention": {"k_proj": {"bias": 0.0001049041748046875, "kernel": 0.20703125}, "out_proj": {"bias": 0.0164794921875, "kernel": 0.392578125}, "q_proj": {"bias": 0.01165771484375, "kernel": 0.208984375}, "v_proj": {"bias": 0.031494140625, "kernel": 0.404296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0234375, "kernel": 0.45703125}, "output_dense": {"bias": 0.016357421875, "kernel": 0.326171875}}, "final_layer_norm": {"bias": 0.04248046875, "scale": 0.044921875}, "layer_norm": {"bias": 0.0673828125, "scale": 0.08447265625}}, "29": {"attention": {"k_proj": {"bias": 9.918212890625e-05, "kernel": 0.267578125}, "out_proj": {"bias": 0.0157470703125, "kernel": 0.28515625}, "q_proj": {"bias": 0.01495361328125, "kernel": 0.265625}, "v_proj": {"bias": 0.02978515625, "kernel": 0.31640625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0234375, "kernel": 0.494140625}, "output_dense": {"bias": 0.01531982421875, "kernel": 0.296875}}, "final_layer_norm": {"bias": 0.03955078125, "scale": 0.03515625}, "layer_norm": {"bias": 0.0654296875, "scale": 0.061279296875}}, "3": {"attention": {"k_proj": {"bias": 0.00012111663818359375, "kernel": 0.0986328125}, "out_proj": {"bias": 0.016845703125, "kernel": 0.314453125}, "q_proj": {"bias": 0.00726318359375, "kernel": 0.0888671875}, "v_proj": {"bias": 0.0283203125, "kernel": 0.2470703125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0242919921875, "kernel": 0.3828125}, "output_dense": {"bias": 0.0150146484375, "kernel": 0.271484375}}, "final_layer_norm": {"bias": 0.0458984375, "scale": 0.03125}, "layer_norm": {"bias": 0.0498046875, "scale": 0.0380859375}}, "30": {"attention": {"k_proj": {"bias": 0.0001220703125, "kernel": 0.13671875}, "out_proj": {"bias": 0.0152587890625, "kernel": 0.328125}, "q_proj": {"bias": 0.006378173828125, "kernel": 0.138671875}, "v_proj": {"bias": 0.029296875, "kernel": 0.3671875}}, "feed_forward": {"intermediate_dense": {"bias": 0.023681640625, "kernel": 0.51953125}, "output_dense": {"bias": 0.01446533203125, "kernel": 0.26171875}}, "final_layer_norm": {"bias": 0.04150390625, "scale": 0.03564453125}, "layer_norm": {"bias": 0.04931640625, "scale": 0.037109375}}, "31": {"attention": {"k_proj": {"bias": 0.00010347366333007812, "kernel": 0.14453125}, "out_proj": {"bias": 0.0140380859375, "kernel": 0.29296875}, "q_proj": {"bias": 0.006378173828125, "kernel": 0.134765625}, "v_proj": {"bias": 0.0238037109375, "kernel": 0.314453125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02392578125, "kernel": 0.51953125}, "output_dense": {"bias": 0.01385498046875, "kernel": 0.2578125}}, "final_layer_norm": {"bias": 0.0390625, "scale": 0.03662109375}, "layer_norm": {"bias": 0.039306640625, "scale": 0.0291748046875}}, "32": {"attention": {"k_proj": {"bias": 8.296966552734375e-05, "kernel": 0.15625}, "out_proj": {"bias": 0.01263427734375, "kernel": 0.28125}, "q_proj": {"bias": 0.0079345703125, "kernel": 0.1533203125}, "v_proj": {"bias": 0.0264892578125, "kernel": 0.4921875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0216064453125, "kernel": 0.431640625}, "output_dense": {"bias": 0.01129150390625, "kernel": 0.212890625}}, "final_layer_norm": {"bias": 0.04150390625, "scale": 0.03271484375}, "layer_norm": {"bias": 0.046630859375, "scale": 0.05419921875}}, "33": {"attention": {"k_proj": {"bias": 5.9604644775390625e-05, "kernel": 0.166015625}, "out_proj": {"bias": 0.01092529296875, "kernel": 0.2275390625}, "q_proj": {"bias": 0.008544921875, "kernel": 0.166015625}, "v_proj": {"bias": 0.023193359375, "kernel": 0.34765625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0196533203125, "kernel": 0.390625}, "output_dense": {"bias": 0.00897216796875, "kernel": 0.1875}}, "final_layer_norm": {"bias": 0.04345703125, "scale": 0.0361328125}, "layer_norm": {"bias": 0.039794921875, "scale": 0.0498046875}}, "34": {"attention": {"k_proj": {"bias": 0.0002346038818359375, "kernel": 0.158203125}, "out_proj": {"bias": 0.0081787109375, "kernel": 0.181640625}, "q_proj": {"bias": 0.006927490234375, "kernel": 0.14453125}, "v_proj": {"bias": 0.0177001953125, "kernel": 0.25390625}}, "feed_forward": {"intermediate_dense": {"bias": 0.01434326171875, "kernel": 0.291015625}, "output_dense": {"bias": 0.0072021484375, "kernel": 0.1748046875}}, "final_layer_norm": {"bias": 0.028076171875, "scale": 0.025146484375}, "layer_norm": {"bias": 0.03369140625, "scale": 0.026611328125}}, "35": {"attention": {"k_proj": {"bias": 0.0001506805419921875, "kernel": 0.10791015625}, "out_proj": {"bias": 0.00640869140625, "kernel": 0.2109375}, "q_proj": {"bias": 0.004852294921875, "kernel": 0.10791015625}, "v_proj": {"bias": 0.01177978515625, "kernel": 0.21484375}}, "feed_forward": {"intermediate_dense": {"bias": 0.010498046875, "kernel": 0.2119140625}, "output_dense": {"bias": 0.005889892578125, "kernel": 0.15234375}}, "final_layer_norm": {"bias": 0.0206298828125, "scale": 0.0220947265625}, "layer_norm": {"bias": 0.024169921875, "scale": 0.02880859375}}, "36": {"attention": {"k_proj": {"bias": 4.410743713378906e-05, "kernel": 0.1005859375}, "out_proj": {"bias": 0.005645751953125, "kernel": 0.1552734375}, "q_proj": {"bias": 0.00445556640625, "kernel": 0.095703125}, "v_proj": {"bias": 0.00946044921875, "kernel": 0.14453125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0089111328125, "kernel": 0.177734375}, "output_dense": {"bias": 0.0050048828125, "kernel": 0.111328125}}, "final_layer_norm": {"bias": 0.017578125, "scale": 0.01513671875}, "layer_norm": {"bias": 0.0191650390625, "scale": 0.01806640625}}, "37": {"attention": {"k_proj": {"bias": 9.441375732421875e-05, "kernel": 0.0849609375}, "out_proj": {"bias": 0.004913330078125, "kernel": 0.11474609375}, "q_proj": {"bias": 0.00390625, "kernel": 0.0830078125}, "v_proj": {"bias": 0.00897216796875, "kernel": 0.1318359375}}, "feed_forward": {"intermediate_dense": {"bias": 0.00823974609375, "kernel": 0.16796875}, "output_dense": {"bias": 0.004241943359375, "kernel": 0.09716796875}}, "final_layer_norm": {"bias": 0.015869140625, "scale": 0.01434326171875}, "layer_norm": {"bias": 0.019287109375, "scale": 0.015869140625}}, "38": {"attention": {"k_proj": {"bias": 5.650520324707031e-05, "kernel": 0.09130859375}, "out_proj": {"bias": 0.0040283203125, "kernel": 0.11865234375}, "q_proj": {"bias": 0.00396728515625, "kernel": 0.08642578125}, "v_proj": {"bias": 0.007354736328125, "kernel": 0.1279296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0072021484375, "kernel": 0.150390625}, "output_dense": {"bias": 0.0034637451171875, "kernel": 0.09423828125}}, "final_layer_norm": {"bias": 0.0152587890625, "scale": 0.0146484375}, "layer_norm": {"bias": 0.0162353515625, "scale": 0.0135498046875}}, "39": {"attention": {"k_proj": {"bias": 5.316734313964844e-05, "kernel": 0.09619140625}, "out_proj": {"bias": 0.0030975341796875, "kernel": 0.09619140625}, "q_proj": {"bias": 0.00408935546875, "kernel": 0.0908203125}, "v_proj": {"bias": 0.006011962890625, "kernel": 0.10986328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.005401611328125, "kernel": 0.12109375}, "output_dense": {"bias": 0.0025634765625, "kernel": 0.08642578125}}, "final_layer_norm": {"bias": 0.01202392578125, "scale": 0.01226806640625}, "layer_norm": {"bias": 0.0150146484375, "scale": 0.01556396484375}}, "4": {"attention": {"k_proj": {"bias": 0.000148773193359375, "kernel": 0.10498046875}, "out_proj": {"bias": 0.015869140625, "kernel": 0.361328125}, "q_proj": {"bias": 0.0072021484375, "kernel": 0.1005859375}, "v_proj": {"bias": 0.026123046875, "kernel": 0.3046875}}, "feed_forward": {"intermediate_dense": {"bias": 0.023193359375, "kernel": 0.36328125}, "output_dense": {"bias": 0.014404296875, "kernel": 0.29296875}}, "final_layer_norm": {"bias": 0.042724609375, "scale": 0.034423828125}, "layer_norm": {"bias": 0.0478515625, "scale": 0.060546875}}, "40": {"attention": {"k_proj": {"bias": 5.269050598144531e-05, "kernel": 0.046875}, "out_proj": {"bias": 0.0025787353515625, "kernel": 0.080078125}, "q_proj": {"bias": 0.0020294189453125, "kernel": 0.0458984375}, "v_proj": {"bias": 0.004150390625, "kernel": 0.07080078125}}, "feed_forward": {"intermediate_dense": {"bias": 0.004302978515625, "kernel": 0.09326171875}, "output_dense": {"bias": 0.0023040771484375, "kernel": 0.060791015625}}, "final_layer_norm": {"bias": 0.0087890625, "scale": 0.011474609375}, "layer_norm": {"bias": 0.00823974609375, "scale": 0.007781982421875}}, "41": {"attention": {"k_proj": {"bias": 4.0531158447265625e-05, "kernel": 0.0673828125}, "out_proj": {"bias": 0.002044677734375, "kernel": 0.087890625}, "q_proj": {"bias": 0.0025787353515625, "kernel": 0.0634765625}, "v_proj": {"bias": 0.00439453125, "kernel": 0.1044921875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0035858154296875, "kernel": 0.091796875}, "output_dense": {"bias": 0.00168609619140625, "kernel": 0.064453125}}, "final_layer_norm": {"bias": 0.00921630859375, "scale": 0.0106201171875}, "layer_norm": {"bias": 0.010986328125, "scale": 0.0128173828125}}, "42": {"attention": {"k_proj": {"bias": 1.1801719665527344e-05, "kernel": 0.02099609375}, "out_proj": {"bias": 0.001678466796875, "kernel": 0.048828125}, "q_proj": {"bias": 0.0009307861328125, "kernel": 0.02197265625}, "v_proj": {"bias": 0.002349853515625, "kernel": 0.0478515625}}, "feed_forward": {"intermediate_dense": {"bias": 0.002685546875, "kernel": 0.07421875}, "output_dense": {"bias": 0.0014801025390625, "kernel": 0.053955078125}}, "final_layer_norm": {"bias": 0.0059814453125, "scale": 0.00885009765625}, "layer_norm": {"bias": 0.0045166015625, "scale": 0.00439453125}}, "43": {"attention": {"k_proj": {"bias": 8.046627044677734e-06, "kernel": 0.01806640625}, "out_proj": {"bias": 0.0015106201171875, "kernel": 0.035888671875}, "q_proj": {"bias": 0.00095367431640625, "kernel": 0.0208740234375}, "v_proj": {"bias": 0.00171661376953125, "kernel": 0.031005859375}}, "feed_forward": {"intermediate_dense": {"bias": 0.002777099609375, "kernel": 0.083984375}, "output_dense": {"bias": 0.00125885009765625, "kernel": 0.052734375}}, "final_layer_norm": {"bias": 0.00689697265625, "scale": 0.007110595703125}, "layer_norm": {"bias": 0.0036468505859375, "scale": 0.005706787109375}}, "44": {"attention": {"k_proj": {"bias": 1.3113021850585938e-05, "kernel": 0.023681640625}, "out_proj": {"bias": 0.001312255859375, "kernel": 0.033447265625}, "q_proj": {"bias": 0.000946044921875, "kernel": 0.021484375}, "v_proj": {"bias": 0.0017547607421875, "kernel": 0.03515625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0023956298828125, "kernel": 0.0771484375}, "output_dense": {"bias": 0.001129150390625, "kernel": 0.052001953125}}, "final_layer_norm": {"bias": 0.005706787109375, "scale": 0.005859375}, "layer_norm": {"bias": 0.0042724609375, "scale": 0.004730224609375}}, "45": {"attention": {"k_proj": {"bias": 1.4424324035644531e-05, "kernel": 0.01239013671875}, "out_proj": {"bias": 0.00110626220703125, "kernel": 0.0267333984375}, "q_proj": {"bias": 0.00136566162109375, "kernel": 0.029541015625}, "v_proj": {"bias": 0.00142669677734375, "kernel": 0.027587890625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0019989013671875, "kernel": 0.06396484375}, "output_dense": {"bias": 0.000881195068359375, "kernel": 0.05615234375}}, "final_layer_norm": {"bias": 0.006072998046875, "scale": 0.006591796875}, "layer_norm": {"bias": 0.004791259765625, "scale": 0.00494384765625}}, "46": {"attention": {"k_proj": {"bias": 5.745887756347656e-05, "kernel": 0.006439208984375}, "out_proj": {"bias": 0.000888824462890625, "kernel": 0.0289306640625}, "q_proj": {"bias": 0.000591278076171875, "kernel": 0.011962890625}, "v_proj": {"bias": 0.0010986328125, "kernel": 0.0233154296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.00141143798828125, "kernel": 0.03955078125}, "output_dense": {"bias": 0.000873565673828125, "kernel": 0.0478515625}}, "final_layer_norm": {"bias": 0.00433349609375, "scale": 0.00433349609375}, "layer_norm": {"bias": 0.0036468505859375, "scale": 0.003814697265625}}, "47": {"attention": {"k_proj": {"bias": 0.00011301040649414062, "kernel": 0.003997802734375}, "out_proj": {"bias": 0.000896453857421875, "kernel": 0.06640625}, "q_proj": {"bias": 0.00014591217041015625, "kernel": 0.00286865234375}, "v_proj": {"bias": 0.00118255615234375, "kernel": 0.0230712890625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0010528564453125, "kernel": 0.0252685546875}, "output_dense": {"bias": 0.000881195068359375, "kernel": 0.1787109375}}, "final_layer_norm": {"bias": 0.005950927734375, "scale": 0.00677490234375}, "layer_norm": {"bias": 0.005859375, "scale": 0.005767822265625}}, "5": {"attention": {"k_proj": {"bias": 6.4849853515625e-05, "kernel": 0.123046875}, "out_proj": {"bias": 0.0159912109375, "kernel": 0.203125}, "q_proj": {"bias": 0.007598876953125, "kernel": 0.12158203125}, "v_proj": {"bias": 0.02685546875, "kernel": 0.1953125}}, "feed_forward": {"intermediate_dense": {"bias": 0.023193359375, "kernel": 0.33984375}, "output_dense": {"bias": 0.0147705078125, "kernel": 0.275390625}}, "final_layer_norm": {"bias": 0.041015625, "scale": 0.037841796875}, "layer_norm": {"bias": 0.052734375, "scale": 0.04736328125}}, "6": {"attention": {"k_proj": {"bias": 7.152557373046875e-05, "kernel": 0.1318359375}, "out_proj": {"bias": 0.0152587890625, "kernel": 0.349609375}, "q_proj": {"bias": 0.00823974609375, "kernel": 0.119140625}, "v_proj": {"bias": 0.02685546875, "kernel": 0.31640625}}, "feed_forward": {"intermediate_dense": {"bias": 0.022216796875, "kernel": 0.341796875}, "output_dense": {"bias": 0.014404296875, "kernel": 0.271484375}}, "final_layer_norm": {"bias": 0.0380859375, "scale": 0.0322265625}, "layer_norm": {"bias": 0.049560546875, "scale": 0.058349609375}}, "7": {"attention": {"k_proj": {"bias": 7.62939453125e-05, "kernel": 0.1328125}, "out_proj": {"bias": 0.0150146484375, "kernel": 0.349609375}, "q_proj": {"bias": 0.00921630859375, "kernel": 0.126953125}, "v_proj": {"bias": 0.0255126953125, "kernel": 0.30859375}}, "feed_forward": {"intermediate_dense": {"bias": 0.021484375, "kernel": 0.33984375}, "output_dense": {"bias": 0.01409912109375, "kernel": 0.26171875}}, "final_layer_norm": {"bias": 0.038818359375, "scale": 0.033935546875}, "layer_norm": {"bias": 0.05126953125, "scale": 0.050048828125}}, "8": {"attention": {"k_proj": {"bias": 7.915496826171875e-05, "kernel": 0.12060546875}, "out_proj": {"bias": 0.01507568359375, "kernel": 0.302734375}, "q_proj": {"bias": 0.007568359375, "kernel": 0.11474609375}, "v_proj": {"bias": 0.0262451171875, "kernel": 0.27734375}}, "feed_forward": {"intermediate_dense": {"bias": 0.023193359375, "kernel": 0.361328125}, "output_dense": {"bias": 0.01409912109375, "kernel": 0.275390625}}, "final_layer_norm": {"bias": 0.044677734375, "scale": 0.033447265625}, "layer_norm": {"bias": 0.0498046875, "scale": 0.06103515625}}, "9": {"attention": {"k_proj": {"bias": 0.00011777877807617188, "kernel": 0.1513671875}, "out_proj": {"bias": 0.0140380859375, "kernel": 0.416015625}, "q_proj": {"bias": 0.00830078125, "kernel": 0.1376953125}, "v_proj": {"bias": 0.0238037109375, "kernel": 0.40234375}}, "feed_forward": {"intermediate_dense": {"bias": 0.0213623046875, "kernel": 0.3515625}, "output_dense": {"bias": 0.0135498046875, "kernel": 0.279296875}}, "final_layer_norm": {"bias": 0.037109375, "scale": 0.0322265625}, "layer_norm": {"bias": 0.04443359375, "scale": 0.04833984375}}}, "pos_conv_embed": {"conv": {"bias": 0.034912109375, "weight_g": 0.044189453125, "weight_v": 0.287109375}}}, "feature_extractor": {"conv_layers": {"0": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "1": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "2": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "3": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "4": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "5": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "6": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}}}, "feature_projection": {"layer_norm": {"bias": 0.1337890625, "scale": 0.1611328125}, "projection": {"bias": 0.0556640625, "kernel": 1.0546875}}, "masked_spec_embed": 0.0}}, "layer_param_norm/": {"lm_head": {"bias": 0.7824921607971191, "kernel": 55.72966766357422}, "wav2vec2": {"encoder": {"layer_norm": {"bias": 59.6768798828125, "scale": 74.17054748535156}, "layers": {"0": {"attention": {"k_proj": {"bias": 0.37033993005752563, "kernel": 27.536663055419922}, "out_proj": {"bias": 1.6469175815582275, "kernel": 26.147050857543945}, "q_proj": {"bias": 1.5330281257629395, "kernel": 27.813282012939453}, "v_proj": {"bias": 0.44783300161361694, "kernel": 26.55841064453125}}, "feed_forward": {"intermediate_dense": {"bias": 1.9835489988327026, "kernel": 100.66567993164062}, "output_dense": {"bias": 1.116748571395874, "kernel": 96.76679992675781}}, "final_layer_norm": {"bias": 1.335214376449585, "scale": 19.85782241821289}, "layer_norm": {"bias": 2.923041343688965, "scale": 15.398418426513672}}, "1": {"attention": {"k_proj": {"bias": 0.3767347037792206, "kernel": 41.013240814208984}, "out_proj": {"bias": 1.3653483390808105, "kernel": 43.371070861816406}, "q_proj": {"bias": 3.0925614833831787, "kernel": 41.05661392211914}, "v_proj": {"bias": 0.2924947738647461, "kernel": 41.61189270019531}}, "feed_forward": {"intermediate_dense": {"bias": 1.9882662296295166, "kernel": 98.7442626953125}, "output_dense": {"bias": 0.8527815341949463, "kernel": 87.83541870117188}}, "final_layer_norm": {"bias": 1.3605934381484985, "scale": 19.084806442260742}, "layer_norm": {"bias": 1.9318764209747314, "scale": 17.761367797851562}}, "10": {"attention": {"k_proj": {"bias": 0.4123449921607971, "kernel": 49.44670486450195}, "out_proj": {"bias": 1.3130683898925781, "kernel": 52.27025604248047}, "q_proj": {"bias": 2.48445200920105, "kernel": 49.528873443603516}, "v_proj": {"bias": 0.3416975736618042, "kernel": 52.344085693359375}}, "feed_forward": {"intermediate_dense": {"bias": 1.974550485610962, "kernel": 102.70410919189453}, "output_dense": {"bias": 0.5955485105514526, "kernel": 95.81275939941406}}, "final_layer_norm": {"bias": 2.3762400150299072, "scale": 20.81279754638672}, "layer_norm": {"bias": 1.806241512298584, "scale": 21.429487228393555}}, "11": {"attention": {"k_proj": {"bias": 0.4460787773132324, "kernel": 49.37338638305664}, "out_proj": {"bias": 1.1512463092803955, "kernel": 51.949554443359375}, "q_proj": {"bias": 2.5446064472198486, "kernel": 49.20353698730469}, "v_proj": {"bias": 0.40995872020721436, "kernel": 52.2607536315918}}, "feed_forward": {"intermediate_dense": {"bias": 2.019528388977051, "kernel": 103.56025695800781}, "output_dense": {"bias": 0.5711302757263184, "kernel": 97.53776550292969}}, "final_layer_norm": {"bias": 2.3660812377929688, "scale": 20.919677734375}, "layer_norm": {"bias": 1.7802445888519287, "scale": 22.01519203186035}}, "12": {"attention": {"k_proj": {"bias": 0.4312320351600647, "kernel": 50.07032775878906}, "out_proj": {"bias": 1.126122236251831, "kernel": 51.988826751708984}, "q_proj": {"bias": 2.4090728759765625, "kernel": 49.92729949951172}, "v_proj": {"bias": 0.4024103879928589, "kernel": 52.296756744384766}}, "feed_forward": {"intermediate_dense": {"bias": 2.0548508167266846, "kernel": 104.54823303222656}, "output_dense": {"bias": 0.5548778772354126, "kernel": 99.2693099975586}}, "final_layer_norm": {"bias": 2.2933573722839355, "scale": 20.85626983642578}, "layer_norm": {"bias": 1.8587299585342407, "scale": 22.473487854003906}}, "13": {"attention": {"k_proj": {"bias": 0.4430793821811676, "kernel": 51.76431655883789}, "out_proj": {"bias": 1.1271920204162598, "kernel": 51.86852264404297}, "q_proj": {"bias": 2.359200954437256, "kernel": 51.75225830078125}, "v_proj": {"bias": 0.3906242251396179, "kernel": 51.92781066894531}}, "feed_forward": {"intermediate_dense": {"bias": 2.0941619873046875, "kernel": 105.30806732177734}, "output_dense": {"bias": 0.5719542503356934, "kernel": 99.8712158203125}}, "final_layer_norm": {"bias": 2.2314066886901855, "scale": 21.027400970458984}, "layer_norm": {"bias": 1.9997800588607788, "scale": 22.84510040283203}}, "14": {"attention": {"k_proj": {"bias": 0.43604815006256104, "kernel": 51.92181396484375}, "out_proj": {"bias": 1.2681217193603516, "kernel": 49.762760162353516}, "q_proj": {"bias": 2.4942922592163086, "kernel": 52.049808502197266}, "v_proj": {"bias": 0.36820662021636963, "kernel": 49.26283264160156}}, "feed_forward": {"intermediate_dense": {"bias": 2.1335830688476562, "kernel": 105.94457244873047}, "output_dense": {"bias": 0.6063626408576965, "kernel": 101.24859619140625}}, "final_layer_norm": {"bias": 2.2754664421081543, "scale": 21.145992279052734}, "layer_norm": {"bias": 2.1295526027679443, "scale": 22.584672927856445}}, "15": {"attention": {"k_proj": {"bias": 0.45931702852249146, "kernel": 51.93058776855469}, "out_proj": {"bias": 1.3753983974456787, "kernel": 50.94449234008789}, "q_proj": {"bias": 2.5871663093566895, "kernel": 52.099769592285156}, "v_proj": {"bias": 0.459650456905365, "kernel": 50.5831298828125}}, "feed_forward": {"intermediate_dense": {"bias": 2.132938861846924, "kernel": 105.57443237304688}, "output_dense": {"bias": 0.7701732516288757, "kernel": 101.94094848632812}}, "final_layer_norm": {"bias": 2.327320098876953, "scale": 21.192947387695312}, "layer_norm": {"bias": 2.4148712158203125, "scale": 23.526634216308594}}, "16": {"attention": {"k_proj": {"bias": 0.4008745551109314, "kernel": 51.772621154785156}, "out_proj": {"bias": 1.27531099319458, "kernel": 50.134521484375}, "q_proj": {"bias": 2.667466163635254, "kernel": 51.75814437866211}, "v_proj": {"bias": 0.3768249750137329, "kernel": 49.78093719482422}}, "feed_forward": {"intermediate_dense": {"bias": 2.1094985008239746, "kernel": 106.10227966308594}, "output_dense": {"bias": 0.7860437631607056, "kernel": 102.66590881347656}}, "final_layer_norm": {"bias": 2.337951421737671, "scale": 21.583194732666016}, "layer_norm": {"bias": 2.283249855041504, "scale": 22.168060302734375}}, "17": {"attention": {"k_proj": {"bias": 0.3966267704963684, "kernel": 51.728885650634766}, "out_proj": {"bias": 1.2151354551315308, "kernel": 49.4556884765625}, "q_proj": {"bias": 2.714320182800293, "kernel": 51.81880187988281}, "v_proj": {"bias": 0.42661017179489136, "kernel": 49.11927032470703}}, "feed_forward": {"intermediate_dense": {"bias": 2.1101765632629395, "kernel": 107.14872741699219}, "output_dense": {"bias": 0.8218655586242676, "kernel": 103.06423950195312}}, "final_layer_norm": {"bias": 2.383938789367676, "scale": 22.070323944091797}, "layer_norm": {"bias": 2.222898483276367, "scale": 21.219982147216797}}, "18": {"attention": {"k_proj": {"bias": 0.4409676194190979, "kernel": 52.41611099243164}, "out_proj": {"bias": 1.3447906970977783, "kernel": 50.491905212402344}, "q_proj": {"bias": 2.614685535430908, "kernel": 52.796600341796875}, "v_proj": {"bias": 0.4518332779407501, "kernel": 50.001895904541016}}, "feed_forward": {"intermediate_dense": {"bias": 2.144195556640625, "kernel": 107.41338348388672}, "output_dense": {"bias": 0.9481453895568848, "kernel": 104.72514343261719}}, "final_layer_norm": {"bias": 2.5390124320983887, "scale": 22.15178680419922}, "layer_norm": {"bias": 2.424910068511963, "scale": 23.585906982421875}}, "19": {"attention": {"k_proj": {"bias": 0.38193291425704956, "kernel": 51.511146545410156}, "out_proj": {"bias": 1.3303101062774658, "kernel": 50.10035705566406}, "q_proj": {"bias": 2.930327892303467, "kernel": 51.865638732910156}, "v_proj": {"bias": 0.4086824655532837, "kernel": 49.38078308105469}}, "feed_forward": {"intermediate_dense": {"bias": 2.1912901401519775, "kernel": 107.95254516601562}, "output_dense": {"bias": 1.0248571634292603, "kernel": 105.65098571777344}}, "final_layer_norm": {"bias": 2.4923481941223145, "scale": 22.505674362182617}, "layer_norm": {"bias": 2.2888314723968506, "scale": 22.31826400756836}}, "2": {"attention": {"k_proj": {"bias": 0.454792320728302, "kernel": 47.77275085449219}, "out_proj": {"bias": 1.256988525390625, "kernel": 45.969764709472656}, "q_proj": {"bias": 3.2510807514190674, "kernel": 47.61664581298828}, "v_proj": {"bias": 0.339598685503006, "kernel": 45.72273254394531}}, "feed_forward": {"intermediate_dense": {"bias": 1.9737317562103271, "kernel": 103.32754516601562}, "output_dense": {"bias": 0.7398276329040527, "kernel": 91.11263275146484}}, "final_layer_norm": {"bias": 1.5421981811523438, "scale": 21.561111450195312}, "layer_norm": {"bias": 1.7081801891326904, "scale": 20.852447509765625}}, "20": {"attention": {"k_proj": {"bias": 0.4067543148994446, "kernel": 51.605438232421875}, "out_proj": {"bias": 1.359946370124817, "kernel": 49.45553207397461}, "q_proj": {"bias": 2.8498687744140625, "kernel": 52.224571228027344}, "v_proj": {"bias": 0.36227869987487793, "kernel": 48.43864822387695}}, "feed_forward": {"intermediate_dense": {"bias": 2.1725549697875977, "kernel": 109.17405700683594}, "output_dense": {"bias": 1.1388803720474243, "kernel": 106.40528106689453}}, "final_layer_norm": {"bias": 2.435314655303955, "scale": 23.4317626953125}, "layer_norm": {"bias": 2.231672525405884, "scale": 22.230525970458984}}, "21": {"attention": {"k_proj": {"bias": 0.4161534905433655, "kernel": 51.942527770996094}, "out_proj": {"bias": 1.403618335723877, "kernel": 49.51059341430664}, "q_proj": {"bias": 2.7690629959106445, "kernel": 52.67078399658203}, "v_proj": {"bias": 0.41060006618499756, "kernel": 48.64883041381836}}, "feed_forward": {"intermediate_dense": {"bias": 2.2174296379089355, "kernel": 109.5155029296875}, "output_dense": {"bias": 1.253208041191101, "kernel": 106.88243865966797}}, "final_layer_norm": {"bias": 2.4632763862609863, "scale": 23.175764083862305}, "layer_norm": {"bias": 2.2785892486572266, "scale": 22.234222412109375}}, "22": {"attention": {"k_proj": {"bias": 0.45357397198677063, "kernel": 52.54576110839844}, "out_proj": {"bias": 1.349219560623169, "kernel": 49.533172607421875}, "q_proj": {"bias": 2.8105549812316895, "kernel": 52.86981201171875}, "v_proj": {"bias": 0.3973655700683594, "kernel": 49.33363342285156}}, "feed_forward": {"intermediate_dense": {"bias": 2.1619315147399902, "kernel": 109.95498657226562}, "output_dense": {"bias": 1.3076066970825195, "kernel": 106.3852310180664}}, "final_layer_norm": {"bias": 2.3642821311950684, "scale": 22.684059143066406}, "layer_norm": {"bias": 2.3316237926483154, "scale": 21.545879364013672}}, "23": {"attention": {"k_proj": {"bias": 0.4928613007068634, "kernel": 53.47669219970703}, "out_proj": {"bias": 1.564335823059082, "kernel": 50.98707580566406}, "q_proj": {"bias": 2.7065773010253906, "kernel": 53.582611083984375}, "v_proj": {"bias": 0.5810648202896118, "kernel": 51.54853820800781}}, "feed_forward": {"intermediate_dense": {"bias": 2.131969690322876, "kernel": 109.86410522460938}, "output_dense": {"bias": 1.2769315242767334, "kernel": 107.37890625}}, "final_layer_norm": {"bias": 2.767916679382324, "scale": 22.887813568115234}, "layer_norm": {"bias": 2.824352264404297, "scale": 23.373172760009766}}, "24": {"attention": {"k_proj": {"bias": 0.46056002378463745, "kernel": 52.424072265625}, "out_proj": {"bias": 1.6070430278778076, "kernel": 52.50334167480469}, "q_proj": {"bias": 2.828113079071045, "kernel": 52.40515899658203}, "v_proj": {"bias": 0.5424190163612366, "kernel": 52.51116180419922}}, "feed_forward": {"intermediate_dense": {"bias": 2.2367913722991943, "kernel": 109.35035705566406}, "output_dense": {"bias": 1.3016372919082642, "kernel": 110.30095672607422}}, "final_layer_norm": {"bias": 2.83841872215271, "scale": 22.964658737182617}, "layer_norm": {"bias": 2.56215763092041, "scale": 22.983924865722656}}, "25": {"attention": {"k_proj": {"bias": 0.42509031295776367, "kernel": 52.730464935302734}, "out_proj": {"bias": 1.363797664642334, "kernel": 50.5806884765625}, "q_proj": {"bias": 2.9342763423919678, "kernel": 52.548744201660156}, "v_proj": {"bias": 0.6404213309288025, "kernel": 51.0885009765625}}, "feed_forward": {"intermediate_dense": {"bias": 2.1367578506469727, "kernel": 109.70021057128906}, "output_dense": {"bias": 1.1017413139343262, "kernel": 110.27072143554688}}, "final_layer_norm": {"bias": 2.5763301849365234, "scale": 23.494670867919922}, "layer_norm": {"bias": 2.683134078979492, "scale": 21.88357925415039}}, "26": {"attention": {"k_proj": {"bias": 0.4836847186088562, "kernel": 53.01764678955078}, "out_proj": {"bias": 1.2433912754058838, "kernel": 51.37077331542969}, "q_proj": {"bias": 2.943906784057617, "kernel": 52.80891036987305}, "v_proj": {"bias": 0.5064959526062012, "kernel": 52.004638671875}}, "feed_forward": {"intermediate_dense": {"bias": 2.2763516902923584, "kernel": 109.44652557373047}, "output_dense": {"bias": 1.0912110805511475, "kernel": 107.40899658203125}}, "final_layer_norm": {"bias": 2.1937994956970215, "scale": 22.433353424072266}, "layer_norm": {"bias": 2.497119903564453, "scale": 22.19057273864746}}, "27": {"attention": {"k_proj": {"bias": 0.5808594226837158, "kernel": 53.76898956298828}, "out_proj": {"bias": 1.5447406768798828, "kernel": 52.95805358886719}, "q_proj": {"bias": 2.703345775604248, "kernel": 53.69578552246094}, "v_proj": {"bias": 0.6748642325401306, "kernel": 53.388118743896484}}, "feed_forward": {"intermediate_dense": {"bias": 2.404933452606201, "kernel": 107.8713150024414}, "output_dense": {"bias": 0.9485896825790405, "kernel": 107.17198181152344}}, "final_layer_norm": {"bias": 2.5252954959869385, "scale": 21.88959503173828}, "layer_norm": {"bias": 2.6147172451019287, "scale": 23.32440948486328}}, "28": {"attention": {"k_proj": {"bias": 0.5901432037353516, "kernel": 54.482521057128906}, "out_proj": {"bias": 1.5367379188537598, "kernel": 53.31493377685547}, "q_proj": {"bias": 2.9472482204437256, "kernel": 54.1741943359375}, "v_proj": {"bias": 0.5131911039352417, "kernel": 53.759761810302734}}, "feed_forward": {"intermediate_dense": {"bias": 2.3475265502929688, "kernel": 107.87416076660156}, "output_dense": {"bias": 0.8224154710769653, "kernel": 109.1680908203125}}, "final_layer_norm": {"bias": 2.425306797027588, "scale": 22.337677001953125}, "layer_norm": {"bias": 2.0914058685302734, "scale": 23.993711471557617}}, "29": {"attention": {"k_proj": {"bias": 0.46781182289123535, "kernel": 51.12034606933594}, "out_proj": {"bias": 1.5021522045135498, "kernel": 55.685630798339844}, "q_proj": {"bias": 2.809702157974243, "kernel": 51.00274658203125}, "v_proj": {"bias": 0.4760415554046631, "kernel": 55.703304290771484}}, "feed_forward": {"intermediate_dense": {"bias": 2.297222137451172, "kernel": 108.01033020019531}, "output_dense": {"bias": 0.9597339630126953, "kernel": 113.12825012207031}}, "final_layer_norm": {"bias": 2.5980498790740967, "scale": 23.459980010986328}, "layer_norm": {"bias": 2.245180130004883, "scale": 25.39927864074707}}, "3": {"attention": {"k_proj": {"bias": 0.45006245374679565, "kernel": 52.03215789794922}, "out_proj": {"bias": 1.4254932403564453, "kernel": 48.60858917236328}, "q_proj": {"bias": 2.8560738563537598, "kernel": 52.312644958496094}, "v_proj": {"bias": 0.3246268630027771, "kernel": 48.768699645996094}}, "feed_forward": {"intermediate_dense": {"bias": 1.9663825035095215, "kernel": 104.83622741699219}, "output_dense": {"bias": 0.6984099745750427, "kernel": 94.07957458496094}}, "final_layer_norm": {"bias": 1.8095453977584839, "scale": 21.664737701416016}, "layer_norm": {"bias": 1.9017157554626465, "scale": 22.739452362060547}}, "30": {"attention": {"k_proj": {"bias": 0.5024805665016174, "kernel": 52.825706481933594}, "out_proj": {"bias": 1.3023658990859985, "kernel": 52.053871154785156}, "q_proj": {"bias": 2.907101631164551, "kernel": 52.91836166381836}, "v_proj": {"bias": 0.49308842420578003, "kernel": 52.49382019042969}}, "feed_forward": {"intermediate_dense": {"bias": 2.2399911880493164, "kernel": 108.17861938476562}, "output_dense": {"bias": 0.9140658378601074, "kernel": 112.09104919433594}}, "final_layer_norm": {"bias": 2.4926414489746094, "scale": 24.492368698120117}, "layer_norm": {"bias": 2.316732168197632, "scale": 24.931156158447266}}, "31": {"attention": {"k_proj": {"bias": 0.5412741899490356, "kernel": 51.240806579589844}, "out_proj": {"bias": 1.2333163022994995, "kernel": 52.19988250732422}, "q_proj": {"bias": 2.6581294536590576, "kernel": 51.346168518066406}, "v_proj": {"bias": 0.5469827651977539, "kernel": 52.432586669921875}}, "feed_forward": {"intermediate_dense": {"bias": 2.3097352981567383, "kernel": 106.72758483886719}, "output_dense": {"bias": 1.0891624689102173, "kernel": 109.24717712402344}}, "final_layer_norm": {"bias": 2.2962756156921387, "scale": 24.31252670288086}, "layer_norm": {"bias": 2.3430848121643066, "scale": 24.590187072753906}}, "32": {"attention": {"k_proj": {"bias": 0.4704548716545105, "kernel": 50.39933776855469}, "out_proj": {"bias": 1.2453913688659668, "kernel": 51.57465362548828}, "q_proj": {"bias": 2.8450098037719727, "kernel": 50.34847640991211}, "v_proj": {"bias": 0.419519305229187, "kernel": 51.972320556640625}}, "feed_forward": {"intermediate_dense": {"bias": 2.2569355964660645, "kernel": 105.33018493652344}, "output_dense": {"bias": 1.146787166595459, "kernel": 108.35574340820312}}, "final_layer_norm": {"bias": 2.31538724899292, "scale": 24.518985748291016}, "layer_norm": {"bias": 2.417579174041748, "scale": 24.991830825805664}}, "33": {"attention": {"k_proj": {"bias": 0.48390907049179077, "kernel": 50.28266143798828}, "out_proj": {"bias": 1.280959129333496, "kernel": 51.30557632446289}, "q_proj": {"bias": 2.998173713684082, "kernel": 50.253868103027344}, "v_proj": {"bias": 0.4416005611419678, "kernel": 51.71035385131836}}, "feed_forward": {"intermediate_dense": {"bias": 2.279946804046631, "kernel": 103.67684173583984}, "output_dense": {"bias": 1.1754591464996338, "kernel": 106.81501007080078}}, "final_layer_norm": {"bias": 2.257889747619629, "scale": 24.207740783691406}, "layer_norm": {"bias": 2.5865607261657715, "scale": 25.067882537841797}}, "34": {"attention": {"k_proj": {"bias": 0.45390456914901733, "kernel": 49.26523208618164}, "out_proj": {"bias": 1.5265402793884277, "kernel": 52.46200942993164}, "q_proj": {"bias": 2.913527011871338, "kernel": 49.268951416015625}, "v_proj": {"bias": 0.4023621678352356, "kernel": 52.534793853759766}}, "feed_forward": {"intermediate_dense": {"bias": 2.3745017051696777, "kernel": 102.21138000488281}, "output_dense": {"bias": 1.1227428913116455, "kernel": 105.72161102294922}}, "final_layer_norm": {"bias": 2.20273494720459, "scale": 23.640857696533203}, "layer_norm": {"bias": 2.615731716156006, "scale": 25.498104095458984}}, "35": {"attention": {"k_proj": {"bias": 0.5336894989013672, "kernel": 51.04902648925781}, "out_proj": {"bias": 1.4900906085968018, "kernel": 51.15006637573242}, "q_proj": {"bias": 2.573650360107422, "kernel": 51.323951721191406}, "v_proj": {"bias": 0.4889468550682068, "kernel": 51.176795959472656}}, "feed_forward": {"intermediate_dense": {"bias": 2.502547264099121, "kernel": 100.7535400390625}, "output_dense": {"bias": 1.0254027843475342, "kernel": 104.25007629394531}}, "final_layer_norm": {"bias": 2.295243263244629, "scale": 23.61981964111328}, "layer_norm": {"bias": 2.499863862991333, "scale": 26.093618392944336}}, "36": {"attention": {"k_proj": {"bias": 0.4473738670349121, "kernel": 48.319740295410156}, "out_proj": {"bias": 1.5177178382873535, "kernel": 52.26478958129883}, "q_proj": {"bias": 2.6124308109283447, "kernel": 48.236427307128906}, "v_proj": {"bias": 0.39316776394844055, "kernel": 52.66499328613281}}, "feed_forward": {"intermediate_dense": {"bias": 2.3669564723968506, "kernel": 99.60517120361328}, "output_dense": {"bias": 1.0260361433029175, "kernel": 103.68067932128906}}, "final_layer_norm": {"bias": 2.044203519821167, "scale": 24.13540267944336}, "layer_norm": {"bias": 2.2894434928894043, "scale": 25.63806915283203}}, "37": {"attention": {"k_proj": {"bias": 0.6240901350975037, "kernel": 47.300960540771484}, "out_proj": {"bias": 1.7604304552078247, "kernel": 52.189971923828125}, "q_proj": {"bias": 2.3819239139556885, "kernel": 47.316253662109375}, "v_proj": {"bias": 0.38518577814102173, "kernel": 52.32656478881836}}, "feed_forward": {"intermediate_dense": {"bias": 2.2694783210754395, "kernel": 98.56558227539062}, "output_dense": {"bias": 1.0087021589279175, "kernel": 103.13001251220703}}, "final_layer_norm": {"bias": 1.7903382778167725, "scale": 24.500085830688477}, "layer_norm": {"bias": 2.238807439804077, "scale": 25.563133239746094}}, "38": {"attention": {"k_proj": {"bias": 0.722891092300415, "kernel": 45.45281219482422}, "out_proj": {"bias": 1.4463956356048584, "kernel": 51.495086669921875}, "q_proj": {"bias": 2.2622861862182617, "kernel": 45.454437255859375}, "v_proj": {"bias": 0.42936205863952637, "kernel": 51.565425872802734}}, "feed_forward": {"intermediate_dense": {"bias": 2.2012171745300293, "kernel": 96.4342041015625}, "output_dense": {"bias": 0.9817801713943481, "kernel": 101.30387115478516}}, "final_layer_norm": {"bias": 1.7825044393539429, "scale": 25.214679718017578}, "layer_norm": {"bias": 2.4084482192993164, "scale": 26.425636291503906}}, "39": {"attention": {"k_proj": {"bias": 0.7209377884864807, "kernel": 45.24742889404297}, "out_proj": {"bias": 1.7097458839416504, "kernel": 51.329002380371094}, "q_proj": {"bias": 2.1152358055114746, "kernel": 45.53700637817383}, "v_proj": {"bias": 0.4246286153793335, "kernel": 51.27728271484375}}, "feed_forward": {"intermediate_dense": {"bias": 2.177624225616455, "kernel": 94.23902893066406}, "output_dense": {"bias": 1.0458879470825195, "kernel": 101.17718505859375}}, "final_layer_norm": {"bias": 1.753927230834961, "scale": 25.785865783691406}, "layer_norm": {"bias": 2.33198881149292, "scale": 26.942312240600586}}, "4": {"attention": {"k_proj": {"bias": 0.44548165798187256, "kernel": 54.6234130859375}, "out_proj": {"bias": 1.652343988418579, "kernel": 50.16497039794922}, "q_proj": {"bias": 2.615248918533325, "kernel": 54.93098068237305}, "v_proj": {"bias": 0.34892427921295166, "kernel": 50.320098876953125}}, "feed_forward": {"intermediate_dense": {"bias": 1.9531786441802979, "kernel": 104.50968933105469}, "output_dense": {"bias": 0.8575068712234497, "kernel": 95.54541015625}}, "final_layer_norm": {"bias": 1.9920002222061157, "scale": 21.200180053710938}, "layer_norm": {"bias": 2.054612159729004, "scale": 23.620338439941406}}, "40": {"attention": {"k_proj": {"bias": 0.6590453386306763, "kernel": 44.198089599609375}, "out_proj": {"bias": 1.6252505779266357, "kernel": 49.55699920654297}, "q_proj": {"bias": 1.9674756526947021, "kernel": 44.89208221435547}, "v_proj": {"bias": 0.4587768614292145, "kernel": 49.23614501953125}}, "feed_forward": {"intermediate_dense": {"bias": 2.0333969593048096, "kernel": 92.16896057128906}, "output_dense": {"bias": 1.087776780128479, "kernel": 98.40738677978516}}, "final_layer_norm": {"bias": 1.7852704524993896, "scale": 25.04292106628418}, "layer_norm": {"bias": 2.2756104469299316, "scale": 26.40799903869629}}, "41": {"attention": {"k_proj": {"bias": 1.7133712768554688, "kernel": 41.96858596801758}, "out_proj": {"bias": 1.3790823221206665, "kernel": 51.25593566894531}, "q_proj": {"bias": 1.71382737159729, "kernel": 42.56317901611328}, "v_proj": {"bias": 0.4695759415626526, "kernel": 50.369529724121094}}, "feed_forward": {"intermediate_dense": {"bias": 2.110393524169922, "kernel": 88.92567443847656}, "output_dense": {"bias": 1.1446669101715088, "kernel": 97.37409973144531}}, "final_layer_norm": {"bias": 2.23917293548584, "scale": 28.507984161376953}, "layer_norm": {"bias": 2.22525691986084, "scale": 28.246891021728516}}, "42": {"attention": {"k_proj": {"bias": 0.8601109981536865, "kernel": 38.31235885620117}, "out_proj": {"bias": 1.4427157640457153, "kernel": 45.07648849487305}, "q_proj": {"bias": 1.549715280532837, "kernel": 39.524009704589844}, "v_proj": {"bias": 0.6933339834213257, "kernel": 43.49076461791992}}, "feed_forward": {"intermediate_dense": {"bias": 1.9107489585876465, "kernel": 88.009765625}, "output_dense": {"bias": 1.1978566646575928, "kernel": 95.7593994140625}}, "final_layer_norm": {"bias": 1.9227323532104492, "scale": 29.817535400390625}, "layer_norm": {"bias": 1.6761282682418823, "scale": 26.810440063476562}}, "43": {"attention": {"k_proj": {"bias": 1.247081995010376, "kernel": 34.694725036621094}, "out_proj": {"bias": 1.4174811840057373, "kernel": 41.36320495605469}, "q_proj": {"bias": 1.3773530721664429, "kernel": 35.38981628417969}, "v_proj": {"bias": 0.5787136554718018, "kernel": 39.29212951660156}}, "feed_forward": {"intermediate_dense": {"bias": 1.8936835527420044, "kernel": 87.073974609375}, "output_dense": {"bias": 0.9419379234313965, "kernel": 93.74283599853516}}, "final_layer_norm": {"bias": 1.99924898147583, "scale": 32.0491943359375}, "layer_norm": {"bias": 1.7940990924835205, "scale": 25.131242752075195}}, "44": {"attention": {"k_proj": {"bias": 2.5188145637512207, "kernel": 35.16314697265625}, "out_proj": {"bias": 1.1667875051498413, "kernel": 45.019126892089844}, "q_proj": {"bias": 1.317202091217041, "kernel": 35.58300018310547}, "v_proj": {"bias": 0.38874924182891846, "kernel": 44.14718246459961}}, "feed_forward": {"intermediate_dense": {"bias": 1.9462898969650269, "kernel": 86.07953643798828}, "output_dense": {"bias": 0.859266996383667, "kernel": 91.58765411376953}}, "final_layer_norm": {"bias": 2.0454273223876953, "scale": 34.2881965637207}, "layer_norm": {"bias": 1.6815991401672363, "scale": 25.14142608642578}}, "45": {"attention": {"k_proj": {"bias": 2.081407308578491, "kernel": 34.86139678955078}, "out_proj": {"bias": 1.0356104373931885, "kernel": 48.59937286376953}, "q_proj": {"bias": 1.402512788772583, "kernel": 35.03264617919922}, "v_proj": {"bias": 0.4231463074684143, "kernel": 48.76853942871094}}, "feed_forward": {"intermediate_dense": {"bias": 2.016927719116211, "kernel": 82.93773651123047}, "output_dense": {"bias": 0.9764893054962158, "kernel": 87.24796295166016}}, "final_layer_norm": {"bias": 1.9180456399917603, "scale": 33.143672943115234}, "layer_norm": {"bias": 1.5726068019866943, "scale": 23.782546997070312}}, "46": {"attention": {"k_proj": {"bias": 1.5659263134002686, "kernel": 35.878021240234375}, "out_proj": {"bias": 0.8182340264320374, "kernel": 51.16078186035156}, "q_proj": {"bias": 1.5642974376678467, "kernel": 36.18907165527344}, "v_proj": {"bias": 0.4092414081096649, "kernel": 51.89159393310547}}, "feed_forward": {"intermediate_dense": {"bias": 2.0093321800231934, "kernel": 77.47581481933594}, "output_dense": {"bias": 1.1406863927841187, "kernel": 77.73695373535156}}, "final_layer_norm": {"bias": 1.8108854293823242, "scale": 28.70657730102539}, "layer_norm": {"bias": 1.3991491794586182, "scale": 22.808137893676758}}, "47": {"attention": {"k_proj": {"bias": 0.6173280477523804, "kernel": 38.678985595703125}, "out_proj": {"bias": 0.6758822202682495, "kernel": 46.45281219482422}, "q_proj": {"bias": 1.7084776163101196, "kernel": 39.426841735839844}, "v_proj": {"bias": 0.4932914674282074, "kernel": 47.617279052734375}}, "feed_forward": {"intermediate_dense": {"bias": 1.986911654472351, "kernel": 75.47482299804688}, "output_dense": {"bias": 0.6346586346626282, "kernel": 72.82707214355469}}, "final_layer_norm": {"bias": 1.1888140439987183, "scale": 23.650447845458984}, "layer_norm": {"bias": 1.2521969079971313, "scale": 20.66573715209961}}, "5": {"attention": {"k_proj": {"bias": 0.42588678002357483, "kernel": 50.1945686340332}, "out_proj": {"bias": 1.6038882732391357, "kernel": 51.2144889831543}, "q_proj": {"bias": 2.7522244453430176, "kernel": 50.37500762939453}, "v_proj": {"bias": 0.3343381881713867, "kernel": 51.71652603149414}}, "feed_forward": {"intermediate_dense": {"bias": 1.8887722492218018, "kernel": 104.60663604736328}, "output_dense": {"bias": 0.8976269960403442, "kernel": 94.77360534667969}}, "final_layer_norm": {"bias": 2.1965675354003906, "scale": 21.37998390197754}, "layer_norm": {"bias": 2.0435237884521484, "scale": 22.437192916870117}}, "6": {"attention": {"k_proj": {"bias": 0.4843112528324127, "kernel": 51.87700653076172}, "out_proj": {"bias": 1.5925445556640625, "kernel": 50.83113479614258}, "q_proj": {"bias": 2.7889723777770996, "kernel": 52.3514404296875}, "v_proj": {"bias": 0.3247200846672058, "kernel": 51.107398986816406}}, "feed_forward": {"intermediate_dense": {"bias": 1.8638136386871338, "kernel": 103.7142333984375}, "output_dense": {"bias": 0.752193808555603, "kernel": 94.57742309570312}}, "final_layer_norm": {"bias": 2.5145251750946045, "scale": 20.836563110351562}, "layer_norm": {"bias": 2.0285890102386475, "scale": 23.156789779663086}}, "7": {"attention": {"k_proj": {"bias": 0.5048109889030457, "kernel": 51.46453094482422}, "out_proj": {"bias": 1.4398455619812012, "kernel": 51.139068603515625}, "q_proj": {"bias": 2.550907611846924, "kernel": 51.92047119140625}, "v_proj": {"bias": 0.42719271779060364, "kernel": 50.953285217285156}}, "feed_forward": {"intermediate_dense": {"bias": 1.8739449977874756, "kernel": 103.49991607666016}, "output_dense": {"bias": 0.5876641273498535, "kernel": 94.39387512207031}}, "final_layer_norm": {"bias": 2.416801929473877, "scale": 21.010677337646484}, "layer_norm": {"bias": 1.9788501262664795, "scale": 22.19708824157715}}, "8": {"attention": {"k_proj": {"bias": 0.49711495637893677, "kernel": 51.12122344970703}, "out_proj": {"bias": 1.2548246383666992, "kernel": 51.65118408203125}, "q_proj": {"bias": 2.541980504989624, "kernel": 51.03407287597656}, "v_proj": {"bias": 0.35420340299606323, "kernel": 51.662872314453125}}, "feed_forward": {"intermediate_dense": {"bias": 1.9298467636108398, "kernel": 103.21916198730469}, "output_dense": {"bias": 0.5482766628265381, "kernel": 93.97574615478516}}, "final_layer_norm": {"bias": 2.3526053428649902, "scale": 20.7393856048584}, "layer_norm": {"bias": 1.9221248626708984, "scale": 22.40435028076172}}, "9": {"attention": {"k_proj": {"bias": 0.5231171250343323, "kernel": 52.01068878173828}, "out_proj": {"bias": 1.4968843460083008, "kernel": 52.671897888183594}, "q_proj": {"bias": 2.4629459381103516, "kernel": 52.26807403564453}, "v_proj": {"bias": 0.38445231318473816, "kernel": 52.86597442626953}}, "feed_forward": {"intermediate_dense": {"bias": 2.026733875274658, "kernel": 101.98421478271484}, "output_dense": {"bias": 0.6828575134277344, "kernel": 94.36962890625}}, "final_layer_norm": {"bias": 2.325080156326294, "scale": 20.160720825195312}, "layer_norm": {"bias": 2.0236480236053467, "scale": 24.083864212036133}}}, "pos_conv_embed": {"conv": {"bias": 5.847014427185059, "weight_g": 9.12463665008545, "weight_v": 93.52015686035156}}}, "feature_extractor": {"conv_layers": {"0": {"conv": {"bias": 2.0290679931640625, "kernel": 20.55536460876465}, "layer_norm": {"bias": 4.550922393798828, "scale": 16.167570114135742}}, "1": {"conv": {"bias": 1.7790228128433228, "kernel": 51.24136734008789}, "layer_norm": {"bias": 5.962646961212158, "scale": 23.268157958984375}}, "2": {"conv": {"bias": 1.140576720237732, "kernel": 46.50312042236328}, "layer_norm": {"bias": 4.176670551300049, "scale": 20.370853424072266}}, "3": {"conv": {"bias": 0.6725863218307495, "kernel": 44.397525787353516}, "layer_norm": {"bias": 3.888174533843994, "scale": 17.53795051574707}}, "4": {"conv": {"bias": 0.6373162269592285, "kernel": 41.314056396484375}, "layer_norm": {"bias": 2.385471820831299, "scale": 16.34571647644043}}, "5": {"conv": {"bias": 0.5147221684455872, "kernel": 37.479759216308594}, "layer_norm": {"bias": 2.020900011062622, "scale": 17.064470291137695}}, "6": {"conv": {"bias": 0.4947893023490906, "kernel": 40.64780044555664}, "layer_norm": {"bias": 0.5876954793930054, "scale": 19.058603286743164}}}}, "feature_projection": {"layer_norm": {"bias": 6.376383304595947, "scale": 16.443069458007812}, "projection": {"bias": 1.8670344352722168, "kernel": 37.218414306640625}}, "masked_spec_embed": 11.914372444152832}}, "train/learning_rate": 1.9151924789184704e-05, "train/loss": 0.204779714345932, "train/param_norm": 1241.662353515625, "_runtime": 4765, "_timestamp": 1660121820, "_step": 275600, "_wandb": {"runtime": 4766}}

wandb/run-20220810_073735-23avj35z/logs/debug-internal.log

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

wandb/run-20220810_073735-23avj35z/logs/debug.log

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

wandb/run-20220810_073735-23avj35z/run-23avj35z.wandb

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

wandb/run-20220810_111559-290849gb/files/code/run_flax_speech_recognition_ctc.py

@@ -0,0 +1,1631 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace 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 connectionist temporal classification (CTC) 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 re
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import datasets
+import numpy as np
+from datasets import DatasetDict, load_dataset, load_metric
+from tqdm import tqdm
+
+import flax
+import jax
+import jax.numpy as jnp
+import optax
+import transformers
+import wandb as wandb
+from flax import core, jax_utils, struct, traverse_util
+from flax.jax_utils import unreplicate, pad_shard_unpad
+from flax.training.common_utils import get_metrics, shard, shard_prng_key
+from huggingface_hub import Repository
+from models import Wav2Vec2Config, FlaxWav2Vec2ForCTC
+from optax._src import linear_algebra
+from transformers import (
+    AutoFeatureExtractor,
+    AutoProcessor,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainingArguments,
+    is_tensorboard_available,
+    set_seed,
+)
+from transformers.file_utils import get_full_repo_name
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.17.0.dev0")
+
+require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/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)."
+        },
+    )
+    freeze_feature_encoder: bool = field(
+        default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}
+    )
+    attention_dropout: float = field(
+        default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}
+    )
+    activation_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The hidden activation dropout probability in the embeddings, encoder, and pooler."
+        },
+    )
+    hidden_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The dropout probability for all fully connected layers in the embeddings, encoder, and pooler."
+        },
+    )
+    feat_proj_dropout: float = field(
+        default=0.0,
+        metadata={
+            "help": "The feat proj dropout probability for feature encoder representations."
+        },
+    )
+    final_dropout: float = field(
+        default=0.0,
+        metadata={"help": "The dropout probability for the final projection layer."},
+    )
+    mask_time_prob: float = field(
+        default=0.1,
+        metadata={
+            "help": "The spec aug dropout probability for feature encoder representations."
+        },
+    )
+    mask_time_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the time axis."},
+    )
+    mask_feature_prob: float = field(
+        default=0.0,
+        metadata={
+            "help": "Probability of each feature vector along the feature axis to be chosen as the start of the vector"
+                    "span to be masked. Approximately ``mask_feature_prob * sequence_length // mask_feature_length`` feature bins will be masked along the time axis."
+        },
+    )
+    mask_feature_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the feature axis."},
+    )
+    layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})
+    ctc_loss_reduction: Optional[str] = field(
+        default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}
+    )
+    ctc_zero_infinity: Optional[bool] = field(
+        default=False, metadata={"help": "If True, will try yo aboud the CTC loss goinf to infinity."}
+    )
+
+
+@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."
+        },
+    )
+    max_test_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of test 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": "Truncate audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_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: Optional[int] = field(
+        default=512,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    min_label_length: Optional[int] = field(
+        default=2,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    pad_input_to_multiple_of: Optional[int] = field(
+        default=32000,
+        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."
+        },
+    )
+    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."
+        },
+    )
+    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 training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    do_lower_case: bool = field(
+        default=True,
+        metadata={"help": "Whether the target text should be lower cased."},
+    )
+    wandb_project: str = field(
+        default="flax-speech-recognition-ctc",
+        metadata={"help": "The name of the wandb project."},
+    )
+    wandb_name: str = field(
+        default=None,
+        metadata={"help": "The name of the wandb run."},
+    )
+    wandb_job_type: str = field(
+        default="CTC",
+        metadata={"help": "The name of the wandb job type."},
+    )
+    test_split_name: str = field(
+        default="test",
+        metadata={"help": "The name of the test data set split to use (via the datasets library). Defaults to 'test'"},
+    )
+    remove_punctuation: bool = field(
+        default=False, metadata={"help": "Whether or not to remove punctuation during training."}
+    )
+    skip_steps: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": "Skip this number of steps. Useful to continue training"
+        },
+    )
+
+
+# @flax.struct.dataclass
+@dataclass
+class FlaxTrainingArguments(TrainingArguments):
+    precision: str = field(
+        default="full",
+        metadata={
+            "help": "Whether to enable mixed-precision training. If true, the optimizer is stored in half-precision (bfloat16) and computations are executed in half-precision"
+            "**Note that this only specifies the dtype of the computation and optimizer state. It does not influence the dtype of model parameters.**"
+        },
+    )
+    matmul_precision: str = field(
+        default="default",
+        metadata={
+            "help": "Default floating-point precision of internal computations used in TPU matrix multiplications and convolutions. "
+            "This configuration option controls the default precision for JAX operations that take an optional precision argument (e.g. `lax.conv_general_dilated` and `lax.dot`). "
+            "This configuration option does not change the behaviours of such calls with explicit precision arguments; "
+            "it only changes the behaviors of calls with no such argument provided. "
+            "One of `['highest', 'float32', 'high', 'bfloat16_3x', 'default', 'bfloat16', 'fastest', None]`."
+        },
+    )
+    multisteps: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to use Optax MultiSteps for gradient accumulation. If `False` (default) and `gradient_accumulation_steps > 1`, "
+            "a custom gradient accumulation implementation will be employed."
+        },
+    )
+
+
+def to_fp32(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+
+
+def to_bf16(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t)
+
+
+class MixedPrecisionTrainState(struct.PyTreeNode):
+    """Train state for use with a single Optax optimizer.
+    Adapted from flax train_state https://github.com/google/flax/blob/main/flax/training/train_state.py
+
+    Synopsis::
+
+        state = TrainState.create(
+            apply_fn=model.apply,
+            params=variables['params'],
+            tx=tx)
+        grad_fn = jax.grad(make_loss_fn(state.apply_fn))
+        for batch in data:
+          grads = grad_fn(state.params, batch)
+          state = state.apply_gradients(grads=grads)
+
+    Args:
+      step: Counter starts at 0 and is incremented by every call to
+        `.apply_gradients()`.
+      apply_fn: Usually set to `model.apply()`. Kept in this dataclass for
+        convenience to have a shorter params list for the `train_step()` function
+        in your training loop.
+      params: The parameters to be updated by `tx` and used by `apply_fn`.
+      tx: An Optax gradient transformation.
+      opt_state: The state for `tx`.
+      dropout_rng: PRNG key for stochastic operations.
+      bf16: Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training.
+    """
+
+    step: int
+    apply_fn: Callable = struct.field(pytree_node=False)
+    get_attention_mask_fn: Callable = struct.field(pytree_node=False)
+    params: core.FrozenDict[str, Any]
+    tx: optax.GradientTransformation = struct.field(pytree_node=False)
+    opt_state: optax.OptState
+    dropout_rng: jnp.ndarray
+    max_grad_norm: Optional[float] = 1.0
+
+    def apply_gradients(self, *, grads, to_dtype, **kwargs):
+        """Updates `step`, `params`, `opt_state` and `**kwargs` in return value.
+
+        Note that internally this function calls `.tx.update()` followed by a call
+        to `optax.apply_updates()` to update `params` and `opt_state`.
+
+        Args:
+          grads: Gradients that have the same pytree structure as `.params`.
+          **kwargs: Additional dataclass attributes that should be `.replace()`-ed.
+
+        Returns:
+          An updated instance of `self` with `step` incremented by one, `params`
+          and `opt_state` updated by applying `grads`, and additional attributes
+          replaced as specified by `kwargs`.
+        """
+
+        # clip gradients by global l2 norm
+        casted_max_grad_norm = to_dtype(self.max_grad_norm)
+        g_norm = linear_algebra.global_norm(grads)
+        g_norm = jnp.maximum(casted_max_grad_norm, g_norm)
+        grads = jax.tree_util.tree_map(lambda t: (t / g_norm) * casted_max_grad_norm, grads)
+
+        # perform update step in fp32 and subsequently downcast optimizer states if mixed precision training
+        # grads and opt_state in bf16 (need to upcast), params in fp32 (leave as is)
+        updates, new_opt_state = self.tx.update(to_fp32(grads), to_fp32(self.opt_state), self.params)
+
+        new_params = optax.apply_updates(self.params, updates)
+        return self.replace(
+            step=self.step + 1,
+            params=new_params,
+            opt_state=to_dtype(new_opt_state),
+            **kwargs,
+        )
+
+    @classmethod
+    def create(cls, *, apply_fn, params, tx, to_dtype, **kwargs):
+        """Creates a new instance with `step=0` and initialized `opt_state`."""
+        # downcast optimizer state to bf16 if mixed-precision training
+        opt_state = tx.init(to_dtype(params)) if tx is not None else None
+        return cls(
+            step=0,
+            apply_fn=apply_fn,
+            params=params,
+            tx=tx,
+            opt_state=opt_state,
+            **kwargs,
+        )
+
+    def replicate(self):
+        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
+
+
+@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).
+        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
+    input_padding: Union[bool, str] = "longest"
+    label_padding: Union[bool, str] = "max_length"
+    pad_input_to_multiple_of: Optional[int] = None
+    pad_to_multiple_of_label: Optional[int] = None
+    max_input_length: Optional[float] = None
+    max_label_length: Optional[float] = 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
+        input_features = [{"input_values": feature["input_values"]} for feature in features]
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+
+        # reformat list to dict and set to pytorch format
+        batch = self.processor.feature_extractor.pad(
+            input_features,
+            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_label_length,
+            padding=self.label_padding,
+            pad_to_multiple_of=self.pad_to_multiple_of_label,
+            return_tensors="np",
+        )
+
+        labels = labels_batch["input_ids"]
+        labels = np.ma.array(labels, mask=np.not_equal(labels_batch.attention_mask, 1))
+        labels = labels.filled(fill_value=-100)
+
+        batch["labels"] = labels
+
+        return batch
+
+
+def get_grouped_indices(
+    dataset, batch_size: int, rng: Optional[List[int]] = None, mega_batch_mult: Optional[int] = None
+) -> np.array:
+    """
+    Adapted from the `get_length_grouped_indices` function in the PyTorch Trainer utils file (https://github.com/huggingface/transformers/blob/main/src/transformers/trainer_pt_utils.py#L486)
+    Function that returns a list of indices in which each slice of `batch_size` consecutive indices correspond to elements of similar
+    lengths. To do this, the indices are:
+
+    - randomly permuted (if a JAX rng is specified)
+    - grouped in mega-batches of size `mega_batch_mult * batch_size`
+    - sorted by length in each mega-batch
+
+    The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of
+    maximum length placed first, so that an OOM happens sooner rather than later.
+    """
+    lengths = dataset["input_length"]
+
+    # Default for mega_batch_mult: 50 or the number to get 4 megabatches, whichever is smaller.
+    if mega_batch_mult is None:
+        mega_batch_mult = min(len(lengths) // (batch_size * 4), 50)
+        # Just in case, for tiny datasets
+        if mega_batch_mult == 0:
+            mega_batch_mult = 1
+
+    # We need to use JAX for the random permutation as the PRNG key will be set based on the seed outside of the sampler.
+    num_samples = len(lengths)
+    indices = jax.random.permutation(rng, np.arange(num_samples)) if rng is not None else np.arange(num_samples)
+
+    megabatch_size = mega_batch_mult * batch_size
+    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
+    megabatches = [list(sorted(megabatch, key=lambda i: lengths[i], reverse=True)) for megabatch in megabatches]
+
+    # The rest is to get the biggest batch first.
+    # Since each megabatch is sorted by descending length, the longest element is the first
+    megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches]
+    max_idx = np.argmax(megabatch_maximums).item()
+    # Switch to put the longest batch in first position
+    # (note that this is different to the PT grouped sampler in which we only put the longest element in the first position, and not its batch)
+    megabatches[0], megabatches[max_idx] = megabatches[max_idx], megabatches[0]
+
+    megabatches = np.array([i for megabatch in megabatches for i in megabatch])
+
+    return megabatches
+
+
+def generate_batch_splits(samples_idx: np.ndarray, batch_size: int, drop_last=True) -> np.ndarray:
+    """Generate batches of data for a specified batch size from sample indices. If the dataset size is not divisible by
+    the batch size and `drop_last` is `True`, the last incomplete batch is dropped. Else, it is returned."""
+    num_samples = len(samples_idx)
+    if drop_last:
+        samples_to_remove = num_samples % batch_size
+        if samples_to_remove != 0:
+            samples_idx = samples_idx[:-samples_to_remove]
+        sections_split = num_samples // batch_size
+        samples_idx = samples_idx.reshape((sections_split, batch_size))
+    else:
+        sections_split = math.ceil(num_samples / batch_size)
+        samples_idx = np.array_split(samples_idx, sections_split)
+    return samples_idx
+
+
+def write_train_metric(summary_writer, train_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)
+
+
+def write_eval_metric(summary_writer, eval_metrics, step, pred_str=None):
+    for metric_name, value in eval_metrics.items():
+        summary_writer.scalar(f"eval_{metric_name}", value, step)
+
+    if pred_str is not None:
+        # write output actual predictions for debugging
+        summary_writer.text("eval_predictions", "\n".join(pred_str), step)
+
+
+def write_wandb_log(metrics, step, prefix=None):
+    if jax.process_index() == 0:
+        log_metrics = {}
+        for k, v in metrics.items():
+            if "layer" in k:
+                log_metrics[f"{k}/"] = v
+            elif prefix is not None:
+                log_metrics[f"{prefix}/{k}"] = v
+            else:
+                log_metrics[k] = v
+        wandb.log(log_metrics, step)
+
+
+def write_wandb_pred(pred_str, label_str, step, num_log=50, prefix="eval"):
+    if jax.process_index() == 0:
+        # convert str data to a wandb compatible format
+        str_data = [[label_str[i], pred_str[i]] for i in range(len(pred_str))]
+        # we'll log the first 50 predictions for each epoch
+        wandb.log(
+            {
+                f"{prefix}/step_{int(step / 1000)}k": wandb.Table(
+                    columns=["label_str", "pred_str"], data=str_data[:num_log]
+                )
+            },
+            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 ctc_loss(
+    logits,
+    logits_attention_mask,
+    labels,
+    blank_id,
+    loss_reduction="mean",
+    output_emission_dict=False,
+    log_epsilon=-100000.0,
+):
+    """Computes CTC loss.
+    This function performs forward computation over an FSA with `N * 2` states
+    where `N` is the max number of labels. The states are split into two groups:
+    Phi states and emission states. a phi-state accepts repetition of
+    phi (blank)-symbols and transits to emission state when the correct label is
+    observed. An emission state accepts repetition of the label and transits to
+    the next phi states at any time (so called epsilon-transition).
+    Below, `B` denotes the batch size, `T` denotes the time steps in `logits`,
+    and `N` denotes the time steps in `labels`.
+    Args:
+      logits: (B, T, K)-array containing log-probabilities of each class.
+      logitpaddings: (B, T)-array. Padding indicators for `logits`.
+      labels: (B, N)-array containing reference integer labels.
+      labelpaddings: (B, N)-array. Padding indicators for `labels`. Currently,
+        `labels` must be right-padded, i.e. each row of `labelpaddings` must be
+        repetition of zeroes, followed by repetition of ones.
+      blank_id: Id for blank token.
+      loss_reduction: one of "mean", "sum", "default"
+        - "none": no reduction is applied.
+        - "mean": output loss will be divided by target lengths and then the
+          mean over the batch is taken.
+        - "sum": output loss are summed over batch
+      output_emission_dict: whether to output additional information about the emission probs
+    Returns:
+      A pair of `(per_seq_loss, aux)`.
+      per_seq_loss:
+        (B,)-array containing loss values for each sequence in the batch.
+      aux: Dictionary containing interim variables used for computing losses.
+        aux['logalpha_phi']: (T, B, N+1)-array. Log-forward-probabilities of each
+          phi-state corresponding to the n-th label.
+        aux['logalpha_emit']: (T, B, N)-array. Log-forward-probabilities of each
+          emission-state corresponding to the n-th label.
+        aux['logprobs_phi']: (T, B, 1)-array. Probability of the phi-symbol
+          corresponding to each time frame.
+        aux['logprobs_emit']: (T, B, N)-array. Probability of the n-th label
+          corresponding to each time frame.
+    """
+    # label paddings are indicated by -100
+    labelpaddings = labels < 0
+    # logit paddings are the inverse of attention_mask
+    logitpaddings = ~logits_attention_mask
+
+    # Copied from https://github.com/tensorflow/lingvo/blob/master/lingvo/jax/layers/ctc_objectives.py
+    batchsize, unused_maxinputlen, num_classes = logits.shape
+    batchsize_, maxlabellen = labels.shape
+
+    logprobs = jax.nn.log_softmax(logits)
+    labellens = maxlabellen - jnp.sum(labelpaddings, axis=1).astype(jnp.int32)
+
+    # repeat[b, n] == 1.0 when label[b, n] == label[b, n+1].
+    repeat = (labels[:, :-1] == labels[:, 1:]).astype(jnp.float32)
+    repeat = jnp.pad(repeat, ((0, 0), (0, 1)))
+
+    logprobs_phi = logprobs[:, :, blank_id : blank_id + 1]  # [B, T, 1]
+    logprobs_phi = jnp.transpose(logprobs_phi, (1, 0, 2))  # [T, B, 1]
+
+    one_hot = jax.nn.one_hot(labels, num_classes=num_classes)  # [B, N, K]
+    logprobs_emit = jnp.einsum("btk,bnk->btn", logprobs, one_hot)
+    logprobs_emit = jnp.transpose(logprobs_emit, (1, 0, 2))  # [T, B, N]
+
+    logalpha_phi_init = jnp.ones((batchsize, maxlabellen + 1)) * log_epsilon  # [B, N]
+    logalpha_phi_init = logalpha_phi_init.at[:, 0].set(0.0)
+    logalpha_emit_init = jnp.ones((batchsize, maxlabellen)) * log_epsilon  # [B, N]
+
+    def loop_body(prev, x):
+        prev_phi, prev_emit = prev
+        # emit-to-phi epsilon transition, except if the next label is repetition
+        prev_phi_orig = prev_phi
+        prev_phi = prev_phi.at[:, 1:].set(jnp.logaddexp(prev_phi[:, 1:], prev_emit + log_epsilon * repeat))
+
+        logprob_emit, logprob_phi, pad = x
+
+        # phi-to-emit transition
+        next_emit = jnp.logaddexp(prev_phi[:, :-1] + logprob_emit, prev_emit + logprob_emit)
+        # self-loop transition
+        next_phi = prev_phi + logprob_phi
+        # emit-to-phi blank transition only when the next label is repetition
+        next_phi = next_phi.at[:, 1:].set(
+            jnp.logaddexp(next_phi[:, 1:], prev_emit + logprob_phi + log_epsilon * (1.0 - repeat))
+        )
+
+        pad = pad.reshape((batchsize, 1))
+        next_emit = pad * prev_emit + (1.0 - pad) * next_emit
+        next_phi = pad * prev_phi_orig + (1.0 - pad) * next_phi
+
+        return (next_phi, next_emit), (next_phi, next_emit)
+
+    xs = (logprobs_emit, logprobs_phi, logitpaddings.transpose((1, 0)))
+    _, (logalpha_phi, logalpha_emit) = jax.lax.scan(loop_body, (logalpha_phi_init, logalpha_emit_init), xs)
+
+    # last row needs to be updated with the last epsilon transition
+    logalpha_phi_last = logalpha_phi[-1].at[:, 1:].set(jnp.logaddexp(logalpha_phi[-1, :, 1:], logalpha_emit[-1]))
+    logalpha_phi = logalpha_phi.at[-1].set(logalpha_phi_last)
+
+    # extract per_seq_loss
+    one_hot = jax.nn.one_hot(labellens, num_classes=maxlabellen + 1)  # [B, N+1]
+    per_seq_loss = -jnp.einsum("bn,bn->b", logalpha_phi_last, one_hot)
+
+    if loss_reduction == "mean":
+        target_lengths = labelpaddings.shape[-1] - labelpaddings.sum(axis=-1)
+        loss = (per_seq_loss / target_lengths).mean()
+    elif loss_reduction == "sum":
+        loss = per_seq_loss.sum()
+    else:
+        loss = per_seq_loss
+
+    if not output_emission_dict:
+        return loss
+
+    return loss, {
+        "logalpha_phi": logalpha_phi,
+        "logalpha_emit": logalpha_emit,
+        "logprobs_phi": logprobs_phi,
+        "logprobs_emit": logprobs_emit,
+    }
+
+
+def make_dataset(data_args, seed=42):
+    # Pre-processing dataset
+    import re
+
+    def map_nst(entry):
+        text = entry["text"].lower()
+        text = text.replace("(...vær stille under dette opptaket...)", "")
+        text = re.sub('[áàâ]', 'a', text)
+        text = re.sub('[ä]', 'æ', text)
+        text = re.sub('[éèëê]', 'e', text)
+        text = re.sub('[íìïî]', 'i', text)
+        text = re.sub('[óòöô]', 'o', text)
+        text = re.sub('[ö]', 'ø', text)
+        text = re.sub('[ç]', 'c', text)
+        text = re.sub('[úùüû]', 'u', text)
+        # text = re.sub('\\(?=(Punktum|Komma|Utropstegn|Spørsmålstegn))', ' ', text)
+        text = re.sub('\s+', ' ', text)
+        return {"text": text}
+
+    def filter_nst(entry):
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.match(entry["type"], "pIW|CA"):
+            return False  # Spelling out words
+        return True
+
+    def filter_npsc(entry):
+        # False if there are digits in the text
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.search("\d", entry["text"]):
+            return False
+        return True
+
+    def map_npsc(entry):
+        batch = {"text": entry["text"].lower()}
+        batch["text"] = re.sub('[áàâ]', 'a', batch["text"])
+        batch["text"] = re.sub('[ä]', 'æ', batch["text"])
+        batch["text"] = re.sub('[éèëê]', 'e', batch["text"])
+        batch["text"] = re.sub('[íìïî]', 'i', batch["text"])
+        batch["text"] = re.sub('[óòöô]', 'o', batch["text"])
+        batch["text"] = re.sub('[ö]', 'ø', batch["text"])
+        batch["text"] = re.sub('[ç]', 'c', batch["text"])
+        batch["text"] = re.sub('[úùüû]', 'u', batch["text"])
+        batch["text"] = re.sub('\s', ' ', batch["text"])
+        batch["text"] = re.sub('<ee>', 'eee', batch["text"])
+        batch["text"] = re.sub('<qq>', 'qqq', batch["text"])
+        batch["text"] = re.sub('<mm>', 'mmm', batch["text"])
+        batch["text"] = re.sub('<inaudible>', 'xxx', batch["text"])
+        # batch["text"] = re.sub('<inaudible>', '?', batch["text"])
+        if "<" in batch["text"]:
+            raise ValueError(batch["text"])
+        return batch
+
+    nst = datasets.load_dataset("NbAiLab/NST", "no-close")
+    npsc = datasets.load_dataset("NbAiLab/NPSC", "16K_mp3")
+    # TODO NST_hesitate
+
+    split = len(npsc[data_args.train_split_name]) / (len(npsc[data_args.train_split_name]) + len(npsc[data_args.eval_split_name]))  # Use same train/val ratio as NPSC
+    nst_train = nst[data_args.train_split_name].train_test_split(train_size=split, seed=seed)
+    nst[data_args.train_split_name] = nst_train["train"]
+    nst[data_args.eval_split_name] = nst_train["test"]
+
+    nst = nst.filter(filter_nst).map(
+        map_nst,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NST",
+    ).shuffle(seed=seed)
+    npsc = npsc.filter(filter_npsc).map(
+        map_npsc,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NPSC",
+    ).shuffle(seed=seed)
+
+    npsc_base = npsc.remove_columns([col for col in npsc[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+    nst_base = nst.remove_columns([col for col in nst[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+
+    combined = {}
+    for split in data_args.train_split_name, data_args.eval_split_name, data_args.test_split_name:
+        probs = np.array([len(nst_base[split]), len(npsc_base[split])])  # Weight by number of examples
+        probs = (probs / probs.sum()).tolist()
+        comb = datasets.interleave_datasets([nst_base[split], npsc_base[split]], probabilities=probs, seed=seed)
+        combined[split] = comb
+
+    return datasets.DatasetDict(**combined)
+
+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, FlaxTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # 2. 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()
+
+    # Set up wandb run
+    if jax.process_index() == 0:
+        wandb.init(project=data_args.wandb_project, name=data_args.wandb_name, job_type=data_args.wandb_job_type)
+
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set the default TPU matmul precision and display the number of devices
+    jax.config.update("jax_default_matmul_precision", training_args.matmul_precision)
+    logger.info(f"JAX devices: {jax.device_count()}, matmul precision: {training_args.matmul_precision}")
+
+    # 4. Load dataset
+
+    set_seed(training_args.seed)
+    raw_datasets = make_dataset(data_args, seed=training_args.seed)
+   
+    # raw_datasets = DatasetDict()
+
+    # if training_args.do_train:
+    #     raw_datasets[data_args.train_split_name] = 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[data_args.eval_split_name] = 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 training_args.do_predict:
+    #     test_split = data_args.test_split_name.split("+")
+    #     for split in test_split:
+    #         raw_datasets[split] = load_dataset(
+    #             data_args.dataset_name,
+    #             data_args.dataset_config_name,
+    #             split=split,
+    #             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 and not training_args.do_predict:
+        raise ValueError(
+            "Cannot not train, not do evaluation and not do prediction. At least one of "
+            "training, evaluation or prediction has to be done."
+        )
+
+    # if not training, there is no need to run multiple epochs
+    if not training_args.do_train:
+        training_args.num_train_epochs = 1
+
+    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
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    config = Wav2Vec2Config.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,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    # update config according to training args, model args, and tokenizer attributes
+    config.update(
+        {
+            "feat_proj_dropout": model_args.feat_proj_dropout,
+            "attention_dropout": model_args.attention_dropout,
+            "hidden_dropout": model_args.hidden_dropout,
+            "final_dropout": model_args.final_dropout,
+            "mask_time_prob": model_args.mask_time_prob,
+            "mask_time_length": model_args.mask_time_length,
+            "mask_feature_prob": model_args.mask_feature_prob,
+            "mask_feature_length": model_args.mask_feature_length,
+            "gradient_checkpointing": training_args.gradient_checkpointing,
+            "layerdrop": model_args.layerdrop,
+            "ctc_loss_reduction": model_args.ctc_loss_reduction,
+            "ctc_zero_infinity": model_args.ctc_zero_infinity,
+            "pad_token_id": tokenizer.pad_token_id,
+            "vocab_size": tokenizer.vocab_size,  # len(tokenizer),
+            "activation_dropout": model_args.activation_dropout,
+        }
+    )
+
+    if tokenizer.do_lower_case and data_args.dataset_name != "librispeech_asr":
+        raise ValueError(
+            "Setting the tokenizer attribute `do_lower_case` to `True` converts all input strings to "
+            "uppercase prior to tokenization. This should only be done when the tokenizer is built on an uppercased corpus,"
+            "i.e. for the dataset `librispeech_asr` only. If your dataset is not `librispeech_asr`, the tokenizer is mostly likely "
+            "built on an lowercased corpus. In this case, set `tokenizer.do_lower_case` to ``False`."
+        )
+
+    if training_args.precision == "full_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = True
+    elif training_args.precision == "half_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = False
+    else:
+        dtype = jnp.float32
+        training_args.mixed_precision = False
+
+    try:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+        )
+    except:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+            from_pt=True,
+        )
+
+    # 6. Resample speech dataset ALWAYS
+    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_target_length = data_args.max_label_length
+    min_target_length = data_args.min_label_length
+    pad_input_to_multiple_of = data_args.pad_input_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
+    dataset_name = data_args.dataset_name
+    chars_to_ignore = ', ? . ! - ; : " “ % ‘ ” ?'.split(" ")
+    chars_to_ignore_regex = f'[{"".join(chars_to_ignore)}]'
+    # gigaspeech_punctuation = {" <comma>": ",", " <period>": ".", " <questionmark>": "?", " <exclamationpoint>": "!"}
+    # gigaspeech_disfluencies = ["<other>", "<sil>"]
+    # swb_disfluencies = ["[noise]", "[laughter]", "[silence]", "<a_aside>", "<b_aside>", "<e_aside>", "[laughter-",
+    #                 "[vocalized-noise]", "_1"]
+    # swb_punctuations = ["{", "}", "[", "]-", "]"]
+    # earnings_disfluencies = ["<crosstalk>", "<affirmative>", "<inaudible>", "inaudible", "<laugh>", "<unk>"]
+    ignore_segments = ["ignore_time_segment_in_scoring", "<noise>", "<music>", "[noise]", "[laughter]", "[silence]",
+                       "[vocalized-noise]", "<crosstalk>", "<affirmative>", "<inaudible>", "<laugh>", "<other>", "<sil>", ""]
+
+    if training_args.do_train and data_args.max_train_samples is not None:
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].select(range(data_args.max_train_samples))
+
+    if training_args.do_eval and data_args.max_eval_samples is not None:
+        raw_datasets[data_args.eval_split_name] = raw_datasets[data_args.eval_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_predict and data_args.max_test_samples is not None:
+        raw_datasets[data_args.test_split_name] = raw_datasets[data_args.test_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_train and data_args.remove_punctuation:
+
+        def remove_punctuation(batch):
+            batch[text_column_name] = (
+                re.sub(chars_to_ignore_regex, "", batch[text_column_name]).replace("'", "").replace('"', "")
+            )
+
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].map(
+            remove_punctuation,
+            num_proc=data_args.preprocessing_num_workers,
+            desc="removing punctuation from train split",
+        )
+
+    # filter data where the targets are ignored in scoring
+    def is_target_labels(input_str):
+        return input_str.lower() not in ignore_segments
+
+    raw_datasets = raw_datasets.filter(
+            is_target_labels,
+            num_proc=num_workers,
+            input_columns=[text_column_name],
+            desc="filtering data where the targets are ignored in scoring",
+        )
+
+    def prepare_dataset(batch):
+        # process audio
+        try:
+            sample = batch[audio_column_name]
+        except ValueError:
+            sample = {"array": np.array([0.]), "sampling_rate": feature_extractor.sampling_rate}
+        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
+        # process audio length
+        batch[model_input_name] = inputs.input_values[0]
+        batch["input_length"] = len(batch["input_values"])
+
+        # process targets
+        input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
+
+        # if dataset_name == "google/xtreme_s":
+        #     # Finally, we tokenize the processed text
+        #     batch["labels"] = tokenizer(input_str).input_ids
+        #     batch["labels_length"] = len(batch["labels"])
+        #     return batch
+
+        # # Common Voice 9
+        # if input_str.startswith('"') and input_str.endswith('"'):
+        #     # we can remove trailing quotation marks as they do not affect the transcription
+        #     input_str = input_str[1:-1]
+        # # normalize quotation marks
+        # input_str = re.sub(r'["“”]', '"', input_str)
+        # # normalize apostrophes
+        # input_str = re.sub(r"[’']", "'", input_str)
+        # # normalize hyphens
+        # input_str = re.sub(r"[—–]", "-", input_str)
+        # # replace double quotation marks with single
+        # input_str = input_str.replace('""', '"')
+        # if dataset_name == "mozilla-foundation/common_voice_9_0" and len(input_str):
+        #     # for CV9, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # TEDLIUM-3
+        # # delete the <unk> token from the text and replace spaced apostrophes with un-spaced
+        # input_str = input_str.replace("<unk>", "").replace(" '", "'")
+
+        # # GigaSpeech
+        # for disfluency in gigaspeech_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # convert spelled out punctuation to symbolic form
+        # for punctuation, replacement in gigaspeech_punctuation.items():
+        #     input_str = input_str.replace(punctuation, replacement)
+        # if dataset_name == "speechcolab/gigaspeech" and len(input_str):
+        #     # for GS, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # SWB
+        # for disfluency in swb_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # remove parenthesised text (test data only)
+        # input_str = re.sub("[\(].*?[\)]", "", input_str)
+        # for punctuation in swb_punctuations:
+        #     input_str = input_str.replace(punctuation, "")
+        # # replace anomalous words with their correct transcriptions
+        # split_str = input_str.split("/")
+        # if len(split_str) > 1:
+        #     input_str = " ".join(
+        #         [" ".join([" ".join(i.split(" ")[:-1]) for i in split_str])] + [split_str[-1].split(" ")[-1]])
+
+        # # Earnings 22
+        # for disfluency in earnings_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # replace mal-formatted ellipsis
+        # input_str = input_str.replace("…", ".")
+
+        # JIWER compliance
+        # remove multiple spaces
+        input_str = re.sub(r"\s\s+", " ", input_str)
+        # strip trailing spaces
+        input_str = input_str.strip()
+
+        # Finally, we tokenize the processed text
+        batch["labels"] = tokenizer(input_str).input_ids
+        batch["labels_length"] = len(batch["labels"])
+        return batch
+
+    vectorized_datasets = raw_datasets.map(
+        prepare_dataset,
+        remove_columns=next(iter(raw_datasets.values())).column_names,
+        num_proc=num_workers,
+        desc="preprocess dataset",
+    )
+
+    # filter data with inputs shorter than min_input_length or longer than max_input_length
+    def is_audio_in_length_range(length):
+        return length > min_input_length and length < max_input_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_audio_in_length_range,
+        num_proc=num_workers,
+        input_columns=["input_length"],
+    )
+
+    # filter data with targets shorter than min_target_length or longer than max_target_length
+    def is_labels_in_length_range(length):
+        return length > min_target_length  # and length < max_target_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_labels_in_length_range,
+        num_proc=num_workers,
+        input_columns=["labels_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 Metrics
+    wer_metric = load_metric("wer")
+    cer_metric = load_metric("cer")
+
+    def compute_metrics(pred_ids: List[List[int]], label_ids: List[List[int]]):
+        padded_ids = np.where(np.asarray(label_ids) == -100, tokenizer.pad_token_id, np.asarray(label_ids))
+
+        pred_str = tokenizer.batch_decode(pred_ids)
+        # we do not want to group tokens when computing the metrics
+        label_str = tokenizer.batch_decode(padded_ids, group_tokens=False)
+
+        wer = wer_metric.compute(predictions=pred_str, references=label_str)
+        cer = cer_metric.compute(predictions=pred_str, references=label_str)
+
+        return {"wer": wer, "cer": cer}, pred_str, label_str
+
+    # 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,
+        input_padding="longest",
+        pad_input_to_multiple_of=pad_input_to_multiple_of,
+        max_label_length=data_args.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."
+        )
+
+    # 10. Handle the repository creation
+    if training_args.push_to_hub:
+        with open(os.path.join(training_args.output_dir, ".gitattributes"), "r+") as f:
+            git_lfs_extensions = f.read()
+            if "*.wandb" not in git_lfs_extensions:
+                f.write("*.wandb filter=lfs diff=lfs merge=lfs -text")
+        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
+        repo = Repository(training_args.output_dir, clone_from=repo_name)
+
+    # 11. Initialize our training
+    rng = jax.random.PRNGKey(training_args.seed)
+    rng, dropout_rng = jax.random.split(rng)
+
+    # Store some constants
+    max_steps = int(training_args.max_steps)
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
+    batch_size_per_update = train_batch_size * gradient_accumulation_steps
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
+    to_dtype = to_bf16 if training_args.mixed_precision else to_fp32
+
+    if training_args.do_train:
+        num_train_samples = len(vectorized_datasets[data_args.train_split_name])
+        steps_per_epoch = num_train_samples // batch_size_per_update
+        if max_steps > 0:
+            num_epochs = -(training_args.max_steps // -steps_per_epoch)
+            total_train_steps = max_steps
+        else:
+            num_epochs = int(training_args.num_train_epochs)
+            total_train_steps = steps_per_epoch * num_epochs
+
+        # Create learning rate schedule
+        # Create learning rate schedule
+        linear_decay_lr_schedule_fn = create_learning_rate_fn(
+            total_train_steps,
+            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.
+        # Note that this mask is specifically adapted for FlaxWav2Vec2 and FlaxBart.
+        # For FlaxT5, one should correct the layer norm parameter naming
+        # accordingly - see `run_t5_mlm_flax.py` e.g.
+        def decay_mask_fn(params):
+            flat_params = traverse_util.flatten_dict(params)
+            layer_norm_params = [
+                (name, "scale")
+                for name in ["layer_norm", "self_attn_layer_norm", "layernorm_embedding", "final_layer_norm"]
+            ]
+            flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_params) for path in flat_params}
+            return traverse_util.unflatten_dict(flat_mask)
+
+        if training_args.adafactor:
+            # Create Adafactor optimizer
+            optim = optax.adafactor(
+                learning_rate=linear_decay_lr_schedule_fn,
+                dtype_momentum=jnp.bfloat16 if training_args.mixed_precision else jnp.float32,
+                weight_decay_rate=training_args.weight_decay,
+                weight_decay_mask=decay_mask_fn,
+            )
+        else:
+            # Create AdamW optimizer
+            optim = 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,
+            )
+
+        # Optax MultiSteps for gradient accumulation. We'll only call this optimizer transformation if gradient accumulation is required (i.e. gradient accumulation steps > 1)
+        if training_args.multisteps and gradient_accumulation_steps > 1:
+            optim = optax.MultiSteps(optim, gradient_accumulation_steps, use_grad_mean=False)
+    else:
+        num_epochs = 0
+        total_train_steps = 0
+        num_train_samples = 0
+        optim = None
+
+    # Setup train state
+    state = MixedPrecisionTrainState.create(
+        apply_fn=model.__call__,
+        get_attention_mask_fn=model._get_feature_vector_attention_mask,
+        params=model.params,
+        tx=optim,
+        to_dtype=to_dtype,
+        dropout_rng=dropout_rng,
+        max_grad_norm=training_args.max_grad_norm,
+    )
+
+    # Replicate the train state on each device
+    state = state.replicate()
+    blank_id = model.config.pad_token_id
+
+    # Define gradient update step fn
+    def train_step(state, batch):
+        # only one single rng per grad step, with or without accumulation, as the graph should be identical over one effective training batch
+        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
+
+        def compute_loss(params, minibatch):
+            labels = minibatch.pop("labels")
+            logits = state.apply_fn(
+                **minibatch,
+                params=params,
+                dropout_rng=dropout_rng,
+                freeze_feature_encoder=model_args.freeze_feature_encoder,
+                train=True,
+            )[0]
+            logits_mask = state.get_attention_mask_fn(logits.shape[1], batch["attention_mask"])
+            loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+            return loss
+
+        grad_fn = jax.value_and_grad(compute_loss)
+
+        if gradient_accumulation_steps == 1 or training_args.multisteps:
+            loss, grad = grad_fn(to_dtype(state.params), batch)
+
+        # Custom gradient accumulation
+        else:
+            # add a first dimension over gradient_accumulation_steps for minibatch slices
+            batch = jax.tree_util.tree_map(
+                lambda x: x.reshape(
+                    gradient_accumulation_steps, training_args.per_device_train_batch_size, *x.shape[1::]
+                ),
+                batch,
+            )
+
+            def accum_minibatch_step(accum_grad, minibatch):
+                # compute loss, num labels and grad over minibatch and accumulate
+                loss, grad = grad_fn(to_dtype(state.params), minibatch)
+                return jax.tree_util.tree_map(jnp.add, accum_grad, grad), loss
+
+            # create an initial state for accumulating losses, num labels and gradients
+            init_grad = jax.tree_util.tree_map(jnp.zeros_like, to_dtype(state.params))
+            # loop accum minibatch step over the number of gradient accumulation steps
+            grad, loss = jax.lax.scan(accum_minibatch_step, init_grad, batch)
+
+        # update state
+        new_state = state.apply_gradients(
+            grads=grad,
+            dropout_rng=new_dropout_rng,
+            to_dtype=to_dtype,
+        )
+
+        # compute gradient norms over all layers and globally for detailed monitoring
+        layer_grad_norm = jax.tree_util.tree_map(jnp.linalg.norm, grad)
+        logs = {
+            "layer_grad_norm": layer_grad_norm,
+            "grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm)),
+        }
+
+        # compute parameter norms over all layers and globally for detailed monitoring
+        layer_param_norm = jax.tree_util.tree_map(jnp.linalg.norm, new_state.params)
+        logs["layer_param_norm"] = layer_param_norm
+        logs["param_norm"] = jnp.linalg.norm(jax.tree_util.tree_leaves(layer_param_norm))
+
+        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+        metrics.update(logs)
+
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+
+        return new_state, metrics
+
+    # Define eval fn
+    def eval_step(params, batch):
+        labels = batch.pop("labels")
+        logits = model(**batch, params=params, train=False)[0]
+
+        logits_mask = model._get_feature_vector_attention_mask(logits.shape[1], batch["attention_mask"])
+        loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+        pred_ids = jnp.argmax(logits, axis=-1)
+
+        # summarize metrics
+        metrics = {"loss": loss}
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+        return metrics, pred_ids
+
+    # Create parallel version of the train and eval step
+    if training_args.do_train:
+        p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
+
+    if training_args.do_eval:
+        p_eval_step = jax.pmap(eval_step, "batch")
+
+    def run_evaluation(step):
+        if training_args.do_eval:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the eval dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[data_args.eval_split_name], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
+                samples = [vectorized_datasets[data_args.eval_split_name][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, step, prefix="eval")
+            write_wandb_pred(pred_str, label_str, step)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, step, pred_str=pred_str)
+
+    def save_checkpoint(step):
+        # save 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"{wandb.run.id}: saving weights and logs of step {int(step / 1000)}k", blocking=False)
+
+    skip_epochs = data_args.skip_steps // (num_train_samples // batch_size_per_update)
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_train_samples}")
+    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"  Num gradient accumulation steps = {gradient_accumulation_steps}")
+    logger.info(f"  Total train batch size (w. parallel & distributed) = {batch_size_per_update}")
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+    logger.info(f"  Gradient checkpointing: {config.gradient_checkpointing}")
+    logger.info(f"  Use scan: {config.use_scan}")
+    logger.info(f"  Fuse matmuls: {config.fuse_matmuls}")
+    logger.info(f"  Skipping: {data_args.skip_steps} steps ({skip_epochs} epochs)")
+
+    train_time = cur_step = 0
+    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
+    for epoch in epochs:
+        if training_args.do_train:
+            # ======================== Training ================================
+            train_start = time.time()
+
+            if epoch < skip_epochs:
+                logger.info(f"Skipping epoch {epoch + 1}")
+                continue
+
+            # Create sampling rng
+            rng, input_rng = jax.random.split(rng)
+
+            # Generate an epoch by randomly shuffling sampling indices from the train dataset and grouping by length
+            train_samples_idx = get_grouped_indices(vectorized_datasets[data_args.train_split_name], batch_size_per_update, input_rng)
+            train_batch_idx = generate_batch_splits(train_samples_idx, batch_size_per_update)
+
+            if data_args.skip_steps > cur_step:
+                logger.info(f"Skipping {data_args.skip_steps - (epoch * (num_train_samples // batch_size_per_update))} steps...")
+            # Gather the indices for creating the batch and do a training step
+            for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1), 1):
+                cur_step = epoch * (num_train_samples // batch_size_per_update) + step
+                if cur_step <= data_args.skip_steps:
+                    continue
+
+                samples = [vectorized_datasets[data_args.train_split_name][int(idx)] for idx in batch_idx]              
+                batch = data_collator(samples)
+                batch = shard(batch.data)
+                try:
+                    state, train_metric = p_train_step(state, batch)
+                except TypeError as e:
+                    logger.warning("Encountered following error: \n", e)
+
+
+                if cur_step % training_args.logging_steps == 0:
+                    # Save metrics
+                    train_metric = unreplicate(train_metric)
+                    train_time += time.time() - train_start
+                    # need to upcast all device arrays to fp32 for wandb logging (jnp.bfloat16 not supported) -> do this here OR in train_step
+                    write_wandb_log(to_fp32(train_metric), cur_step, prefix=data_args.train_split_name)
+                    # we won't log to tensorboard for now (it is fiddly logging param and grad norms on a layer-by-layer basis)
+                    # if has_tensorboard and jax.process_index() == 0:
+                    # write_train_metric(summary_writer, train_metrics, train_time, cur_step)
+
+                    epochs.write(
+                        f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']}, Gradient Norm: {train_metric['grad_norm']})"
+                    )
+
+                if cur_step % total_train_steps == 0:
+                    break
+
+                if training_args.eval_steps and cur_step % training_args.eval_steps == 0:
+                    run_evaluation(cur_step)
+
+                if cur_step % training_args.save_steps == 0:
+                    save_checkpoint(cur_step)
+
+            if training_args.eval_steps == 0 and (epoch + 1) != num_epochs:
+                # run evaluation at the end of the epoch if eval steps are not specified
+                run_evaluation(cur_step)
+                save_checkpoint(cur_step)
+
+    if training_args.do_train:
+        save_checkpoint(cur_step)
+
+    cur_step = max_steps if max_steps > 0 else cur_step  # set step to max steps so that eval happens in alignment with training
+
+    if training_args.do_eval:
+        run_evaluation(cur_step)
+
+    # TODO: collapse 'do_predict' into the run_evaluation function
+    if training_args.do_predict:
+        for split in [data_args.test_split_name]:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the test dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[split], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc=f"Predicting {split}...", position=2)):
+                samples = [vectorized_datasets[split][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({cur_step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, cur_step, prefix=split)
+            write_wandb_pred(pred_str, label_str, cur_step, prefix=split)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, cur_step, pred_str=pred_str)
+
+
+if __name__ == "__main__":
+    main()

wandb/run-20220810_111559-290849gb/files/config.yaml

@@ -0,0 +1,33 @@
+wandb_version: 1
+
+_wandb:
+  desc: null
+  value:
+    cli_version: 0.12.9
+    code_path: code/run_flax_speech_recognition_ctc.py
+    framework: huggingface
+    huggingface_version: 4.21.0
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    python_version: 3.8.10
+    start_time: 1660130159
+    t:
+      1:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      2:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      3:
+      - 13
+      4: 3.8.10
+      5: 0.12.9
+      6: 4.21.0
+      8:
+      - 5

wandb/run-20220810_111559-290849gb/files/diff.patch

@@ -0,0 +1,52 @@
+diff --git a/special_tokens_map.json b/special_tokens_map.json
+index 218961f..c11fc15 100644
+--- a/special_tokens_map.json
++++ b/special_tokens_map.json
+@@ -399,6 +399,20 @@
+       "rstrip": false,
+       "single_word": false
+     },
++    {
++      "content": "</s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "<s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
+     {
+       "content": "</s>",
+       "lstrip": false,
+diff --git a/wandb/debug-internal.log b/wandb/debug-internal.log
+index 23926ef..ad68f93 120000
+--- a/wandb/debug-internal.log
++++ b/wandb/debug-internal.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug-internal.log
+
++run-20220810_111559-290849gb/logs/debug-internal.log
+
+diff --git a/wandb/debug.log b/wandb/debug.log
+index 279853d..8db277f 120000
+--- a/wandb/debug.log
++++ b/wandb/debug.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug.log
+
++run-20220810_111559-290849gb/logs/debug.log
+
+diff --git a/wandb/latest-run b/wandb/latest-run
+index f069a7a..052e8bb 120000
+--- a/wandb/latest-run
++++ b/wandb/latest-run
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4
+
++run-20220810_111559-290849gb
+

wandb/run-20220810_111559-290849gb/files/output.log

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

wandb/run-20220810_111559-290849gb/files/requirements.txt

@@ -0,0 +1,158 @@
+absl-py==1.0.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+appdirs==1.4.4
+astunparse==1.6.3
+async-timeout==4.0.2
+attrs==21.4.0
+audioread==2.1.9
+backcall==0.2.0
+cachetools==4.2.4
+certifi==2021.10.8
+cffi==1.15.1
+charset-normalizer==2.0.10
+chex==0.1.3
+click==8.0.3
+cloud-tpu-client==0.10
+cloud-tpu-profiler==2.4.0
+clu==0.0.6
+colorama==0.4.5
+commonmark==0.9.1
+configparser==5.2.0
+contextlib2==21.6.0
+cycler==0.11.0
+datasets==2.4.0
+decorator==5.1.0
+dill==0.3.4
+dm-tree==0.1.6
+docker-pycreds==0.4.0
+etils==0.6.0
+exceptiongroup==1.0.0rc8
+filelock==3.4.2
+flatbuffers==2.0
+flax==0.5.3
+fonttools==4.28.5
+frozenlist==1.2.0
+fsspec==2021.11.1
+future==0.18.2
+gast==0.4.0
+gitdb==4.0.9
+gitpython==3.1.26
+google-api-core==1.31.5
+google-api-python-client==1.8.0
+google-auth-httplib2==0.1.0
+google-auth-oauthlib==0.4.6
+google-auth==2.3.3
+google-pasta==0.2.0
+googleapis-common-protos==1.54.0
+grpcio==1.43.0
+h5py==3.6.0
+httplib2==0.20.2
+huggingface-hub==0.2.1
+hypothesis==6.53.0
+idna==3.3
+importlib-metadata==4.10.0
+importlib-resources==5.4.0
+ipython==7.31.0
+jax==0.3.15
+jaxlib==0.3.15
+jedi==0.18.1
+jiwer==2.3.0
+joblib==1.1.0
+keras-preprocessing==1.1.2
+keras==2.7.0
+kiwisolver==1.3.2
+libclang==12.0.0
+librosa==0.9.2
+libtpu-nightly==0.1.dev20220722
+llvmlite==0.39.0
+markdown==3.3.6
+matplotlib-inline==0.1.3
+matplotlib==3.5.1
+ml-collections==0.1.0
+msgpack==1.0.3
+multidict==5.2.0
+multiprocess==0.70.12.2
+numba==0.56.0
+numpy==1.22.0
+oauth2client==4.1.3
+oauthlib==3.1.1
+opt-einsum==3.3.0
+optax==0.1.3
+packaging==21.3
+pandas==1.3.5
+parso==0.8.3
+pathtools==0.1.2
+pexpect==4.8.0
+pickleshare==0.7.5
+pillow==9.0.0
+pip==22.2.2
+pkg-resources==0.0.0
+pooch==1.6.0
+promise==2.3
+prompt-toolkit==3.0.24
+protobuf==3.19.1
+psutil==5.9.0
+ptyprocess==0.7.0
+pyarrow==6.0.1
+pyasn1-modules==0.2.8
+pyasn1==0.4.8
+pycparser==2.21
+pyctcdecode==0.4.0
+pygments==2.11.1
+pygtrie==2.5.0
+pyparsing==3.0.6
+python-dateutil==2.8.2
+python-levenshtein==0.12.2
+pytz==2021.3
+pyyaml==6.0
+regex==2021.11.10
+requests-oauthlib==1.3.0
+requests==2.27.0
+resampy==0.3.1
+responses==0.18.0
+rich==11.2.0
+rsa==4.8
+sacremoses==0.0.46
+scikit-learn==1.1.1
+scipy==1.7.3
+sentry-sdk==1.5.2
+setuptools==44.0.0
+shortuuid==1.0.8
+six==1.16.0
+smmap==5.0.0
+sortedcontainers==2.4.0
+soundfile==0.10.3.post1
+sox==1.4.1
+subprocess32==3.5.4
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.0
+tensorboard==2.7.0
+tensorflow-cpu==2.7.0
+tensorflow-datasets==4.4.0
+tensorflow-estimator==2.7.0
+tensorflow-io-gcs-filesystem==0.23.1
+tensorflow-metadata==1.5.0
+tensorflow==2.7.0
+tensorstore==0.1.21
+termcolor==1.1.0
+threadpoolctl==3.1.0
+tokenizers==0.11.2
+toolz==0.11.2
+torch==1.12.0
+torchaudio==0.12.0+cpu
+tqdm==4.62.3
+traitlets==5.1.1
+transformers==4.21.0
+typing-extensions==4.3.0
+uritemplate==3.0.1
+urllib3==1.26.7
+wandb==0.12.9
+wcwidth==0.2.5
+werkzeug==2.0.2
+wheel==0.37.1
+wrapt==1.13.3
+xxhash==2.0.2
+yarl==1.7.2
+yaspin==2.1.0
+zipp==3.7.0

wandb/run-20220810_111559-290849gb/files/wandb-metadata.json

@@ -0,0 +1,70 @@
+{
+    "os": "Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29",
+    "python": "3.8.10",
+    "heartbeatAt": "2022-08-10T11:16:02.847385",
+    "startedAt": "2022-08-10T11:15:59.241818",
+    "docker": null,
+    "cpu_count": 96,
+    "cuda": null,
+    "args": [
+        "--model_name_or_path=./",
+        "--hub_model_id=NbAiLab/wav2vec2-1b-npsc-nst-tpu",
+        "--tokenizer_name=./",
+        "--output_dir=./",
+        "--overwrite_output_dir",
+        "--num_train_epochs=40",
+        "--per_device_train_batch_size=2",
+        "--per_device_eval_batch_size=2",
+        "--gradient_accumulation_steps=1",
+        "--precision=full_mixed",
+        "--matmul_precision=bfloat16",
+        "--multisteps",
+        "--learning_rate=6.394633237505332e-05",
+        "--skip_steps=275000",
+        "--warmup_steps=2000",
+        "--length_column_name=input_length",
+        "--evaluation_strategy=steps",
+        "--text_column_name=text",
+        "--save_steps=5000",
+        "--eval_steps=5000",
+        "--logging_steps=100",
+        "--layerdrop=0.041",
+        "--attention_dropout=0.094",
+        "--activation_dropout=0.055",
+        "--hidden_dropout=0.047",
+        "--save_total_limit=5",
+        "--freeze_feature_encoder",
+        "--feat_proj_dropout=0.04",
+        "--mask_time_prob=0.082",
+        "--mask_time_length=10",
+        "--mask_feature_prob=0.25",
+        "--mask_feature_length=64",
+        "--gradient_checkpointing",
+        "--min_duration_in_seconds=0.5",
+        "--max_duration_in_seconds=30.0",
+        "--use_auth_token",
+        "--seed=42",
+        "--group_by_length",
+        "--do_train",
+        "--do_eval",
+        "--push_to_hub",
+        "--preprocessing_num_workers=32",
+        "--ctc_zero_infinity",
+        "--do_lower_case",
+        "--wandb_project=wav2vec2",
+        "--wandb_name=wav2vec2-1b-npsc-nst-tpu (cont.)",
+        "--remove_punctuation"
+    ],
+    "state": "running",
+    "program": "run_flax_speech_recognition_ctc.py",
+    "codePath": "run_flax_speech_recognition_ctc.py",
+    "git": {
+        "remote": "https://huggingface.co/NbAiLab/wav2vec2-1b-npsc-nst-tpu",
+        "commit": "f624ac4bfedfbf56891676d7c5f2e37b4c8e0745"
+    },
+    "email": "versae@gmail.com",
+    "root": "/data/wav2vec2-1b-npsc-nst-tpu",
+    "host": "t1v-n-eedfb410-w-0",
+    "username": "javierr",
+    "executable": "/data/flax/bin/python"
+}

wandb/run-20220810_111559-290849gb/files/wandb-summary.json

@@ -0,0 +1 @@
+{"train/grad_norm": 6.5625, "layer_grad_norm/": {"lm_head": {"bias": 0.031982421875, "kernel": 4.625}, "wav2vec2": {"encoder": {"layer_norm": {"bias": 0.0556640625, "scale": 0.06103515625}, "layers": {"0": {"attention": {"k_proj": {"bias": 5.9604644775390625e-05, "kernel": 0.023681640625}, "out_proj": {"bias": 0.04150390625, "kernel": 0.2431640625}, "q_proj": {"bias": 0.002899169921875, "kernel": 0.031005859375}, "v_proj": {"bias": 0.037109375, "kernel": 0.265625}}, "feed_forward": {"intermediate_dense": {"bias": 0.04443359375, "kernel": 0.515625}, "output_dense": {"bias": 0.0159912109375, "kernel": 0.439453125}}, "final_layer_norm": {"bias": 0.146484375, "scale": 0.322265625}, "layer_norm": {"bias": 0.0703125, "scale": 0.07080078125}}, "1": {"attention": {"k_proj": {"bias": 3.4332275390625e-05, "kernel": 0.03955078125}, "out_proj": {"bias": 0.0177001953125, "kernel": 0.134765625}, "q_proj": {"bias": 0.0035247802734375, "kernel": 0.0439453125}, "v_proj": {"bias": 0.02880859375, "kernel": 0.111328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.022216796875, "kernel": 0.3359375}, "output_dense": {"bias": 0.0157470703125, "kernel": 0.259765625}}, "final_layer_norm": {"bias": 0.046142578125, "scale": 0.05712890625}, "layer_norm": {"bias": 0.05712890625, "scale": 0.039794921875}}, "10": {"attention": {"k_proj": {"bias": 3.600120544433594e-05, "kernel": 0.123046875}, "out_proj": {"bias": 0.01416015625, "kernel": 0.2001953125}, "q_proj": {"bias": 0.0078125, "kernel": 0.12255859375}, "v_proj": {"bias": 0.022705078125, "kernel": 0.2001953125}}, "feed_forward": {"intermediate_dense": {"bias": 0.021484375, "kernel": 0.328125}, "output_dense": {"bias": 0.013671875, "kernel": 0.2734375}}, "final_layer_norm": {"bias": 0.0390625, "scale": 0.0322265625}, "layer_norm": {"bias": 0.04931640625, "scale": 0.0341796875}}, "11": {"attention": {"k_proj": {"bias": 8.344650268554688e-05, "kernel": 0.158203125}, "out_proj": {"bias": 0.0142822265625, "kernel": 0.28125}, "q_proj": {"bias": 0.0087890625, "kernel": 0.130859375}, "v_proj": {"bias": 0.024658203125, "kernel": 0.28515625}}, "feed_forward": {"intermediate_dense": {"bias": 0.01953125, "kernel": 0.310546875}, "output_dense": {"bias": 0.013916015625, "kernel": 0.244140625}}, "final_layer_norm": {"bias": 0.03271484375, "scale": 0.0308837890625}, "layer_norm": {"bias": 0.05029296875, "scale": 0.0439453125}}, "12": {"attention": {"k_proj": {"bias": 4.982948303222656e-05, "kernel": 0.123046875}, "out_proj": {"bias": 0.0147705078125, "kernel": 0.244140625}, "q_proj": {"bias": 0.0081787109375, "kernel": 0.1162109375}, "v_proj": {"bias": 0.023681640625, "kernel": 0.2294921875}}, "feed_forward": {"intermediate_dense": {"bias": 0.021484375, "kernel": 0.32421875}, "output_dense": {"bias": 0.014404296875, "kernel": 0.255859375}}, "final_layer_norm": {"bias": 0.0390625, "scale": 0.04248046875}, "layer_norm": {"bias": 0.046630859375, "scale": 0.0546875}}, "13": {"attention": {"k_proj": {"bias": 0.00012493133544921875, "kernel": 0.15625}, "out_proj": {"bias": 0.01519775390625, "kernel": 0.330078125}, "q_proj": {"bias": 0.0111083984375, "kernel": 0.158203125}, "v_proj": {"bias": 0.026611328125, "kernel": 0.3046875}}, "feed_forward": {"intermediate_dense": {"bias": 0.022216796875, "kernel": 0.34375}, "output_dense": {"bias": 0.01513671875, "kernel": 0.3125}}, "final_layer_norm": {"bias": 0.040283203125, "scale": 0.032958984375}, "layer_norm": {"bias": 0.051513671875, "scale": 0.091796875}}, "14": {"attention": {"k_proj": {"bias": 9.441375732421875e-05, "kernel": 0.1005859375}, "out_proj": {"bias": 0.015625, "kernel": 0.2412109375}, "q_proj": {"bias": 0.006256103515625, "kernel": 0.099609375}, "v_proj": {"bias": 0.0235595703125, "kernel": 0.2275390625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0257568359375, "kernel": 0.39453125}, "output_dense": {"bias": 0.015380859375, "kernel": 0.33984375}}, "final_layer_norm": {"bias": 0.05126953125, "scale": 0.05517578125}, "layer_norm": {"bias": 0.041748046875, "scale": 0.03076171875}}, "15": {"attention": {"k_proj": {"bias": 0.0003070831298828125, "kernel": 0.1806640625}, "out_proj": {"bias": 0.015625, "kernel": 0.5078125}, "q_proj": {"bias": 0.0106201171875, "kernel": 0.173828125}, "v_proj": {"bias": 0.026611328125, "kernel": 0.361328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.024169921875, "kernel": 0.376953125}, "output_dense": {"bias": 0.01556396484375, "kernel": 0.349609375}}, "final_layer_norm": {"bias": 0.041015625, "scale": 0.033447265625}, "layer_norm": {"bias": 0.048095703125, "scale": 0.072265625}}, "16": {"attention": {"k_proj": {"bias": 6.389617919921875e-05, "kernel": 0.1025390625}, "out_proj": {"bias": 0.016357421875, "kernel": 0.267578125}, "q_proj": {"bias": 0.0057373046875, "kernel": 0.1005859375}, "v_proj": {"bias": 0.0238037109375, "kernel": 0.220703125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0223388671875, "kernel": 0.359375}, "output_dense": {"bias": 0.0159912109375, "kernel": 0.341796875}}, "final_layer_norm": {"bias": 0.0380859375, "scale": 0.033935546875}, "layer_norm": {"bias": 0.043212890625, "scale": 0.034912109375}}, "17": {"attention": {"k_proj": {"bias": 4.57763671875e-05, "kernel": 0.0927734375}, "out_proj": {"bias": 0.0172119140625, "kernel": 0.23046875}, "q_proj": {"bias": 0.005889892578125, "kernel": 0.087890625}, "v_proj": {"bias": 0.0244140625, "kernel": 0.2177734375}}, "feed_forward": {"intermediate_dense": {"bias": 0.024169921875, "kernel": 0.390625}, "output_dense": {"bias": 0.01708984375, "kernel": 0.353515625}}, "final_layer_norm": {"bias": 0.041259765625, "scale": 0.036376953125}, "layer_norm": {"bias": 0.0439453125, "scale": 0.0341796875}}, "18": {"attention": {"k_proj": {"bias": 0.000247955322265625, "kernel": 0.126953125}, "out_proj": {"bias": 0.017578125, "kernel": 0.369140625}, "q_proj": {"bias": 0.0076904296875, "kernel": 0.1337890625}, "v_proj": {"bias": 0.027587890625, "kernel": 0.298828125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02587890625, "kernel": 0.44921875}, "output_dense": {"bias": 0.0174560546875, "kernel": 0.41015625}}, "final_layer_norm": {"bias": 0.04443359375, "scale": 0.03857421875}, "layer_norm": {"bias": 0.048583984375, "scale": 0.039794921875}}, "19": {"attention": {"k_proj": {"bias": 8.678436279296875e-05, "kernel": 0.140625}, "out_proj": {"bias": 0.017822265625, "kernel": 0.28125}, "q_proj": {"bias": 0.009033203125, "kernel": 0.140625}, "v_proj": {"bias": 0.0286865234375, "kernel": 0.283203125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02587890625, "kernel": 0.474609375}, "output_dense": {"bias": 0.0174560546875, "kernel": 0.421875}}, "final_layer_norm": {"bias": 0.041748046875, "scale": 0.0380859375}, "layer_norm": {"bias": 0.052734375, "scale": 0.04052734375}}, "2": {"attention": {"k_proj": {"bias": 4.982948303222656e-05, "kernel": 0.07421875}, "out_proj": {"bias": 0.0177001953125, "kernel": 0.2060546875}, "q_proj": {"bias": 0.006195068359375, "kernel": 0.06982421875}, "v_proj": {"bias": 0.03173828125, "kernel": 0.181640625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0234375, "kernel": 0.390625}, "output_dense": {"bias": 0.01556396484375, "kernel": 0.275390625}}, "final_layer_norm": {"bias": 0.047119140625, "scale": 0.03173828125}, "layer_norm": {"bias": 0.0556640625, "scale": 0.07275390625}}, "20": {"attention": {"k_proj": {"bias": 2.110004425048828e-05, "kernel": 0.095703125}, "out_proj": {"bias": 0.0185546875, "kernel": 0.142578125}, "q_proj": {"bias": 0.005157470703125, "kernel": 0.0947265625}, "v_proj": {"bias": 0.0263671875, "kernel": 0.140625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0250244140625, "kernel": 0.4765625}, "output_dense": {"bias": 0.018310546875, "kernel": 0.390625}}, "final_layer_norm": {"bias": 0.041015625, "scale": 0.04443359375}, "layer_norm": {"bias": 0.046142578125, "scale": 0.038330078125}}, "21": {"attention": {"k_proj": {"bias": 4.00543212890625e-05, "kernel": 0.1259765625}, "out_proj": {"bias": 0.0189208984375, "kernel": 0.2216796875}, "q_proj": {"bias": 0.006927490234375, "kernel": 0.12890625}, "v_proj": {"bias": 0.02734375, "kernel": 0.203125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0267333984375, "kernel": 0.51953125}, "output_dense": {"bias": 0.0185546875, "kernel": 0.41796875}}, "final_layer_norm": {"bias": 0.04541015625, "scale": 0.04736328125}, "layer_norm": {"bias": 0.044189453125, "scale": 0.054443359375}}, "22": {"attention": {"k_proj": {"bias": 3.3855438232421875e-05, "kernel": 0.1181640625}, "out_proj": {"bias": 0.019775390625, "kernel": 0.240234375}, "q_proj": {"bias": 0.006011962890625, "kernel": 0.11279296875}, "v_proj": {"bias": 0.028076171875, "kernel": 0.21875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0269775390625, "kernel": 0.515625}, "output_dense": {"bias": 0.0194091796875, "kernel": 0.400390625}}, "final_layer_norm": {"bias": 0.046142578125, "scale": 0.047119140625}, "layer_norm": {"bias": 0.049560546875, "scale": 0.0458984375}}, "23": {"attention": {"k_proj": {"bias": 0.0001087188720703125, "kernel": 0.16015625}, "out_proj": {"bias": 0.0198974609375, "kernel": 0.443359375}, "q_proj": {"bias": 0.008544921875, "kernel": 0.1630859375}, "v_proj": {"bias": 0.03173828125, "kernel": 0.35546875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0263671875, "kernel": 0.53125}, "output_dense": {"bias": 0.01953125, "kernel": 0.400390625}}, "final_layer_norm": {"bias": 0.044677734375, "scale": 0.04638671875}, "layer_norm": {"bias": 0.05615234375, "scale": 0.056396484375}}, "24": {"attention": {"k_proj": {"bias": 6.246566772460938e-05, "kernel": 0.169921875}, "out_proj": {"bias": 0.0191650390625, "kernel": 0.36328125}, "q_proj": {"bias": 0.00933837890625, "kernel": 0.18359375}, "v_proj": {"bias": 0.03271484375, "kernel": 0.328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02685546875, "kernel": 0.5390625}, "output_dense": {"bias": 0.01904296875, "kernel": 0.37890625}}, "final_layer_norm": {"bias": 0.04736328125, "scale": 0.04345703125}, "layer_norm": {"bias": 0.0625, "scale": 0.041015625}}, "25": {"attention": {"k_proj": {"bias": 6.079673767089844e-05, "kernel": 0.169921875}, "out_proj": {"bias": 0.0196533203125, "kernel": 0.3125}, "q_proj": {"bias": 0.00860595703125, "kernel": 0.16015625}, "v_proj": {"bias": 0.03271484375, "kernel": 0.32421875}}, "feed_forward": {"intermediate_dense": {"bias": 0.02734375, "kernel": 0.55859375}, "output_dense": {"bias": 0.01953125, "kernel": 0.375}}, "final_layer_norm": {"bias": 0.050537109375, "scale": 0.0478515625}, "layer_norm": {"bias": 0.06005859375, "scale": 0.06298828125}}, "26": {"attention": {"k_proj": {"bias": 5.9604644775390625e-05, "kernel": 0.169921875}, "out_proj": {"bias": 0.01953125, "kernel": 0.29296875}, "q_proj": {"bias": 0.01025390625, "kernel": 0.177734375}, "v_proj": {"bias": 0.0341796875, "kernel": 0.296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.026611328125, "kernel": 0.51171875}, "output_dense": {"bias": 0.01904296875, "kernel": 0.353515625}}, "final_layer_norm": {"bias": 0.0478515625, "scale": 0.04443359375}, "layer_norm": {"bias": 0.060791015625, "scale": 0.06396484375}}, "27": {"attention": {"k_proj": {"bias": 0.00011396408081054688, "kernel": 0.2021484375}, "out_proj": {"bias": 0.01806640625, "kernel": 0.44921875}, "q_proj": {"bias": 0.01068115234375, "kernel": 0.2138671875}, "v_proj": {"bias": 0.03466796875, "kernel": 0.435546875}}, "feed_forward": {"intermediate_dense": {"bias": 0.02734375, "kernel": 0.515625}, "output_dense": {"bias": 0.0181884765625, "kernel": 0.36328125}}, "final_layer_norm": {"bias": 0.05078125, "scale": 0.045654296875}, "layer_norm": {"bias": 0.06640625, "scale": 0.04931640625}}, "28": {"attention": {"k_proj": {"bias": 0.0001049041748046875, "kernel": 0.20703125}, "out_proj": {"bias": 0.0164794921875, "kernel": 0.392578125}, "q_proj": {"bias": 0.01165771484375, "kernel": 0.208984375}, "v_proj": {"bias": 0.031494140625, "kernel": 0.404296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0234375, "kernel": 0.45703125}, "output_dense": {"bias": 0.016357421875, "kernel": 0.326171875}}, "final_layer_norm": {"bias": 0.04248046875, "scale": 0.044921875}, "layer_norm": {"bias": 0.0673828125, "scale": 0.08447265625}}, "29": {"attention": {"k_proj": {"bias": 9.918212890625e-05, "kernel": 0.267578125}, "out_proj": {"bias": 0.0157470703125, "kernel": 0.28515625}, "q_proj": {"bias": 0.01495361328125, "kernel": 0.265625}, "v_proj": {"bias": 0.02978515625, "kernel": 0.31640625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0234375, "kernel": 0.494140625}, "output_dense": {"bias": 0.01531982421875, "kernel": 0.296875}}, "final_layer_norm": {"bias": 0.03955078125, "scale": 0.03515625}, "layer_norm": {"bias": 0.0654296875, "scale": 0.061279296875}}, "3": {"attention": {"k_proj": {"bias": 0.00012111663818359375, "kernel": 0.0986328125}, "out_proj": {"bias": 0.016845703125, "kernel": 0.314453125}, "q_proj": {"bias": 0.00726318359375, "kernel": 0.0888671875}, "v_proj": {"bias": 0.0283203125, "kernel": 0.2470703125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0242919921875, "kernel": 0.3828125}, "output_dense": {"bias": 0.0150146484375, "kernel": 0.271484375}}, "final_layer_norm": {"bias": 0.0458984375, "scale": 0.03125}, "layer_norm": {"bias": 0.0498046875, "scale": 0.0380859375}}, "30": {"attention": {"k_proj": {"bias": 0.0001220703125, "kernel": 0.13671875}, "out_proj": {"bias": 0.0152587890625, "kernel": 0.328125}, "q_proj": {"bias": 0.006378173828125, "kernel": 0.138671875}, "v_proj": {"bias": 0.029296875, "kernel": 0.3671875}}, "feed_forward": {"intermediate_dense": {"bias": 0.023681640625, "kernel": 0.51953125}, "output_dense": {"bias": 0.01446533203125, "kernel": 0.26171875}}, "final_layer_norm": {"bias": 0.04150390625, "scale": 0.03564453125}, "layer_norm": {"bias": 0.04931640625, "scale": 0.037109375}}, "31": {"attention": {"k_proj": {"bias": 0.00010347366333007812, "kernel": 0.14453125}, "out_proj": {"bias": 0.0140380859375, "kernel": 0.29296875}, "q_proj": {"bias": 0.006378173828125, "kernel": 0.134765625}, "v_proj": {"bias": 0.0238037109375, "kernel": 0.314453125}}, "feed_forward": {"intermediate_dense": {"bias": 0.02392578125, "kernel": 0.51953125}, "output_dense": {"bias": 0.01385498046875, "kernel": 0.2578125}}, "final_layer_norm": {"bias": 0.0390625, "scale": 0.03662109375}, "layer_norm": {"bias": 0.039306640625, "scale": 0.0291748046875}}, "32": {"attention": {"k_proj": {"bias": 8.296966552734375e-05, "kernel": 0.15625}, "out_proj": {"bias": 0.01263427734375, "kernel": 0.28125}, "q_proj": {"bias": 0.0079345703125, "kernel": 0.1533203125}, "v_proj": {"bias": 0.0264892578125, "kernel": 0.4921875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0216064453125, "kernel": 0.431640625}, "output_dense": {"bias": 0.01129150390625, "kernel": 0.212890625}}, "final_layer_norm": {"bias": 0.04150390625, "scale": 0.03271484375}, "layer_norm": {"bias": 0.046630859375, "scale": 0.05419921875}}, "33": {"attention": {"k_proj": {"bias": 5.9604644775390625e-05, "kernel": 0.166015625}, "out_proj": {"bias": 0.01092529296875, "kernel": 0.2275390625}, "q_proj": {"bias": 0.008544921875, "kernel": 0.166015625}, "v_proj": {"bias": 0.023193359375, "kernel": 0.34765625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0196533203125, "kernel": 0.390625}, "output_dense": {"bias": 0.00897216796875, "kernel": 0.1875}}, "final_layer_norm": {"bias": 0.04345703125, "scale": 0.0361328125}, "layer_norm": {"bias": 0.039794921875, "scale": 0.0498046875}}, "34": {"attention": {"k_proj": {"bias": 0.0002346038818359375, "kernel": 0.158203125}, "out_proj": {"bias": 0.0081787109375, "kernel": 0.181640625}, "q_proj": {"bias": 0.006927490234375, "kernel": 0.14453125}, "v_proj": {"bias": 0.0177001953125, "kernel": 0.25390625}}, "feed_forward": {"intermediate_dense": {"bias": 0.01434326171875, "kernel": 0.291015625}, "output_dense": {"bias": 0.0072021484375, "kernel": 0.1748046875}}, "final_layer_norm": {"bias": 0.028076171875, "scale": 0.025146484375}, "layer_norm": {"bias": 0.03369140625, "scale": 0.026611328125}}, "35": {"attention": {"k_proj": {"bias": 0.0001506805419921875, "kernel": 0.10791015625}, "out_proj": {"bias": 0.00640869140625, "kernel": 0.2109375}, "q_proj": {"bias": 0.004852294921875, "kernel": 0.10791015625}, "v_proj": {"bias": 0.01177978515625, "kernel": 0.21484375}}, "feed_forward": {"intermediate_dense": {"bias": 0.010498046875, "kernel": 0.2119140625}, "output_dense": {"bias": 0.005889892578125, "kernel": 0.15234375}}, "final_layer_norm": {"bias": 0.0206298828125, "scale": 0.0220947265625}, "layer_norm": {"bias": 0.024169921875, "scale": 0.02880859375}}, "36": {"attention": {"k_proj": {"bias": 4.410743713378906e-05, "kernel": 0.1005859375}, "out_proj": {"bias": 0.005645751953125, "kernel": 0.1552734375}, "q_proj": {"bias": 0.00445556640625, "kernel": 0.095703125}, "v_proj": {"bias": 0.00946044921875, "kernel": 0.14453125}}, "feed_forward": {"intermediate_dense": {"bias": 0.0089111328125, "kernel": 0.177734375}, "output_dense": {"bias": 0.0050048828125, "kernel": 0.111328125}}, "final_layer_norm": {"bias": 0.017578125, "scale": 0.01513671875}, "layer_norm": {"bias": 0.0191650390625, "scale": 0.01806640625}}, "37": {"attention": {"k_proj": {"bias": 9.441375732421875e-05, "kernel": 0.0849609375}, "out_proj": {"bias": 0.004913330078125, "kernel": 0.11474609375}, "q_proj": {"bias": 0.00390625, "kernel": 0.0830078125}, "v_proj": {"bias": 0.00897216796875, "kernel": 0.1318359375}}, "feed_forward": {"intermediate_dense": {"bias": 0.00823974609375, "kernel": 0.16796875}, "output_dense": {"bias": 0.004241943359375, "kernel": 0.09716796875}}, "final_layer_norm": {"bias": 0.015869140625, "scale": 0.01434326171875}, "layer_norm": {"bias": 0.019287109375, "scale": 0.015869140625}}, "38": {"attention": {"k_proj": {"bias": 5.650520324707031e-05, "kernel": 0.09130859375}, "out_proj": {"bias": 0.0040283203125, "kernel": 0.11865234375}, "q_proj": {"bias": 0.00396728515625, "kernel": 0.08642578125}, "v_proj": {"bias": 0.007354736328125, "kernel": 0.1279296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0072021484375, "kernel": 0.150390625}, "output_dense": {"bias": 0.0034637451171875, "kernel": 0.09423828125}}, "final_layer_norm": {"bias": 0.0152587890625, "scale": 0.0146484375}, "layer_norm": {"bias": 0.0162353515625, "scale": 0.0135498046875}}, "39": {"attention": {"k_proj": {"bias": 5.316734313964844e-05, "kernel": 0.09619140625}, "out_proj": {"bias": 0.0030975341796875, "kernel": 0.09619140625}, "q_proj": {"bias": 0.00408935546875, "kernel": 0.0908203125}, "v_proj": {"bias": 0.006011962890625, "kernel": 0.10986328125}}, "feed_forward": {"intermediate_dense": {"bias": 0.005401611328125, "kernel": 0.12109375}, "output_dense": {"bias": 0.0025634765625, "kernel": 0.08642578125}}, "final_layer_norm": {"bias": 0.01202392578125, "scale": 0.01226806640625}, "layer_norm": {"bias": 0.0150146484375, "scale": 0.01556396484375}}, "4": {"attention": {"k_proj": {"bias": 0.000148773193359375, "kernel": 0.10498046875}, "out_proj": {"bias": 0.015869140625, "kernel": 0.361328125}, "q_proj": {"bias": 0.0072021484375, "kernel": 0.1005859375}, "v_proj": {"bias": 0.026123046875, "kernel": 0.3046875}}, "feed_forward": {"intermediate_dense": {"bias": 0.023193359375, "kernel": 0.36328125}, "output_dense": {"bias": 0.014404296875, "kernel": 0.29296875}}, "final_layer_norm": {"bias": 0.042724609375, "scale": 0.034423828125}, "layer_norm": {"bias": 0.0478515625, "scale": 0.060546875}}, "40": {"attention": {"k_proj": {"bias": 5.269050598144531e-05, "kernel": 0.046875}, "out_proj": {"bias": 0.0025787353515625, "kernel": 0.080078125}, "q_proj": {"bias": 0.0020294189453125, "kernel": 0.0458984375}, "v_proj": {"bias": 0.004150390625, "kernel": 0.07080078125}}, "feed_forward": {"intermediate_dense": {"bias": 0.004302978515625, "kernel": 0.09326171875}, "output_dense": {"bias": 0.0023040771484375, "kernel": 0.060791015625}}, "final_layer_norm": {"bias": 0.0087890625, "scale": 0.011474609375}, "layer_norm": {"bias": 0.00823974609375, "scale": 0.007781982421875}}, "41": {"attention": {"k_proj": {"bias": 4.0531158447265625e-05, "kernel": 0.0673828125}, "out_proj": {"bias": 0.002044677734375, "kernel": 0.087890625}, "q_proj": {"bias": 0.0025787353515625, "kernel": 0.0634765625}, "v_proj": {"bias": 0.00439453125, "kernel": 0.1044921875}}, "feed_forward": {"intermediate_dense": {"bias": 0.0035858154296875, "kernel": 0.091796875}, "output_dense": {"bias": 0.00168609619140625, "kernel": 0.064453125}}, "final_layer_norm": {"bias": 0.00921630859375, "scale": 0.0106201171875}, "layer_norm": {"bias": 0.010986328125, "scale": 0.0128173828125}}, "42": {"attention": {"k_proj": {"bias": 1.1801719665527344e-05, "kernel": 0.02099609375}, "out_proj": {"bias": 0.001678466796875, "kernel": 0.048828125}, "q_proj": {"bias": 0.0009307861328125, "kernel": 0.02197265625}, "v_proj": {"bias": 0.002349853515625, "kernel": 0.0478515625}}, "feed_forward": {"intermediate_dense": {"bias": 0.002685546875, "kernel": 0.07421875}, "output_dense": {"bias": 0.0014801025390625, "kernel": 0.053955078125}}, "final_layer_norm": {"bias": 0.0059814453125, "scale": 0.00885009765625}, "layer_norm": {"bias": 0.0045166015625, "scale": 0.00439453125}}, "43": {"attention": {"k_proj": {"bias": 8.046627044677734e-06, "kernel": 0.01806640625}, "out_proj": {"bias": 0.0015106201171875, "kernel": 0.035888671875}, "q_proj": {"bias": 0.00095367431640625, "kernel": 0.0208740234375}, "v_proj": {"bias": 0.00171661376953125, "kernel": 0.031005859375}}, "feed_forward": {"intermediate_dense": {"bias": 0.002777099609375, "kernel": 0.083984375}, "output_dense": {"bias": 0.00125885009765625, "kernel": 0.052734375}}, "final_layer_norm": {"bias": 0.00689697265625, "scale": 0.007110595703125}, "layer_norm": {"bias": 0.0036468505859375, "scale": 0.005706787109375}}, "44": {"attention": {"k_proj": {"bias": 1.3113021850585938e-05, "kernel": 0.023681640625}, "out_proj": {"bias": 0.001312255859375, "kernel": 0.033447265625}, "q_proj": {"bias": 0.000946044921875, "kernel": 0.021484375}, "v_proj": {"bias": 0.0017547607421875, "kernel": 0.03515625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0023956298828125, "kernel": 0.0771484375}, "output_dense": {"bias": 0.001129150390625, "kernel": 0.052001953125}}, "final_layer_norm": {"bias": 0.005706787109375, "scale": 0.005859375}, "layer_norm": {"bias": 0.0042724609375, "scale": 0.004730224609375}}, "45": {"attention": {"k_proj": {"bias": 1.4424324035644531e-05, "kernel": 0.01239013671875}, "out_proj": {"bias": 0.00110626220703125, "kernel": 0.0267333984375}, "q_proj": {"bias": 0.00136566162109375, "kernel": 0.029541015625}, "v_proj": {"bias": 0.00142669677734375, "kernel": 0.027587890625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0019989013671875, "kernel": 0.06396484375}, "output_dense": {"bias": 0.000881195068359375, "kernel": 0.05615234375}}, "final_layer_norm": {"bias": 0.006072998046875, "scale": 0.006591796875}, "layer_norm": {"bias": 0.004791259765625, "scale": 0.00494384765625}}, "46": {"attention": {"k_proj": {"bias": 5.745887756347656e-05, "kernel": 0.006439208984375}, "out_proj": {"bias": 0.000888824462890625, "kernel": 0.0289306640625}, "q_proj": {"bias": 0.000591278076171875, "kernel": 0.011962890625}, "v_proj": {"bias": 0.0010986328125, "kernel": 0.0233154296875}}, "feed_forward": {"intermediate_dense": {"bias": 0.00141143798828125, "kernel": 0.03955078125}, "output_dense": {"bias": 0.000873565673828125, "kernel": 0.0478515625}}, "final_layer_norm": {"bias": 0.00433349609375, "scale": 0.00433349609375}, "layer_norm": {"bias": 0.0036468505859375, "scale": 0.003814697265625}}, "47": {"attention": {"k_proj": {"bias": 0.00011301040649414062, "kernel": 0.003997802734375}, "out_proj": {"bias": 0.000896453857421875, "kernel": 0.06640625}, "q_proj": {"bias": 0.00014591217041015625, "kernel": 0.00286865234375}, "v_proj": {"bias": 0.00118255615234375, "kernel": 0.0230712890625}}, "feed_forward": {"intermediate_dense": {"bias": 0.0010528564453125, "kernel": 0.0252685546875}, "output_dense": {"bias": 0.000881195068359375, "kernel": 0.1787109375}}, "final_layer_norm": {"bias": 0.005950927734375, "scale": 0.00677490234375}, "layer_norm": {"bias": 0.005859375, "scale": 0.005767822265625}}, "5": {"attention": {"k_proj": {"bias": 6.4849853515625e-05, "kernel": 0.123046875}, "out_proj": {"bias": 0.0159912109375, "kernel": 0.203125}, "q_proj": {"bias": 0.007598876953125, "kernel": 0.12158203125}, "v_proj": {"bias": 0.02685546875, "kernel": 0.1953125}}, "feed_forward": {"intermediate_dense": {"bias": 0.023193359375, "kernel": 0.33984375}, "output_dense": {"bias": 0.0147705078125, "kernel": 0.275390625}}, "final_layer_norm": {"bias": 0.041015625, "scale": 0.037841796875}, "layer_norm": {"bias": 0.052734375, "scale": 0.04736328125}}, "6": {"attention": {"k_proj": {"bias": 7.152557373046875e-05, "kernel": 0.1318359375}, "out_proj": {"bias": 0.0152587890625, "kernel": 0.349609375}, "q_proj": {"bias": 0.00823974609375, "kernel": 0.119140625}, "v_proj": {"bias": 0.02685546875, "kernel": 0.31640625}}, "feed_forward": {"intermediate_dense": {"bias": 0.022216796875, "kernel": 0.341796875}, "output_dense": {"bias": 0.014404296875, "kernel": 0.271484375}}, "final_layer_norm": {"bias": 0.0380859375, "scale": 0.0322265625}, "layer_norm": {"bias": 0.049560546875, "scale": 0.058349609375}}, "7": {"attention": {"k_proj": {"bias": 7.62939453125e-05, "kernel": 0.1328125}, "out_proj": {"bias": 0.0150146484375, "kernel": 0.349609375}, "q_proj": {"bias": 0.00921630859375, "kernel": 0.126953125}, "v_proj": {"bias": 0.0255126953125, "kernel": 0.30859375}}, "feed_forward": {"intermediate_dense": {"bias": 0.021484375, "kernel": 0.33984375}, "output_dense": {"bias": 0.01409912109375, "kernel": 0.26171875}}, "final_layer_norm": {"bias": 0.038818359375, "scale": 0.033935546875}, "layer_norm": {"bias": 0.05126953125, "scale": 0.050048828125}}, "8": {"attention": {"k_proj": {"bias": 7.915496826171875e-05, "kernel": 0.12060546875}, "out_proj": {"bias": 0.01507568359375, "kernel": 0.302734375}, "q_proj": {"bias": 0.007568359375, "kernel": 0.11474609375}, "v_proj": {"bias": 0.0262451171875, "kernel": 0.27734375}}, "feed_forward": {"intermediate_dense": {"bias": 0.023193359375, "kernel": 0.361328125}, "output_dense": {"bias": 0.01409912109375, "kernel": 0.275390625}}, "final_layer_norm": {"bias": 0.044677734375, "scale": 0.033447265625}, "layer_norm": {"bias": 0.0498046875, "scale": 0.06103515625}}, "9": {"attention": {"k_proj": {"bias": 0.00011777877807617188, "kernel": 0.1513671875}, "out_proj": {"bias": 0.0140380859375, "kernel": 0.416015625}, "q_proj": {"bias": 0.00830078125, "kernel": 0.1376953125}, "v_proj": {"bias": 0.0238037109375, "kernel": 0.40234375}}, "feed_forward": {"intermediate_dense": {"bias": 0.0213623046875, "kernel": 0.3515625}, "output_dense": {"bias": 0.0135498046875, "kernel": 0.279296875}}, "final_layer_norm": {"bias": 0.037109375, "scale": 0.0322265625}, "layer_norm": {"bias": 0.04443359375, "scale": 0.04833984375}}}, "pos_conv_embed": {"conv": {"bias": 0.034912109375, "weight_g": 0.044189453125, "weight_v": 0.287109375}}}, "feature_extractor": {"conv_layers": {"0": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "1": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "2": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "3": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "4": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "5": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}, "6": {"conv": {"bias": 0.0, "kernel": 0.0}, "layer_norm": {"bias": 0.0, "scale": 0.0}}}}, "feature_projection": {"layer_norm": {"bias": 0.1337890625, "scale": 0.1611328125}, "projection": {"bias": 0.0556640625, "kernel": 1.0546875}}, "masked_spec_embed": 0.0}}, "layer_param_norm/": {"lm_head": {"bias": 0.7824921607971191, "kernel": 55.72966766357422}, "wav2vec2": {"encoder": {"layer_norm": {"bias": 59.6768798828125, "scale": 74.17054748535156}, "layers": {"0": {"attention": {"k_proj": {"bias": 0.37033993005752563, "kernel": 27.536663055419922}, "out_proj": {"bias": 1.6469175815582275, "kernel": 26.147050857543945}, "q_proj": {"bias": 1.5330281257629395, "kernel": 27.813282012939453}, "v_proj": {"bias": 0.44783300161361694, "kernel": 26.55841064453125}}, "feed_forward": {"intermediate_dense": {"bias": 1.9835489988327026, "kernel": 100.66567993164062}, "output_dense": {"bias": 1.116748571395874, "kernel": 96.76679992675781}}, "final_layer_norm": {"bias": 1.335214376449585, "scale": 19.85782241821289}, "layer_norm": {"bias": 2.923041343688965, "scale": 15.398418426513672}}, "1": {"attention": {"k_proj": {"bias": 0.3767347037792206, "kernel": 41.013240814208984}, "out_proj": {"bias": 1.3653483390808105, "kernel": 43.371070861816406}, "q_proj": {"bias": 3.0925614833831787, "kernel": 41.05661392211914}, "v_proj": {"bias": 0.2924947738647461, "kernel": 41.61189270019531}}, "feed_forward": {"intermediate_dense": {"bias": 1.9882662296295166, "kernel": 98.7442626953125}, "output_dense": {"bias": 0.8527815341949463, "kernel": 87.83541870117188}}, "final_layer_norm": {"bias": 1.3605934381484985, "scale": 19.084806442260742}, "layer_norm": {"bias": 1.9318764209747314, "scale": 17.761367797851562}}, "10": {"attention": {"k_proj": {"bias": 0.4123449921607971, "kernel": 49.44670486450195}, "out_proj": {"bias": 1.3130683898925781, "kernel": 52.27025604248047}, "q_proj": {"bias": 2.48445200920105, "kernel": 49.528873443603516}, "v_proj": {"bias": 0.3416975736618042, "kernel": 52.344085693359375}}, "feed_forward": {"intermediate_dense": {"bias": 1.974550485610962, "kernel": 102.70410919189453}, "output_dense": {"bias": 0.5955485105514526, "kernel": 95.81275939941406}}, "final_layer_norm": {"bias": 2.3762400150299072, "scale": 20.81279754638672}, "layer_norm": {"bias": 1.806241512298584, "scale": 21.429487228393555}}, "11": {"attention": {"k_proj": {"bias": 0.4460787773132324, "kernel": 49.37338638305664}, "out_proj": {"bias": 1.1512463092803955, "kernel": 51.949554443359375}, "q_proj": {"bias": 2.5446064472198486, "kernel": 49.20353698730469}, "v_proj": {"bias": 0.40995872020721436, "kernel": 52.2607536315918}}, "feed_forward": {"intermediate_dense": {"bias": 2.019528388977051, "kernel": 103.56025695800781}, "output_dense": {"bias": 0.5711302757263184, "kernel": 97.53776550292969}}, "final_layer_norm": {"bias": 2.3660812377929688, "scale": 20.919677734375}, "layer_norm": {"bias": 1.7802445888519287, "scale": 22.01519203186035}}, "12": {"attention": {"k_proj": {"bias": 0.4312320351600647, "kernel": 50.07032775878906}, "out_proj": {"bias": 1.126122236251831, "kernel": 51.988826751708984}, "q_proj": {"bias": 2.4090728759765625, "kernel": 49.92729949951172}, "v_proj": {"bias": 0.4024103879928589, "kernel": 52.296756744384766}}, "feed_forward": {"intermediate_dense": {"bias": 2.0548508167266846, "kernel": 104.54823303222656}, "output_dense": {"bias": 0.5548778772354126, "kernel": 99.2693099975586}}, "final_layer_norm": {"bias": 2.2933573722839355, "scale": 20.85626983642578}, "layer_norm": {"bias": 1.8587299585342407, "scale": 22.473487854003906}}, "13": {"attention": {"k_proj": {"bias": 0.4430793821811676, "kernel": 51.76431655883789}, "out_proj": {"bias": 1.1271920204162598, "kernel": 51.86852264404297}, "q_proj": {"bias": 2.359200954437256, "kernel": 51.75225830078125}, "v_proj": {"bias": 0.3906242251396179, "kernel": 51.92781066894531}}, "feed_forward": {"intermediate_dense": {"bias": 2.0941619873046875, "kernel": 105.30806732177734}, "output_dense": {"bias": 0.5719542503356934, "kernel": 99.8712158203125}}, "final_layer_norm": {"bias": 2.2314066886901855, "scale": 21.027400970458984}, "layer_norm": {"bias": 1.9997800588607788, "scale": 22.84510040283203}}, "14": {"attention": {"k_proj": {"bias": 0.43604815006256104, "kernel": 51.92181396484375}, "out_proj": {"bias": 1.2681217193603516, "kernel": 49.762760162353516}, "q_proj": {"bias": 2.4942922592163086, "kernel": 52.049808502197266}, "v_proj": {"bias": 0.36820662021636963, "kernel": 49.26283264160156}}, "feed_forward": {"intermediate_dense": {"bias": 2.1335830688476562, "kernel": 105.94457244873047}, "output_dense": {"bias": 0.6063626408576965, "kernel": 101.24859619140625}}, "final_layer_norm": {"bias": 2.2754664421081543, "scale": 21.145992279052734}, "layer_norm": {"bias": 2.1295526027679443, "scale": 22.584672927856445}}, "15": {"attention": {"k_proj": {"bias": 0.45931702852249146, "kernel": 51.93058776855469}, "out_proj": {"bias": 1.3753983974456787, "kernel": 50.94449234008789}, "q_proj": {"bias": 2.5871663093566895, "kernel": 52.099769592285156}, "v_proj": {"bias": 0.459650456905365, "kernel": 50.5831298828125}}, "feed_forward": {"intermediate_dense": {"bias": 2.132938861846924, "kernel": 105.57443237304688}, "output_dense": {"bias": 0.7701732516288757, "kernel": 101.94094848632812}}, "final_layer_norm": {"bias": 2.327320098876953, "scale": 21.192947387695312}, "layer_norm": {"bias": 2.4148712158203125, "scale": 23.526634216308594}}, "16": {"attention": {"k_proj": {"bias": 0.4008745551109314, "kernel": 51.772621154785156}, "out_proj": {"bias": 1.27531099319458, "kernel": 50.134521484375}, "q_proj": {"bias": 2.667466163635254, "kernel": 51.75814437866211}, "v_proj": {"bias": 0.3768249750137329, "kernel": 49.78093719482422}}, "feed_forward": {"intermediate_dense": {"bias": 2.1094985008239746, "kernel": 106.10227966308594}, "output_dense": {"bias": 0.7860437631607056, "kernel": 102.66590881347656}}, "final_layer_norm": {"bias": 2.337951421737671, "scale": 21.583194732666016}, "layer_norm": {"bias": 2.283249855041504, "scale": 22.168060302734375}}, "17": {"attention": {"k_proj": {"bias": 0.3966267704963684, "kernel": 51.728885650634766}, "out_proj": {"bias": 1.2151354551315308, "kernel": 49.4556884765625}, "q_proj": {"bias": 2.714320182800293, "kernel": 51.81880187988281}, "v_proj": {"bias": 0.42661017179489136, "kernel": 49.11927032470703}}, "feed_forward": {"intermediate_dense": {"bias": 2.1101765632629395, "kernel": 107.14872741699219}, "output_dense": {"bias": 0.8218655586242676, "kernel": 103.06423950195312}}, "final_layer_norm": {"bias": 2.383938789367676, "scale": 22.070323944091797}, "layer_norm": {"bias": 2.222898483276367, "scale": 21.219982147216797}}, "18": {"attention": {"k_proj": {"bias": 0.4409676194190979, "kernel": 52.41611099243164}, "out_proj": {"bias": 1.3447906970977783, "kernel": 50.491905212402344}, "q_proj": {"bias": 2.614685535430908, "kernel": 52.796600341796875}, "v_proj": {"bias": 0.4518332779407501, "kernel": 50.001895904541016}}, "feed_forward": {"intermediate_dense": {"bias": 2.144195556640625, "kernel": 107.41338348388672}, "output_dense": {"bias": 0.9481453895568848, "kernel": 104.72514343261719}}, "final_layer_norm": {"bias": 2.5390124320983887, "scale": 22.15178680419922}, "layer_norm": {"bias": 2.424910068511963, "scale": 23.585906982421875}}, "19": {"attention": {"k_proj": {"bias": 0.38193291425704956, "kernel": 51.511146545410156}, "out_proj": {"bias": 1.3303101062774658, "kernel": 50.10035705566406}, "q_proj": {"bias": 2.930327892303467, "kernel": 51.865638732910156}, "v_proj": {"bias": 0.4086824655532837, "kernel": 49.38078308105469}}, "feed_forward": {"intermediate_dense": {"bias": 2.1912901401519775, "kernel": 107.95254516601562}, "output_dense": {"bias": 1.0248571634292603, "kernel": 105.65098571777344}}, "final_layer_norm": {"bias": 2.4923481941223145, "scale": 22.505674362182617}, "layer_norm": {"bias": 2.2888314723968506, "scale": 22.31826400756836}}, "2": {"attention": {"k_proj": {"bias": 0.454792320728302, "kernel": 47.77275085449219}, "out_proj": {"bias": 1.256988525390625, "kernel": 45.969764709472656}, "q_proj": {"bias": 3.2510807514190674, "kernel": 47.61664581298828}, "v_proj": {"bias": 0.339598685503006, "kernel": 45.72273254394531}}, "feed_forward": {"intermediate_dense": {"bias": 1.9737317562103271, "kernel": 103.32754516601562}, "output_dense": {"bias": 0.7398276329040527, "kernel": 91.11263275146484}}, "final_layer_norm": {"bias": 1.5421981811523438, "scale": 21.561111450195312}, "layer_norm": {"bias": 1.7081801891326904, "scale": 20.852447509765625}}, "20": {"attention": {"k_proj": {"bias": 0.4067543148994446, "kernel": 51.605438232421875}, "out_proj": {"bias": 1.359946370124817, "kernel": 49.45553207397461}, "q_proj": {"bias": 2.8498687744140625, "kernel": 52.224571228027344}, "v_proj": {"bias": 0.36227869987487793, "kernel": 48.43864822387695}}, "feed_forward": {"intermediate_dense": {"bias": 2.1725549697875977, "kernel": 109.17405700683594}, "output_dense": {"bias": 1.1388803720474243, "kernel": 106.40528106689453}}, "final_layer_norm": {"bias": 2.435314655303955, "scale": 23.4317626953125}, "layer_norm": {"bias": 2.231672525405884, "scale": 22.230525970458984}}, "21": {"attention": {"k_proj": {"bias": 0.4161534905433655, "kernel": 51.942527770996094}, "out_proj": {"bias": 1.403618335723877, "kernel": 49.51059341430664}, "q_proj": {"bias": 2.7690629959106445, "kernel": 52.67078399658203}, "v_proj": {"bias": 0.41060006618499756, "kernel": 48.64883041381836}}, "feed_forward": {"intermediate_dense": {"bias": 2.2174296379089355, "kernel": 109.5155029296875}, "output_dense": {"bias": 1.253208041191101, "kernel": 106.88243865966797}}, "final_layer_norm": {"bias": 2.4632763862609863, "scale": 23.175764083862305}, "layer_norm": {"bias": 2.2785892486572266, "scale": 22.234222412109375}}, "22": {"attention": {"k_proj": {"bias": 0.45357397198677063, "kernel": 52.54576110839844}, "out_proj": {"bias": 1.349219560623169, "kernel": 49.533172607421875}, "q_proj": {"bias": 2.8105549812316895, "kernel": 52.86981201171875}, "v_proj": {"bias": 0.3973655700683594, "kernel": 49.33363342285156}}, "feed_forward": {"intermediate_dense": {"bias": 2.1619315147399902, "kernel": 109.95498657226562}, "output_dense": {"bias": 1.3076066970825195, "kernel": 106.3852310180664}}, "final_layer_norm": {"bias": 2.3642821311950684, "scale": 22.684059143066406}, "layer_norm": {"bias": 2.3316237926483154, "scale": 21.545879364013672}}, "23": {"attention": {"k_proj": {"bias": 0.4928613007068634, "kernel": 53.47669219970703}, "out_proj": {"bias": 1.564335823059082, "kernel": 50.98707580566406}, "q_proj": {"bias": 2.7065773010253906, "kernel": 53.582611083984375}, "v_proj": {"bias": 0.5810648202896118, "kernel": 51.54853820800781}}, "feed_forward": {"intermediate_dense": {"bias": 2.131969690322876, "kernel": 109.86410522460938}, "output_dense": {"bias": 1.2769315242767334, "kernel": 107.37890625}}, "final_layer_norm": {"bias": 2.767916679382324, "scale": 22.887813568115234}, "layer_norm": {"bias": 2.824352264404297, "scale": 23.373172760009766}}, "24": {"attention": {"k_proj": {"bias": 0.46056002378463745, "kernel": 52.424072265625}, "out_proj": {"bias": 1.6070430278778076, "kernel": 52.50334167480469}, "q_proj": {"bias": 2.828113079071045, "kernel": 52.40515899658203}, "v_proj": {"bias": 0.5424190163612366, "kernel": 52.51116180419922}}, "feed_forward": {"intermediate_dense": {"bias": 2.2367913722991943, "kernel": 109.35035705566406}, "output_dense": {"bias": 1.3016372919082642, "kernel": 110.30095672607422}}, "final_layer_norm": {"bias": 2.83841872215271, "scale": 22.964658737182617}, "layer_norm": {"bias": 2.56215763092041, "scale": 22.983924865722656}}, "25": {"attention": {"k_proj": {"bias": 0.42509031295776367, "kernel": 52.730464935302734}, "out_proj": {"bias": 1.363797664642334, "kernel": 50.5806884765625}, "q_proj": {"bias": 2.9342763423919678, "kernel": 52.548744201660156}, "v_proj": {"bias": 0.6404213309288025, "kernel": 51.0885009765625}}, "feed_forward": {"intermediate_dense": {"bias": 2.1367578506469727, "kernel": 109.70021057128906}, "output_dense": {"bias": 1.1017413139343262, "kernel": 110.27072143554688}}, "final_layer_norm": {"bias": 2.5763301849365234, "scale": 23.494670867919922}, "layer_norm": {"bias": 2.683134078979492, "scale": 21.88357925415039}}, "26": {"attention": {"k_proj": {"bias": 0.4836847186088562, "kernel": 53.01764678955078}, "out_proj": {"bias": 1.2433912754058838, "kernel": 51.37077331542969}, "q_proj": {"bias": 2.943906784057617, "kernel": 52.80891036987305}, "v_proj": {"bias": 0.5064959526062012, "kernel": 52.004638671875}}, "feed_forward": {"intermediate_dense": {"bias": 2.2763516902923584, "kernel": 109.44652557373047}, "output_dense": {"bias": 1.0912110805511475, "kernel": 107.40899658203125}}, "final_layer_norm": {"bias": 2.1937994956970215, "scale": 22.433353424072266}, "layer_norm": {"bias": 2.497119903564453, "scale": 22.19057273864746}}, "27": {"attention": {"k_proj": {"bias": 0.5808594226837158, "kernel": 53.76898956298828}, "out_proj": {"bias": 1.5447406768798828, "kernel": 52.95805358886719}, "q_proj": {"bias": 2.703345775604248, "kernel": 53.69578552246094}, "v_proj": {"bias": 0.6748642325401306, "kernel": 53.388118743896484}}, "feed_forward": {"intermediate_dense": {"bias": 2.404933452606201, "kernel": 107.8713150024414}, "output_dense": {"bias": 0.9485896825790405, "kernel": 107.17198181152344}}, "final_layer_norm": {"bias": 2.5252954959869385, "scale": 21.88959503173828}, "layer_norm": {"bias": 2.6147172451019287, "scale": 23.32440948486328}}, "28": {"attention": {"k_proj": {"bias": 0.5901432037353516, "kernel": 54.482521057128906}, "out_proj": {"bias": 1.5367379188537598, "kernel": 53.31493377685547}, "q_proj": {"bias": 2.9472482204437256, "kernel": 54.1741943359375}, "v_proj": {"bias": 0.5131911039352417, "kernel": 53.759761810302734}}, "feed_forward": {"intermediate_dense": {"bias": 2.3475265502929688, "kernel": 107.87416076660156}, "output_dense": {"bias": 0.8224154710769653, "kernel": 109.1680908203125}}, "final_layer_norm": {"bias": 2.425306797027588, "scale": 22.337677001953125}, "layer_norm": {"bias": 2.0914058685302734, "scale": 23.993711471557617}}, "29": {"attention": {"k_proj": {"bias": 0.46781182289123535, "kernel": 51.12034606933594}, "out_proj": {"bias": 1.5021522045135498, "kernel": 55.685630798339844}, "q_proj": {"bias": 2.809702157974243, "kernel": 51.00274658203125}, "v_proj": {"bias": 0.4760415554046631, "kernel": 55.703304290771484}}, "feed_forward": {"intermediate_dense": {"bias": 2.297222137451172, "kernel": 108.01033020019531}, "output_dense": {"bias": 0.9597339630126953, "kernel": 113.12825012207031}}, "final_layer_norm": {"bias": 2.5980498790740967, "scale": 23.459980010986328}, "layer_norm": {"bias": 2.245180130004883, "scale": 25.39927864074707}}, "3": {"attention": {"k_proj": {"bias": 0.45006245374679565, "kernel": 52.03215789794922}, "out_proj": {"bias": 1.4254932403564453, "kernel": 48.60858917236328}, "q_proj": {"bias": 2.8560738563537598, "kernel": 52.312644958496094}, "v_proj": {"bias": 0.3246268630027771, "kernel": 48.768699645996094}}, "feed_forward": {"intermediate_dense": {"bias": 1.9663825035095215, "kernel": 104.83622741699219}, "output_dense": {"bias": 0.6984099745750427, "kernel": 94.07957458496094}}, "final_layer_norm": {"bias": 1.8095453977584839, "scale": 21.664737701416016}, "layer_norm": {"bias": 1.9017157554626465, "scale": 22.739452362060547}}, "30": {"attention": {"k_proj": {"bias": 0.5024805665016174, "kernel": 52.825706481933594}, "out_proj": {"bias": 1.3023658990859985, "kernel": 52.053871154785156}, "q_proj": {"bias": 2.907101631164551, "kernel": 52.91836166381836}, "v_proj": {"bias": 0.49308842420578003, "kernel": 52.49382019042969}}, "feed_forward": {"intermediate_dense": {"bias": 2.2399911880493164, "kernel": 108.17861938476562}, "output_dense": {"bias": 0.9140658378601074, "kernel": 112.09104919433594}}, "final_layer_norm": {"bias": 2.4926414489746094, "scale": 24.492368698120117}, "layer_norm": {"bias": 2.316732168197632, "scale": 24.931156158447266}}, "31": {"attention": {"k_proj": {"bias": 0.5412741899490356, "kernel": 51.240806579589844}, "out_proj": {"bias": 1.2333163022994995, "kernel": 52.19988250732422}, "q_proj": {"bias": 2.6581294536590576, "kernel": 51.346168518066406}, "v_proj": {"bias": 0.5469827651977539, "kernel": 52.432586669921875}}, "feed_forward": {"intermediate_dense": {"bias": 2.3097352981567383, "kernel": 106.72758483886719}, "output_dense": {"bias": 1.0891624689102173, "kernel": 109.24717712402344}}, "final_layer_norm": {"bias": 2.2962756156921387, "scale": 24.31252670288086}, "layer_norm": {"bias": 2.3430848121643066, "scale": 24.590187072753906}}, "32": {"attention": {"k_proj": {"bias": 0.4704548716545105, "kernel": 50.39933776855469}, "out_proj": {"bias": 1.2453913688659668, "kernel": 51.57465362548828}, "q_proj": {"bias": 2.8450098037719727, "kernel": 50.34847640991211}, "v_proj": {"bias": 0.419519305229187, "kernel": 51.972320556640625}}, "feed_forward": {"intermediate_dense": {"bias": 2.2569355964660645, "kernel": 105.33018493652344}, "output_dense": {"bias": 1.146787166595459, "kernel": 108.35574340820312}}, "final_layer_norm": {"bias": 2.31538724899292, "scale": 24.518985748291016}, "layer_norm": {"bias": 2.417579174041748, "scale": 24.991830825805664}}, "33": {"attention": {"k_proj": {"bias": 0.48390907049179077, "kernel": 50.28266143798828}, "out_proj": {"bias": 1.280959129333496, "kernel": 51.30557632446289}, "q_proj": {"bias": 2.998173713684082, "kernel": 50.253868103027344}, "v_proj": {"bias": 0.4416005611419678, "kernel": 51.71035385131836}}, "feed_forward": {"intermediate_dense": {"bias": 2.279946804046631, "kernel": 103.67684173583984}, "output_dense": {"bias": 1.1754591464996338, "kernel": 106.81501007080078}}, "final_layer_norm": {"bias": 2.257889747619629, "scale": 24.207740783691406}, "layer_norm": {"bias": 2.5865607261657715, "scale": 25.067882537841797}}, "34": {"attention": {"k_proj": {"bias": 0.45390456914901733, "kernel": 49.26523208618164}, "out_proj": {"bias": 1.5265402793884277, "kernel": 52.46200942993164}, "q_proj": {"bias": 2.913527011871338, "kernel": 49.268951416015625}, "v_proj": {"bias": 0.4023621678352356, "kernel": 52.534793853759766}}, "feed_forward": {"intermediate_dense": {"bias": 2.3745017051696777, "kernel": 102.21138000488281}, "output_dense": {"bias": 1.1227428913116455, "kernel": 105.72161102294922}}, "final_layer_norm": {"bias": 2.20273494720459, "scale": 23.640857696533203}, "layer_norm": {"bias": 2.615731716156006, "scale": 25.498104095458984}}, "35": {"attention": {"k_proj": {"bias": 0.5336894989013672, "kernel": 51.04902648925781}, "out_proj": {"bias": 1.4900906085968018, "kernel": 51.15006637573242}, "q_proj": {"bias": 2.573650360107422, "kernel": 51.323951721191406}, "v_proj": {"bias": 0.4889468550682068, "kernel": 51.176795959472656}}, "feed_forward": {"intermediate_dense": {"bias": 2.502547264099121, "kernel": 100.7535400390625}, "output_dense": {"bias": 1.0254027843475342, "kernel": 104.25007629394531}}, "final_layer_norm": {"bias": 2.295243263244629, "scale": 23.61981964111328}, "layer_norm": {"bias": 2.499863862991333, "scale": 26.093618392944336}}, "36": {"attention": {"k_proj": {"bias": 0.4473738670349121, "kernel": 48.319740295410156}, "out_proj": {"bias": 1.5177178382873535, "kernel": 52.26478958129883}, "q_proj": {"bias": 2.6124308109283447, "kernel": 48.236427307128906}, "v_proj": {"bias": 0.39316776394844055, "kernel": 52.66499328613281}}, "feed_forward": {"intermediate_dense": {"bias": 2.3669564723968506, "kernel": 99.60517120361328}, "output_dense": {"bias": 1.0260361433029175, "kernel": 103.68067932128906}}, "final_layer_norm": {"bias": 2.044203519821167, "scale": 24.13540267944336}, "layer_norm": {"bias": 2.2894434928894043, "scale": 25.63806915283203}}, "37": {"attention": {"k_proj": {"bias": 0.6240901350975037, "kernel": 47.300960540771484}, "out_proj": {"bias": 1.7604304552078247, "kernel": 52.189971923828125}, "q_proj": {"bias": 2.3819239139556885, "kernel": 47.316253662109375}, "v_proj": {"bias": 0.38518577814102173, "kernel": 52.32656478881836}}, "feed_forward": {"intermediate_dense": {"bias": 2.2694783210754395, "kernel": 98.56558227539062}, "output_dense": {"bias": 1.0087021589279175, "kernel": 103.13001251220703}}, "final_layer_norm": {"bias": 1.7903382778167725, "scale": 24.500085830688477}, "layer_norm": {"bias": 2.238807439804077, "scale": 25.563133239746094}}, "38": {"attention": {"k_proj": {"bias": 0.722891092300415, "kernel": 45.45281219482422}, "out_proj": {"bias": 1.4463956356048584, "kernel": 51.495086669921875}, "q_proj": {"bias": 2.2622861862182617, "kernel": 45.454437255859375}, "v_proj": {"bias": 0.42936205863952637, "kernel": 51.565425872802734}}, "feed_forward": {"intermediate_dense": {"bias": 2.2012171745300293, "kernel": 96.4342041015625}, "output_dense": {"bias": 0.9817801713943481, "kernel": 101.30387115478516}}, "final_layer_norm": {"bias": 1.7825044393539429, "scale": 25.214679718017578}, "layer_norm": {"bias": 2.4084482192993164, "scale": 26.425636291503906}}, "39": {"attention": {"k_proj": {"bias": 0.7209377884864807, "kernel": 45.24742889404297}, "out_proj": {"bias": 1.7097458839416504, "kernel": 51.329002380371094}, "q_proj": {"bias": 2.1152358055114746, "kernel": 45.53700637817383}, "v_proj": {"bias": 0.4246286153793335, "kernel": 51.27728271484375}}, "feed_forward": {"intermediate_dense": {"bias": 2.177624225616455, "kernel": 94.23902893066406}, "output_dense": {"bias": 1.0458879470825195, "kernel": 101.17718505859375}}, "final_layer_norm": {"bias": 1.753927230834961, "scale": 25.785865783691406}, "layer_norm": {"bias": 2.33198881149292, "scale": 26.942312240600586}}, "4": {"attention": {"k_proj": {"bias": 0.44548165798187256, "kernel": 54.6234130859375}, "out_proj": {"bias": 1.652343988418579, "kernel": 50.16497039794922}, "q_proj": {"bias": 2.615248918533325, "kernel": 54.93098068237305}, "v_proj": {"bias": 0.34892427921295166, "kernel": 50.320098876953125}}, "feed_forward": {"intermediate_dense": {"bias": 1.9531786441802979, "kernel": 104.50968933105469}, "output_dense": {"bias": 0.8575068712234497, "kernel": 95.54541015625}}, "final_layer_norm": {"bias": 1.9920002222061157, "scale": 21.200180053710938}, "layer_norm": {"bias": 2.054612159729004, "scale": 23.620338439941406}}, "40": {"attention": {"k_proj": {"bias": 0.6590453386306763, "kernel": 44.198089599609375}, "out_proj": {"bias": 1.6252505779266357, "kernel": 49.55699920654297}, "q_proj": {"bias": 1.9674756526947021, "kernel": 44.89208221435547}, "v_proj": {"bias": 0.4587768614292145, "kernel": 49.23614501953125}}, "feed_forward": {"intermediate_dense": {"bias": 2.0333969593048096, "kernel": 92.16896057128906}, "output_dense": {"bias": 1.087776780128479, "kernel": 98.40738677978516}}, "final_layer_norm": {"bias": 1.7852704524993896, "scale": 25.04292106628418}, "layer_norm": {"bias": 2.2756104469299316, "scale": 26.40799903869629}}, "41": {"attention": {"k_proj": {"bias": 1.7133712768554688, "kernel": 41.96858596801758}, "out_proj": {"bias": 1.3790823221206665, "kernel": 51.25593566894531}, "q_proj": {"bias": 1.71382737159729, "kernel": 42.56317901611328}, "v_proj": {"bias": 0.4695759415626526, "kernel": 50.369529724121094}}, "feed_forward": {"intermediate_dense": {"bias": 2.110393524169922, "kernel": 88.92567443847656}, "output_dense": {"bias": 1.1446669101715088, "kernel": 97.37409973144531}}, "final_layer_norm": {"bias": 2.23917293548584, "scale": 28.507984161376953}, "layer_norm": {"bias": 2.22525691986084, "scale": 28.246891021728516}}, "42": {"attention": {"k_proj": {"bias": 0.8601109981536865, "kernel": 38.31235885620117}, "out_proj": {"bias": 1.4427157640457153, "kernel": 45.07648849487305}, "q_proj": {"bias": 1.549715280532837, "kernel": 39.524009704589844}, "v_proj": {"bias": 0.6933339834213257, "kernel": 43.49076461791992}}, "feed_forward": {"intermediate_dense": {"bias": 1.9107489585876465, "kernel": 88.009765625}, "output_dense": {"bias": 1.1978566646575928, "kernel": 95.7593994140625}}, "final_layer_norm": {"bias": 1.9227323532104492, "scale": 29.817535400390625}, "layer_norm": {"bias": 1.6761282682418823, "scale": 26.810440063476562}}, "43": {"attention": {"k_proj": {"bias": 1.247081995010376, "kernel": 34.694725036621094}, "out_proj": {"bias": 1.4174811840057373, "kernel": 41.36320495605469}, "q_proj": {"bias": 1.3773530721664429, "kernel": 35.38981628417969}, "v_proj": {"bias": 0.5787136554718018, "kernel": 39.29212951660156}}, "feed_forward": {"intermediate_dense": {"bias": 1.8936835527420044, "kernel": 87.073974609375}, "output_dense": {"bias": 0.9419379234313965, "kernel": 93.74283599853516}}, "final_layer_norm": {"bias": 1.99924898147583, "scale": 32.0491943359375}, "layer_norm": {"bias": 1.7940990924835205, "scale": 25.131242752075195}}, "44": {"attention": {"k_proj": {"bias": 2.5188145637512207, "kernel": 35.16314697265625}, "out_proj": {"bias": 1.1667875051498413, "kernel": 45.019126892089844}, "q_proj": {"bias": 1.317202091217041, "kernel": 35.58300018310547}, "v_proj": {"bias": 0.38874924182891846, "kernel": 44.14718246459961}}, "feed_forward": {"intermediate_dense": {"bias": 1.9462898969650269, "kernel": 86.07953643798828}, "output_dense": {"bias": 0.859266996383667, "kernel": 91.58765411376953}}, "final_layer_norm": {"bias": 2.0454273223876953, "scale": 34.2881965637207}, "layer_norm": {"bias": 1.6815991401672363, "scale": 25.14142608642578}}, "45": {"attention": {"k_proj": {"bias": 2.081407308578491, "kernel": 34.86139678955078}, "out_proj": {"bias": 1.0356104373931885, "kernel": 48.59937286376953}, "q_proj": {"bias": 1.402512788772583, "kernel": 35.03264617919922}, "v_proj": {"bias": 0.4231463074684143, "kernel": 48.76853942871094}}, "feed_forward": {"intermediate_dense": {"bias": 2.016927719116211, "kernel": 82.93773651123047}, "output_dense": {"bias": 0.9764893054962158, "kernel": 87.24796295166016}}, "final_layer_norm": {"bias": 1.9180456399917603, "scale": 33.143672943115234}, "layer_norm": {"bias": 1.5726068019866943, "scale": 23.782546997070312}}, "46": {"attention": {"k_proj": {"bias": 1.5659263134002686, "kernel": 35.878021240234375}, "out_proj": {"bias": 0.8182340264320374, "kernel": 51.16078186035156}, "q_proj": {"bias": 1.5642974376678467, "kernel": 36.18907165527344}, "v_proj": {"bias": 0.4092414081096649, "kernel": 51.89159393310547}}, "feed_forward": {"intermediate_dense": {"bias": 2.0093321800231934, "kernel": 77.47581481933594}, "output_dense": {"bias": 1.1406863927841187, "kernel": 77.73695373535156}}, "final_layer_norm": {"bias": 1.8108854293823242, "scale": 28.70657730102539}, "layer_norm": {"bias": 1.3991491794586182, "scale": 22.808137893676758}}, "47": {"attention": {"k_proj": {"bias": 0.6173280477523804, "kernel": 38.678985595703125}, "out_proj": {"bias": 0.6758822202682495, "kernel": 46.45281219482422}, "q_proj": {"bias": 1.7084776163101196, "kernel": 39.426841735839844}, "v_proj": {"bias": 0.4932914674282074, "kernel": 47.617279052734375}}, "feed_forward": {"intermediate_dense": {"bias": 1.986911654472351, "kernel": 75.47482299804688}, "output_dense": {"bias": 0.6346586346626282, "kernel": 72.82707214355469}}, "final_layer_norm": {"bias": 1.1888140439987183, "scale": 23.650447845458984}, "layer_norm": {"bias": 1.2521969079971313, "scale": 20.66573715209961}}, "5": {"attention": {"k_proj": {"bias": 0.42588678002357483, "kernel": 50.1945686340332}, "out_proj": {"bias": 1.6038882732391357, "kernel": 51.2144889831543}, "q_proj": {"bias": 2.7522244453430176, "kernel": 50.37500762939453}, "v_proj": {"bias": 0.3343381881713867, "kernel": 51.71652603149414}}, "feed_forward": {"intermediate_dense": {"bias": 1.8887722492218018, "kernel": 104.60663604736328}, "output_dense": {"bias": 0.8976269960403442, "kernel": 94.77360534667969}}, "final_layer_norm": {"bias": 2.1965675354003906, "scale": 21.37998390197754}, "layer_norm": {"bias": 2.0435237884521484, "scale": 22.437192916870117}}, "6": {"attention": {"k_proj": {"bias": 0.4843112528324127, "kernel": 51.87700653076172}, "out_proj": {"bias": 1.5925445556640625, "kernel": 50.83113479614258}, "q_proj": {"bias": 2.7889723777770996, "kernel": 52.3514404296875}, "v_proj": {"bias": 0.3247200846672058, "kernel": 51.107398986816406}}, "feed_forward": {"intermediate_dense": {"bias": 1.8638136386871338, "kernel": 103.7142333984375}, "output_dense": {"bias": 0.752193808555603, "kernel": 94.57742309570312}}, "final_layer_norm": {"bias": 2.5145251750946045, "scale": 20.836563110351562}, "layer_norm": {"bias": 2.0285890102386475, "scale": 23.156789779663086}}, "7": {"attention": {"k_proj": {"bias": 0.5048109889030457, "kernel": 51.46453094482422}, "out_proj": {"bias": 1.4398455619812012, "kernel": 51.139068603515625}, "q_proj": {"bias": 2.550907611846924, "kernel": 51.92047119140625}, "v_proj": {"bias": 0.42719271779060364, "kernel": 50.953285217285156}}, "feed_forward": {"intermediate_dense": {"bias": 1.8739449977874756, "kernel": 103.49991607666016}, "output_dense": {"bias": 0.5876641273498535, "kernel": 94.39387512207031}}, "final_layer_norm": {"bias": 2.416801929473877, "scale": 21.010677337646484}, "layer_norm": {"bias": 1.9788501262664795, "scale": 22.19708824157715}}, "8": {"attention": {"k_proj": {"bias": 0.49711495637893677, "kernel": 51.12122344970703}, "out_proj": {"bias": 1.2548246383666992, "kernel": 51.65118408203125}, "q_proj": {"bias": 2.541980504989624, "kernel": 51.03407287597656}, "v_proj": {"bias": 0.35420340299606323, "kernel": 51.662872314453125}}, "feed_forward": {"intermediate_dense": {"bias": 1.9298467636108398, "kernel": 103.21916198730469}, "output_dense": {"bias": 0.5482766628265381, "kernel": 93.97574615478516}}, "final_layer_norm": {"bias": 2.3526053428649902, "scale": 20.7393856048584}, "layer_norm": {"bias": 1.9221248626708984, "scale": 22.40435028076172}}, "9": {"attention": {"k_proj": {"bias": 0.5231171250343323, "kernel": 52.01068878173828}, "out_proj": {"bias": 1.4968843460083008, "kernel": 52.671897888183594}, "q_proj": {"bias": 2.4629459381103516, "kernel": 52.26807403564453}, "v_proj": {"bias": 0.38445231318473816, "kernel": 52.86597442626953}}, "feed_forward": {"intermediate_dense": {"bias": 2.026733875274658, "kernel": 101.98421478271484}, "output_dense": {"bias": 0.6828575134277344, "kernel": 94.36962890625}}, "final_layer_norm": {"bias": 2.325080156326294, "scale": 20.160720825195312}, "layer_norm": {"bias": 2.0236480236053467, "scale": 24.083864212036133}}}, "pos_conv_embed": {"conv": {"bias": 5.847014427185059, "weight_g": 9.12463665008545, "weight_v": 93.52015686035156}}}, "feature_extractor": {"conv_layers": {"0": {"conv": {"bias": 2.0290679931640625, "kernel": 20.55536460876465}, "layer_norm": {"bias": 4.550922393798828, "scale": 16.167570114135742}}, "1": {"conv": {"bias": 1.7790228128433228, "kernel": 51.24136734008789}, "layer_norm": {"bias": 5.962646961212158, "scale": 23.268157958984375}}, "2": {"conv": {"bias": 1.140576720237732, "kernel": 46.50312042236328}, "layer_norm": {"bias": 4.176670551300049, "scale": 20.370853424072266}}, "3": {"conv": {"bias": 0.6725863218307495, "kernel": 44.397525787353516}, "layer_norm": {"bias": 3.888174533843994, "scale": 17.53795051574707}}, "4": {"conv": {"bias": 0.6373162269592285, "kernel": 41.314056396484375}, "layer_norm": {"bias": 2.385471820831299, "scale": 16.34571647644043}}, "5": {"conv": {"bias": 0.5147221684455872, "kernel": 37.479759216308594}, "layer_norm": {"bias": 2.020900011062622, "scale": 17.064470291137695}}, "6": {"conv": {"bias": 0.4947893023490906, "kernel": 40.64780044555664}, "layer_norm": {"bias": 0.5876954793930054, "scale": 19.058603286743164}}}}, "feature_projection": {"layer_norm": {"bias": 6.376383304595947, "scale": 16.443069458007812}, "projection": {"bias": 1.8670344352722168, "kernel": 37.218414306640625}}, "masked_spec_embed": 11.914372444152832}}, "train/learning_rate": 1.9151924789184704e-05, "train/loss": 0.204779714345932, "train/param_norm": 1241.662353515625, "_runtime": 5032, "_timestamp": 1660135191, "_step": 275600, "_wandb": {"runtime": 5033}}

wandb/run-20220810_111559-290849gb/logs/debug-internal.log

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

wandb/run-20220810_111559-290849gb/logs/debug.log

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

wandb/run-20220810_111559-290849gb/run-290849gb.wandb

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

wandb/run-20220810_145446-1k92sv35/files/code/run_flax_speech_recognition_ctc.py

@@ -0,0 +1,1632 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace 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 connectionist temporal classification (CTC) 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 re
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import datasets
+import numpy as np
+from datasets import DatasetDict, load_dataset, load_metric
+from tqdm import tqdm
+
+import flax
+import jax
+import jax.numpy as jnp
+import optax
+import transformers
+import wandb as wandb
+from flax import core, jax_utils, struct, traverse_util
+from flax.jax_utils import unreplicate, pad_shard_unpad
+from flax.training.common_utils import get_metrics, shard, shard_prng_key
+from huggingface_hub import Repository
+from models import Wav2Vec2Config, FlaxWav2Vec2ForCTC
+from optax._src import linear_algebra
+from transformers import (
+    AutoFeatureExtractor,
+    AutoProcessor,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainingArguments,
+    is_tensorboard_available,
+    set_seed,
+)
+from transformers.file_utils import get_full_repo_name
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.17.0.dev0")
+
+require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/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)."
+        },
+    )
+    freeze_feature_encoder: bool = field(
+        default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}
+    )
+    attention_dropout: float = field(
+        default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}
+    )
+    activation_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The hidden activation dropout probability in the embeddings, encoder, and pooler."
+        },
+    )
+    hidden_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The dropout probability for all fully connected layers in the embeddings, encoder, and pooler."
+        },
+    )
+    feat_proj_dropout: float = field(
+        default=0.0,
+        metadata={
+            "help": "The feat proj dropout probability for feature encoder representations."
+        },
+    )
+    final_dropout: float = field(
+        default=0.0,
+        metadata={"help": "The dropout probability for the final projection layer."},
+    )
+    mask_time_prob: float = field(
+        default=0.1,
+        metadata={
+            "help": "The spec aug dropout probability for feature encoder representations."
+        },
+    )
+    mask_time_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the time axis."},
+    )
+    mask_feature_prob: float = field(
+        default=0.0,
+        metadata={
+            "help": "Probability of each feature vector along the feature axis to be chosen as the start of the vector"
+                    "span to be masked. Approximately ``mask_feature_prob * sequence_length // mask_feature_length`` feature bins will be masked along the time axis."
+        },
+    )
+    mask_feature_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the feature axis."},
+    )
+    layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})
+    ctc_loss_reduction: Optional[str] = field(
+        default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}
+    )
+    ctc_zero_infinity: Optional[bool] = field(
+        default=False, metadata={"help": "If True, will try yo aboud the CTC loss goinf to infinity."}
+    )
+
+
+@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."
+        },
+    )
+    max_test_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of test 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": "Truncate audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_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: Optional[int] = field(
+        default=512,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    min_label_length: Optional[int] = field(
+        default=2,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    pad_input_to_multiple_of: Optional[int] = field(
+        default=32000,
+        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."
+        },
+    )
+    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."
+        },
+    )
+    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 training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    do_lower_case: bool = field(
+        default=True,
+        metadata={"help": "Whether the target text should be lower cased."},
+    )
+    wandb_project: str = field(
+        default="flax-speech-recognition-ctc",
+        metadata={"help": "The name of the wandb project."},
+    )
+    wandb_name: str = field(
+        default=None,
+        metadata={"help": "The name of the wandb run."},
+    )
+    wandb_job_type: str = field(
+        default="CTC",
+        metadata={"help": "The name of the wandb job type."},
+    )
+    test_split_name: str = field(
+        default="test",
+        metadata={"help": "The name of the test data set split to use (via the datasets library). Defaults to 'test'"},
+    )
+    remove_punctuation: bool = field(
+        default=False, metadata={"help": "Whether or not to remove punctuation during training."}
+    )
+    skip_steps: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": "Skip this number of steps. Useful to continue training"
+        },
+    )
+
+
+# @flax.struct.dataclass
+@dataclass
+class FlaxTrainingArguments(TrainingArguments):
+    precision: str = field(
+        default="full",
+        metadata={
+            "help": "Whether to enable mixed-precision training. If true, the optimizer is stored in half-precision (bfloat16) and computations are executed in half-precision"
+            "**Note that this only specifies the dtype of the computation and optimizer state. It does not influence the dtype of model parameters.**"
+        },
+    )
+    matmul_precision: str = field(
+        default="default",
+        metadata={
+            "help": "Default floating-point precision of internal computations used in TPU matrix multiplications and convolutions. "
+            "This configuration option controls the default precision for JAX operations that take an optional precision argument (e.g. `lax.conv_general_dilated` and `lax.dot`). "
+            "This configuration option does not change the behaviours of such calls with explicit precision arguments; "
+            "it only changes the behaviors of calls with no such argument provided. "
+            "One of `['highest', 'float32', 'high', 'bfloat16_3x', 'default', 'bfloat16', 'fastest', None]`."
+        },
+    )
+    multisteps: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to use Optax MultiSteps for gradient accumulation. If `False` (default) and `gradient_accumulation_steps > 1`, "
+            "a custom gradient accumulation implementation will be employed."
+        },
+    )
+
+
+def to_fp32(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+
+
+def to_bf16(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t)
+
+
+class MixedPrecisionTrainState(struct.PyTreeNode):
+    """Train state for use with a single Optax optimizer.
+    Adapted from flax train_state https://github.com/google/flax/blob/main/flax/training/train_state.py
+
+    Synopsis::
+
+        state = TrainState.create(
+            apply_fn=model.apply,
+            params=variables['params'],
+            tx=tx)
+        grad_fn = jax.grad(make_loss_fn(state.apply_fn))
+        for batch in data:
+          grads = grad_fn(state.params, batch)
+          state = state.apply_gradients(grads=grads)
+
+    Args:
+      step: Counter starts at 0 and is incremented by every call to
+        `.apply_gradients()`.
+      apply_fn: Usually set to `model.apply()`. Kept in this dataclass for
+        convenience to have a shorter params list for the `train_step()` function
+        in your training loop.
+      params: The parameters to be updated by `tx` and used by `apply_fn`.
+      tx: An Optax gradient transformation.
+      opt_state: The state for `tx`.
+      dropout_rng: PRNG key for stochastic operations.
+      bf16: Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training.
+    """
+
+    step: int
+    apply_fn: Callable = struct.field(pytree_node=False)
+    get_attention_mask_fn: Callable = struct.field(pytree_node=False)
+    params: core.FrozenDict[str, Any]
+    tx: optax.GradientTransformation = struct.field(pytree_node=False)
+    opt_state: optax.OptState
+    dropout_rng: jnp.ndarray
+    max_grad_norm: Optional[float] = 1.0
+
+    def apply_gradients(self, *, grads, to_dtype, **kwargs):
+        """Updates `step`, `params`, `opt_state` and `**kwargs` in return value.
+
+        Note that internally this function calls `.tx.update()` followed by a call
+        to `optax.apply_updates()` to update `params` and `opt_state`.
+
+        Args:
+          grads: Gradients that have the same pytree structure as `.params`.
+          **kwargs: Additional dataclass attributes that should be `.replace()`-ed.
+
+        Returns:
+          An updated instance of `self` with `step` incremented by one, `params`
+          and `opt_state` updated by applying `grads`, and additional attributes
+          replaced as specified by `kwargs`.
+        """
+
+        # clip gradients by global l2 norm
+        casted_max_grad_norm = to_dtype(self.max_grad_norm)
+        g_norm = linear_algebra.global_norm(grads)
+        g_norm = jnp.maximum(casted_max_grad_norm, g_norm)
+        grads = jax.tree_util.tree_map(lambda t: (t / g_norm) * casted_max_grad_norm, grads)
+
+        # perform update step in fp32 and subsequently downcast optimizer states if mixed precision training
+        # grads and opt_state in bf16 (need to upcast), params in fp32 (leave as is)
+        updates, new_opt_state = self.tx.update(to_fp32(grads), to_fp32(self.opt_state), self.params)
+
+        new_params = optax.apply_updates(self.params, updates)
+        return self.replace(
+            step=self.step + 1,
+            params=new_params,
+            opt_state=to_dtype(new_opt_state),
+            **kwargs,
+        )
+
+    @classmethod
+    def create(cls, *, apply_fn, params, tx, to_dtype, step=0, **kwargs):
+        """Creates a new instance with `step=0` and initialized `opt_state`."""
+        # downcast optimizer state to bf16 if mixed-precision training
+        opt_state = tx.init(to_dtype(params)) if tx is not None else None
+        return cls(
+            step=step,
+            apply_fn=apply_fn,
+            params=params,
+            tx=tx,
+            opt_state=opt_state,
+            **kwargs,
+        )
+
+    def replicate(self):
+        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
+
+
+@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).
+        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
+    input_padding: Union[bool, str] = "longest"
+    label_padding: Union[bool, str] = "max_length"
+    pad_input_to_multiple_of: Optional[int] = None
+    pad_to_multiple_of_label: Optional[int] = None
+    max_input_length: Optional[float] = None
+    max_label_length: Optional[float] = 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
+        input_features = [{"input_values": feature["input_values"]} for feature in features]
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+
+        # reformat list to dict and set to pytorch format
+        batch = self.processor.feature_extractor.pad(
+            input_features,
+            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_label_length,
+            padding=self.label_padding,
+            pad_to_multiple_of=self.pad_to_multiple_of_label,
+            return_tensors="np",
+        )
+
+        labels = labels_batch["input_ids"]
+        labels = np.ma.array(labels, mask=np.not_equal(labels_batch.attention_mask, 1))
+        labels = labels.filled(fill_value=-100)
+
+        batch["labels"] = labels
+
+        return batch
+
+
+def get_grouped_indices(
+    dataset, batch_size: int, rng: Optional[List[int]] = None, mega_batch_mult: Optional[int] = None
+) -> np.array:
+    """
+    Adapted from the `get_length_grouped_indices` function in the PyTorch Trainer utils file (https://github.com/huggingface/transformers/blob/main/src/transformers/trainer_pt_utils.py#L486)
+    Function that returns a list of indices in which each slice of `batch_size` consecutive indices correspond to elements of similar
+    lengths. To do this, the indices are:
+
+    - randomly permuted (if a JAX rng is specified)
+    - grouped in mega-batches of size `mega_batch_mult * batch_size`
+    - sorted by length in each mega-batch
+
+    The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of
+    maximum length placed first, so that an OOM happens sooner rather than later.
+    """
+    lengths = dataset["input_length"]
+
+    # Default for mega_batch_mult: 50 or the number to get 4 megabatches, whichever is smaller.
+    if mega_batch_mult is None:
+        mega_batch_mult = min(len(lengths) // (batch_size * 4), 50)
+        # Just in case, for tiny datasets
+        if mega_batch_mult == 0:
+            mega_batch_mult = 1
+
+    # We need to use JAX for the random permutation as the PRNG key will be set based on the seed outside of the sampler.
+    num_samples = len(lengths)
+    indices = jax.random.permutation(rng, np.arange(num_samples)) if rng is not None else np.arange(num_samples)
+
+    megabatch_size = mega_batch_mult * batch_size
+    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
+    megabatches = [list(sorted(megabatch, key=lambda i: lengths[i], reverse=True)) for megabatch in megabatches]
+
+    # The rest is to get the biggest batch first.
+    # Since each megabatch is sorted by descending length, the longest element is the first
+    megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches]
+    max_idx = np.argmax(megabatch_maximums).item()
+    # Switch to put the longest batch in first position
+    # (note that this is different to the PT grouped sampler in which we only put the longest element in the first position, and not its batch)
+    megabatches[0], megabatches[max_idx] = megabatches[max_idx], megabatches[0]
+
+    megabatches = np.array([i for megabatch in megabatches for i in megabatch])
+
+    return megabatches
+
+
+def generate_batch_splits(samples_idx: np.ndarray, batch_size: int, drop_last=True) -> np.ndarray:
+    """Generate batches of data for a specified batch size from sample indices. If the dataset size is not divisible by
+    the batch size and `drop_last` is `True`, the last incomplete batch is dropped. Else, it is returned."""
+    num_samples = len(samples_idx)
+    if drop_last:
+        samples_to_remove = num_samples % batch_size
+        if samples_to_remove != 0:
+            samples_idx = samples_idx[:-samples_to_remove]
+        sections_split = num_samples // batch_size
+        samples_idx = samples_idx.reshape((sections_split, batch_size))
+    else:
+        sections_split = math.ceil(num_samples / batch_size)
+        samples_idx = np.array_split(samples_idx, sections_split)
+    return samples_idx
+
+
+def write_train_metric(summary_writer, train_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)
+
+
+def write_eval_metric(summary_writer, eval_metrics, step, pred_str=None):
+    for metric_name, value in eval_metrics.items():
+        summary_writer.scalar(f"eval_{metric_name}", value, step)
+
+    if pred_str is not None:
+        # write output actual predictions for debugging
+        summary_writer.text("eval_predictions", "\n".join(pred_str), step)
+
+
+def write_wandb_log(metrics, step, prefix=None):
+    if jax.process_index() == 0:
+        log_metrics = {}
+        for k, v in metrics.items():
+            if "layer" in k:
+                log_metrics[f"{k}/"] = v
+            elif prefix is not None:
+                log_metrics[f"{prefix}/{k}"] = v
+            else:
+                log_metrics[k] = v
+        wandb.log(log_metrics, step)
+
+
+def write_wandb_pred(pred_str, label_str, step, num_log=50, prefix="eval"):
+    if jax.process_index() == 0:
+        # convert str data to a wandb compatible format
+        str_data = [[label_str[i], pred_str[i]] for i in range(len(pred_str))]
+        # we'll log the first 50 predictions for each epoch
+        wandb.log(
+            {
+                f"{prefix}/step_{int(step / 1000)}k": wandb.Table(
+                    columns=["label_str", "pred_str"], data=str_data[:num_log]
+                )
+            },
+            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 ctc_loss(
+    logits,
+    logits_attention_mask,
+    labels,
+    blank_id,
+    loss_reduction="mean",
+    output_emission_dict=False,
+    log_epsilon=-100000.0,
+):
+    """Computes CTC loss.
+    This function performs forward computation over an FSA with `N * 2` states
+    where `N` is the max number of labels. The states are split into two groups:
+    Phi states and emission states. a phi-state accepts repetition of
+    phi (blank)-symbols and transits to emission state when the correct label is
+    observed. An emission state accepts repetition of the label and transits to
+    the next phi states at any time (so called epsilon-transition).
+    Below, `B` denotes the batch size, `T` denotes the time steps in `logits`,
+    and `N` denotes the time steps in `labels`.
+    Args:
+      logits: (B, T, K)-array containing log-probabilities of each class.
+      logitpaddings: (B, T)-array. Padding indicators for `logits`.
+      labels: (B, N)-array containing reference integer labels.
+      labelpaddings: (B, N)-array. Padding indicators for `labels`. Currently,
+        `labels` must be right-padded, i.e. each row of `labelpaddings` must be
+        repetition of zeroes, followed by repetition of ones.
+      blank_id: Id for blank token.
+      loss_reduction: one of "mean", "sum", "default"
+        - "none": no reduction is applied.
+        - "mean": output loss will be divided by target lengths and then the
+          mean over the batch is taken.
+        - "sum": output loss are summed over batch
+      output_emission_dict: whether to output additional information about the emission probs
+    Returns:
+      A pair of `(per_seq_loss, aux)`.
+      per_seq_loss:
+        (B,)-array containing loss values for each sequence in the batch.
+      aux: Dictionary containing interim variables used for computing losses.
+        aux['logalpha_phi']: (T, B, N+1)-array. Log-forward-probabilities of each
+          phi-state corresponding to the n-th label.
+        aux['logalpha_emit']: (T, B, N)-array. Log-forward-probabilities of each
+          emission-state corresponding to the n-th label.
+        aux['logprobs_phi']: (T, B, 1)-array. Probability of the phi-symbol
+          corresponding to each time frame.
+        aux['logprobs_emit']: (T, B, N)-array. Probability of the n-th label
+          corresponding to each time frame.
+    """
+    # label paddings are indicated by -100
+    labelpaddings = labels < 0
+    # logit paddings are the inverse of attention_mask
+    logitpaddings = ~logits_attention_mask
+
+    # Copied from https://github.com/tensorflow/lingvo/blob/master/lingvo/jax/layers/ctc_objectives.py
+    batchsize, unused_maxinputlen, num_classes = logits.shape
+    batchsize_, maxlabellen = labels.shape
+
+    logprobs = jax.nn.log_softmax(logits)
+    labellens = maxlabellen - jnp.sum(labelpaddings, axis=1).astype(jnp.int32)
+
+    # repeat[b, n] == 1.0 when label[b, n] == label[b, n+1].
+    repeat = (labels[:, :-1] == labels[:, 1:]).astype(jnp.float32)
+    repeat = jnp.pad(repeat, ((0, 0), (0, 1)))
+
+    logprobs_phi = logprobs[:, :, blank_id : blank_id + 1]  # [B, T, 1]
+    logprobs_phi = jnp.transpose(logprobs_phi, (1, 0, 2))  # [T, B, 1]
+
+    one_hot = jax.nn.one_hot(labels, num_classes=num_classes)  # [B, N, K]
+    logprobs_emit = jnp.einsum("btk,bnk->btn", logprobs, one_hot)
+    logprobs_emit = jnp.transpose(logprobs_emit, (1, 0, 2))  # [T, B, N]
+
+    logalpha_phi_init = jnp.ones((batchsize, maxlabellen + 1)) * log_epsilon  # [B, N]
+    logalpha_phi_init = logalpha_phi_init.at[:, 0].set(0.0)
+    logalpha_emit_init = jnp.ones((batchsize, maxlabellen)) * log_epsilon  # [B, N]
+
+    def loop_body(prev, x):
+        prev_phi, prev_emit = prev
+        # emit-to-phi epsilon transition, except if the next label is repetition
+        prev_phi_orig = prev_phi
+        prev_phi = prev_phi.at[:, 1:].set(jnp.logaddexp(prev_phi[:, 1:], prev_emit + log_epsilon * repeat))
+
+        logprob_emit, logprob_phi, pad = x
+
+        # phi-to-emit transition
+        next_emit = jnp.logaddexp(prev_phi[:, :-1] + logprob_emit, prev_emit + logprob_emit)
+        # self-loop transition
+        next_phi = prev_phi + logprob_phi
+        # emit-to-phi blank transition only when the next label is repetition
+        next_phi = next_phi.at[:, 1:].set(
+            jnp.logaddexp(next_phi[:, 1:], prev_emit + logprob_phi + log_epsilon * (1.0 - repeat))
+        )
+
+        pad = pad.reshape((batchsize, 1))
+        next_emit = pad * prev_emit + (1.0 - pad) * next_emit
+        next_phi = pad * prev_phi_orig + (1.0 - pad) * next_phi
+
+        return (next_phi, next_emit), (next_phi, next_emit)
+
+    xs = (logprobs_emit, logprobs_phi, logitpaddings.transpose((1, 0)))
+    _, (logalpha_phi, logalpha_emit) = jax.lax.scan(loop_body, (logalpha_phi_init, logalpha_emit_init), xs)
+
+    # last row needs to be updated with the last epsilon transition
+    logalpha_phi_last = logalpha_phi[-1].at[:, 1:].set(jnp.logaddexp(logalpha_phi[-1, :, 1:], logalpha_emit[-1]))
+    logalpha_phi = logalpha_phi.at[-1].set(logalpha_phi_last)
+
+    # extract per_seq_loss
+    one_hot = jax.nn.one_hot(labellens, num_classes=maxlabellen + 1)  # [B, N+1]
+    per_seq_loss = -jnp.einsum("bn,bn->b", logalpha_phi_last, one_hot)
+
+    if loss_reduction == "mean":
+        target_lengths = labelpaddings.shape[-1] - labelpaddings.sum(axis=-1)
+        loss = (per_seq_loss / target_lengths).mean()
+    elif loss_reduction == "sum":
+        loss = per_seq_loss.sum()
+    else:
+        loss = per_seq_loss
+
+    if not output_emission_dict:
+        return loss
+
+    return loss, {
+        "logalpha_phi": logalpha_phi,
+        "logalpha_emit": logalpha_emit,
+        "logprobs_phi": logprobs_phi,
+        "logprobs_emit": logprobs_emit,
+    }
+
+
+def make_dataset(data_args, seed=42):
+    # Pre-processing dataset
+    import re
+
+    def map_nst(entry):
+        text = entry["text"].lower()
+        text = text.replace("(...vær stille under dette opptaket...)", "")
+        text = re.sub('[áàâ]', 'a', text)
+        text = re.sub('[ä]', 'æ', text)
+        text = re.sub('[éèëê]', 'e', text)
+        text = re.sub('[íìïî]', 'i', text)
+        text = re.sub('[óòöô]', 'o', text)
+        text = re.sub('[ö]', 'ø', text)
+        text = re.sub('[ç]', 'c', text)
+        text = re.sub('[úùüû]', 'u', text)
+        # text = re.sub('\\(?=(Punktum|Komma|Utropstegn|Spørsmålstegn))', ' ', text)
+        text = re.sub('\s+', ' ', text)
+        return {"text": text}
+
+    def filter_nst(entry):
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.match(entry["type"], "pIW|CA"):
+            return False  # Spelling out words
+        return True
+
+    def filter_npsc(entry):
+        # False if there are digits in the text
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.search("\d", entry["text"]):
+            return False
+        return True
+
+    def map_npsc(entry):
+        batch = {"text": entry["text"].lower()}
+        batch["text"] = re.sub('[áàâ]', 'a', batch["text"])
+        batch["text"] = re.sub('[ä]', 'æ', batch["text"])
+        batch["text"] = re.sub('[éèëê]', 'e', batch["text"])
+        batch["text"] = re.sub('[íìïî]', 'i', batch["text"])
+        batch["text"] = re.sub('[óòöô]', 'o', batch["text"])
+        batch["text"] = re.sub('[ö]', 'ø', batch["text"])
+        batch["text"] = re.sub('[ç]', 'c', batch["text"])
+        batch["text"] = re.sub('[úùüû]', 'u', batch["text"])
+        batch["text"] = re.sub('\s', ' ', batch["text"])
+        batch["text"] = re.sub('<ee>', 'eee', batch["text"])
+        batch["text"] = re.sub('<qq>', 'qqq', batch["text"])
+        batch["text"] = re.sub('<mm>', 'mmm', batch["text"])
+        batch["text"] = re.sub('<inaudible>', 'xxx', batch["text"])
+        # batch["text"] = re.sub('<inaudible>', '?', batch["text"])
+        if "<" in batch["text"]:
+            raise ValueError(batch["text"])
+        return batch
+
+    nst = datasets.load_dataset("NbAiLab/NST", "no-close")
+    npsc = datasets.load_dataset("NbAiLab/NPSC", "16K_mp3")
+    # TODO NST_hesitate
+
+    split = len(npsc[data_args.train_split_name]) / (len(npsc[data_args.train_split_name]) + len(npsc[data_args.eval_split_name]))  # Use same train/val ratio as NPSC
+    nst_train = nst[data_args.train_split_name].train_test_split(train_size=split, seed=seed)
+    nst[data_args.train_split_name] = nst_train["train"]
+    nst[data_args.eval_split_name] = nst_train["test"]
+
+    nst = nst.filter(filter_nst).map(
+        map_nst,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NST",
+    ).shuffle(seed=seed)
+    npsc = npsc.filter(filter_npsc).map(
+        map_npsc,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NPSC",
+    ).shuffle(seed=seed)
+
+    npsc_base = npsc.remove_columns([col for col in npsc[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+    nst_base = nst.remove_columns([col for col in nst[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+
+    combined = {}
+    for split in data_args.train_split_name, data_args.eval_split_name, data_args.test_split_name:
+        probs = np.array([len(nst_base[split]), len(npsc_base[split])])  # Weight by number of examples
+        probs = (probs / probs.sum()).tolist()
+        comb = datasets.interleave_datasets([nst_base[split], npsc_base[split]], probabilities=probs, seed=seed)
+        combined[split] = comb
+
+    return datasets.DatasetDict(**combined)
+
+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, FlaxTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # 2. 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()
+
+    # Set up wandb run
+    if jax.process_index() == 0:
+        wandb.init(project=data_args.wandb_project, name=data_args.wandb_name, job_type=data_args.wandb_job_type)
+
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set the default TPU matmul precision and display the number of devices
+    jax.config.update("jax_default_matmul_precision", training_args.matmul_precision)
+    logger.info(f"JAX devices: {jax.device_count()}, matmul precision: {training_args.matmul_precision}")
+
+    # 4. Load dataset
+
+    set_seed(training_args.seed)
+    raw_datasets = make_dataset(data_args, seed=training_args.seed)
+   
+    # raw_datasets = DatasetDict()
+
+    # if training_args.do_train:
+    #     raw_datasets[data_args.train_split_name] = 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[data_args.eval_split_name] = 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 training_args.do_predict:
+    #     test_split = data_args.test_split_name.split("+")
+    #     for split in test_split:
+    #         raw_datasets[split] = load_dataset(
+    #             data_args.dataset_name,
+    #             data_args.dataset_config_name,
+    #             split=split,
+    #             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 and not training_args.do_predict:
+        raise ValueError(
+            "Cannot not train, not do evaluation and not do prediction. At least one of "
+            "training, evaluation or prediction has to be done."
+        )
+
+    # if not training, there is no need to run multiple epochs
+    if not training_args.do_train:
+        training_args.num_train_epochs = 1
+
+    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
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    config = Wav2Vec2Config.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,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    # update config according to training args, model args, and tokenizer attributes
+    config.update(
+        {
+            "feat_proj_dropout": model_args.feat_proj_dropout,
+            "attention_dropout": model_args.attention_dropout,
+            "hidden_dropout": model_args.hidden_dropout,
+            "final_dropout": model_args.final_dropout,
+            "mask_time_prob": model_args.mask_time_prob,
+            "mask_time_length": model_args.mask_time_length,
+            "mask_feature_prob": model_args.mask_feature_prob,
+            "mask_feature_length": model_args.mask_feature_length,
+            "gradient_checkpointing": training_args.gradient_checkpointing,
+            "layerdrop": model_args.layerdrop,
+            "ctc_loss_reduction": model_args.ctc_loss_reduction,
+            "ctc_zero_infinity": model_args.ctc_zero_infinity,
+            "pad_token_id": tokenizer.pad_token_id,
+            "vocab_size": tokenizer.vocab_size,  # len(tokenizer),
+            "activation_dropout": model_args.activation_dropout,
+        }
+    )
+
+    if tokenizer.do_lower_case and data_args.dataset_name != "librispeech_asr":
+        raise ValueError(
+            "Setting the tokenizer attribute `do_lower_case` to `True` converts all input strings to "
+            "uppercase prior to tokenization. This should only be done when the tokenizer is built on an uppercased corpus,"
+            "i.e. for the dataset `librispeech_asr` only. If your dataset is not `librispeech_asr`, the tokenizer is mostly likely "
+            "built on an lowercased corpus. In this case, set `tokenizer.do_lower_case` to ``False`."
+        )
+
+    if training_args.precision == "full_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = True
+    elif training_args.precision == "half_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = False
+    else:
+        dtype = jnp.float32
+        training_args.mixed_precision = False
+
+    try:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+        )
+    except:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+            from_pt=True,
+        )
+
+    # 6. Resample speech dataset ALWAYS
+    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_target_length = data_args.max_label_length
+    min_target_length = data_args.min_label_length
+    pad_input_to_multiple_of = data_args.pad_input_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
+    dataset_name = data_args.dataset_name
+    chars_to_ignore = ', ? . ! - ; : " “ % ‘ ” ?'.split(" ")
+    chars_to_ignore_regex = f'[{"".join(chars_to_ignore)}]'
+    # gigaspeech_punctuation = {" <comma>": ",", " <period>": ".", " <questionmark>": "?", " <exclamationpoint>": "!"}
+    # gigaspeech_disfluencies = ["<other>", "<sil>"]
+    # swb_disfluencies = ["[noise]", "[laughter]", "[silence]", "<a_aside>", "<b_aside>", "<e_aside>", "[laughter-",
+    #                 "[vocalized-noise]", "_1"]
+    # swb_punctuations = ["{", "}", "[", "]-", "]"]
+    # earnings_disfluencies = ["<crosstalk>", "<affirmative>", "<inaudible>", "inaudible", "<laugh>", "<unk>"]
+    ignore_segments = ["ignore_time_segment_in_scoring", "<noise>", "<music>", "[noise]", "[laughter]", "[silence]",
+                       "[vocalized-noise]", "<crosstalk>", "<affirmative>", "<inaudible>", "<laugh>", "<other>", "<sil>", ""]
+
+    if training_args.do_train and data_args.max_train_samples is not None:
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].select(range(data_args.max_train_samples))
+
+    if training_args.do_eval and data_args.max_eval_samples is not None:
+        raw_datasets[data_args.eval_split_name] = raw_datasets[data_args.eval_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_predict and data_args.max_test_samples is not None:
+        raw_datasets[data_args.test_split_name] = raw_datasets[data_args.test_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_train and data_args.remove_punctuation:
+
+        def remove_punctuation(batch):
+            batch[text_column_name] = (
+                re.sub(chars_to_ignore_regex, "", batch[text_column_name]).replace("'", "").replace('"', "")
+            )
+
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].map(
+            remove_punctuation,
+            num_proc=data_args.preprocessing_num_workers,
+            desc="removing punctuation from train split",
+        )
+
+    # filter data where the targets are ignored in scoring
+    def is_target_labels(input_str):
+        return input_str.lower() not in ignore_segments
+
+    raw_datasets = raw_datasets.filter(
+            is_target_labels,
+            num_proc=num_workers,
+            input_columns=[text_column_name],
+            desc="filtering data where the targets are ignored in scoring",
+        )
+
+    def prepare_dataset(batch):
+        # process audio
+        try:
+            sample = batch[audio_column_name]
+        except ValueError:
+            sample = {"array": np.array([0.]), "sampling_rate": feature_extractor.sampling_rate}
+        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
+        # process audio length
+        batch[model_input_name] = inputs.input_values[0]
+        batch["input_length"] = len(batch["input_values"])
+
+        # process targets
+        input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
+
+        # if dataset_name == "google/xtreme_s":
+        #     # Finally, we tokenize the processed text
+        #     batch["labels"] = tokenizer(input_str).input_ids
+        #     batch["labels_length"] = len(batch["labels"])
+        #     return batch
+
+        # # Common Voice 9
+        # if input_str.startswith('"') and input_str.endswith('"'):
+        #     # we can remove trailing quotation marks as they do not affect the transcription
+        #     input_str = input_str[1:-1]
+        # # normalize quotation marks
+        # input_str = re.sub(r'["“”]', '"', input_str)
+        # # normalize apostrophes
+        # input_str = re.sub(r"[’']", "'", input_str)
+        # # normalize hyphens
+        # input_str = re.sub(r"[—–]", "-", input_str)
+        # # replace double quotation marks with single
+        # input_str = input_str.replace('""', '"')
+        # if dataset_name == "mozilla-foundation/common_voice_9_0" and len(input_str):
+        #     # for CV9, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # TEDLIUM-3
+        # # delete the <unk> token from the text and replace spaced apostrophes with un-spaced
+        # input_str = input_str.replace("<unk>", "").replace(" '", "'")
+
+        # # GigaSpeech
+        # for disfluency in gigaspeech_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # convert spelled out punctuation to symbolic form
+        # for punctuation, replacement in gigaspeech_punctuation.items():
+        #     input_str = input_str.replace(punctuation, replacement)
+        # if dataset_name == "speechcolab/gigaspeech" and len(input_str):
+        #     # for GS, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # SWB
+        # for disfluency in swb_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # remove parenthesised text (test data only)
+        # input_str = re.sub("[\(].*?[\)]", "", input_str)
+        # for punctuation in swb_punctuations:
+        #     input_str = input_str.replace(punctuation, "")
+        # # replace anomalous words with their correct transcriptions
+        # split_str = input_str.split("/")
+        # if len(split_str) > 1:
+        #     input_str = " ".join(
+        #         [" ".join([" ".join(i.split(" ")[:-1]) for i in split_str])] + [split_str[-1].split(" ")[-1]])
+
+        # # Earnings 22
+        # for disfluency in earnings_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # replace mal-formatted ellipsis
+        # input_str = input_str.replace("…", ".")
+
+        # JIWER compliance
+        # remove multiple spaces
+        input_str = re.sub(r"\s\s+", " ", input_str)
+        # strip trailing spaces
+        input_str = input_str.strip()
+
+        # Finally, we tokenize the processed text
+        batch["labels"] = tokenizer(input_str).input_ids
+        batch["labels_length"] = len(batch["labels"])
+        return batch
+
+    vectorized_datasets = raw_datasets.map(
+        prepare_dataset,
+        remove_columns=next(iter(raw_datasets.values())).column_names,
+        num_proc=num_workers,
+        desc="preprocess dataset",
+    )
+
+    # filter data with inputs shorter than min_input_length or longer than max_input_length
+    def is_audio_in_length_range(length):
+        return length > min_input_length and length < max_input_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_audio_in_length_range,
+        num_proc=num_workers,
+        input_columns=["input_length"],
+    )
+
+    # filter data with targets shorter than min_target_length or longer than max_target_length
+    def is_labels_in_length_range(length):
+        return length > min_target_length  # and length < max_target_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_labels_in_length_range,
+        num_proc=num_workers,
+        input_columns=["labels_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 Metrics
+    wer_metric = load_metric("wer")
+    cer_metric = load_metric("cer")
+
+    def compute_metrics(pred_ids: List[List[int]], label_ids: List[List[int]]):
+        padded_ids = np.where(np.asarray(label_ids) == -100, tokenizer.pad_token_id, np.asarray(label_ids))
+
+        pred_str = tokenizer.batch_decode(pred_ids)
+        # we do not want to group tokens when computing the metrics
+        label_str = tokenizer.batch_decode(padded_ids, group_tokens=False)
+
+        wer = wer_metric.compute(predictions=pred_str, references=label_str)
+        cer = cer_metric.compute(predictions=pred_str, references=label_str)
+
+        return {"wer": wer, "cer": cer}, pred_str, label_str
+
+    # 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,
+        input_padding="longest",
+        pad_input_to_multiple_of=pad_input_to_multiple_of,
+        max_label_length=data_args.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."
+        )
+
+    # 10. Handle the repository creation
+    if training_args.push_to_hub:
+        with open(os.path.join(training_args.output_dir, ".gitattributes"), "r+") as f:
+            git_lfs_extensions = f.read()
+            if "*.wandb" not in git_lfs_extensions:
+                f.write("*.wandb filter=lfs diff=lfs merge=lfs -text")
+        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
+        repo = Repository(training_args.output_dir, clone_from=repo_name)
+
+    # 11. Initialize our training
+    rng = jax.random.PRNGKey(training_args.seed)
+    rng, dropout_rng = jax.random.split(rng)
+
+    # Store some constants
+    max_steps = int(training_args.max_steps)
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
+    batch_size_per_update = train_batch_size * gradient_accumulation_steps
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
+    to_dtype = to_bf16 if training_args.mixed_precision else to_fp32
+
+    if training_args.do_train:
+        num_train_samples = len(vectorized_datasets[data_args.train_split_name])
+        steps_per_epoch = num_train_samples // batch_size_per_update
+        if max_steps > 0:
+            num_epochs = -(training_args.max_steps // -steps_per_epoch)
+            total_train_steps = max_steps
+        else:
+            num_epochs = int(training_args.num_train_epochs)
+            total_train_steps = steps_per_epoch * num_epochs
+
+        # Create learning rate schedule
+        # Create learning rate schedule
+        linear_decay_lr_schedule_fn = create_learning_rate_fn(
+            total_train_steps,
+            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.
+        # Note that this mask is specifically adapted for FlaxWav2Vec2 and FlaxBart.
+        # For FlaxT5, one should correct the layer norm parameter naming
+        # accordingly - see `run_t5_mlm_flax.py` e.g.
+        def decay_mask_fn(params):
+            flat_params = traverse_util.flatten_dict(params)
+            layer_norm_params = [
+                (name, "scale")
+                for name in ["layer_norm", "self_attn_layer_norm", "layernorm_embedding", "final_layer_norm"]
+            ]
+            flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_params) for path in flat_params}
+            return traverse_util.unflatten_dict(flat_mask)
+
+        if training_args.adafactor:
+            # Create Adafactor optimizer
+            optim = optax.adafactor(
+                learning_rate=linear_decay_lr_schedule_fn,
+                dtype_momentum=jnp.bfloat16 if training_args.mixed_precision else jnp.float32,
+                weight_decay_rate=training_args.weight_decay,
+                weight_decay_mask=decay_mask_fn,
+            )
+        else:
+            # Create AdamW optimizer
+            optim = 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,
+            )
+
+        # Optax MultiSteps for gradient accumulation. We'll only call this optimizer transformation if gradient accumulation is required (i.e. gradient accumulation steps > 1)
+        if training_args.multisteps and gradient_accumulation_steps > 1:
+            optim = optax.MultiSteps(optim, gradient_accumulation_steps, use_grad_mean=False)
+    else:
+        num_epochs = 0
+        total_train_steps = 0
+        num_train_samples = 0
+        optim = None
+
+    # Setup train state
+    state = MixedPrecisionTrainState.create(
+        step=data_args.skip_steps,
+        apply_fn=model.__call__,
+        get_attention_mask_fn=model._get_feature_vector_attention_mask,
+        params=model.params,
+        tx=optim,
+        to_dtype=to_dtype,
+        dropout_rng=dropout_rng,
+        max_grad_norm=training_args.max_grad_norm,
+    )
+
+    # Replicate the train state on each device
+    state = state.replicate()
+    blank_id = model.config.pad_token_id
+
+    # Define gradient update step fn
+    def train_step(state, batch):
+        # only one single rng per grad step, with or without accumulation, as the graph should be identical over one effective training batch
+        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
+
+        def compute_loss(params, minibatch):
+            labels = minibatch.pop("labels")
+            logits = state.apply_fn(
+                **minibatch,
+                params=params,
+                dropout_rng=dropout_rng,
+                freeze_feature_encoder=model_args.freeze_feature_encoder,
+                train=True,
+            )[0]
+            logits_mask = state.get_attention_mask_fn(logits.shape[1], batch["attention_mask"])
+            loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+            return loss
+
+        grad_fn = jax.value_and_grad(compute_loss)
+
+        if gradient_accumulation_steps == 1 or training_args.multisteps:
+            loss, grad = grad_fn(to_dtype(state.params), batch)
+
+        # Custom gradient accumulation
+        else:
+            # add a first dimension over gradient_accumulation_steps for minibatch slices
+            batch = jax.tree_util.tree_map(
+                lambda x: x.reshape(
+                    gradient_accumulation_steps, training_args.per_device_train_batch_size, *x.shape[1::]
+                ),
+                batch,
+            )
+
+            def accum_minibatch_step(accum_grad, minibatch):
+                # compute loss, num labels and grad over minibatch and accumulate
+                loss, grad = grad_fn(to_dtype(state.params), minibatch)
+                return jax.tree_util.tree_map(jnp.add, accum_grad, grad), loss
+
+            # create an initial state for accumulating losses, num labels and gradients
+            init_grad = jax.tree_util.tree_map(jnp.zeros_like, to_dtype(state.params))
+            # loop accum minibatch step over the number of gradient accumulation steps
+            grad, loss = jax.lax.scan(accum_minibatch_step, init_grad, batch)
+
+        # update state
+        new_state = state.apply_gradients(
+            grads=grad,
+            dropout_rng=new_dropout_rng,
+            to_dtype=to_dtype,
+        )
+
+        # compute gradient norms over all layers and globally for detailed monitoring
+        layer_grad_norm = jax.tree_util.tree_map(jnp.linalg.norm, grad)
+        logs = {
+            "layer_grad_norm": layer_grad_norm,
+            "grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm)),
+        }
+
+        # compute parameter norms over all layers and globally for detailed monitoring
+        layer_param_norm = jax.tree_util.tree_map(jnp.linalg.norm, new_state.params)
+        logs["layer_param_norm"] = layer_param_norm
+        logs["param_norm"] = jnp.linalg.norm(jax.tree_util.tree_leaves(layer_param_norm))
+
+        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+        metrics.update(logs)
+
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+
+        return new_state, metrics
+
+    # Define eval fn
+    def eval_step(params, batch):
+        labels = batch.pop("labels")
+        logits = model(**batch, params=params, train=False)[0]
+
+        logits_mask = model._get_feature_vector_attention_mask(logits.shape[1], batch["attention_mask"])
+        loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+        pred_ids = jnp.argmax(logits, axis=-1)
+
+        # summarize metrics
+        metrics = {"loss": loss}
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+        return metrics, pred_ids
+
+    # Create parallel version of the train and eval step
+    if training_args.do_train:
+        p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
+
+    if training_args.do_eval:
+        p_eval_step = jax.pmap(eval_step, "batch")
+
+    def run_evaluation(step):
+        if training_args.do_eval:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the eval dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[data_args.eval_split_name], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
+                samples = [vectorized_datasets[data_args.eval_split_name][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, step, prefix="eval")
+            write_wandb_pred(pred_str, label_str, step)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, step, pred_str=pred_str)
+
+    def save_checkpoint(step):
+        # save 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"{wandb.run.id}: saving weights and logs of step {int(step / 1000)}k", blocking=False)
+
+    skip_epochs = data_args.skip_steps // (num_train_samples // batch_size_per_update)
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_train_samples}")
+    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"  Num gradient accumulation steps = {gradient_accumulation_steps}")
+    logger.info(f"  Total train batch size (w. parallel & distributed) = {batch_size_per_update}")
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+    logger.info(f"  Gradient checkpointing: {config.gradient_checkpointing}")
+    logger.info(f"  Use scan: {config.use_scan}")
+    logger.info(f"  Fuse matmuls: {config.fuse_matmuls}")
+    logger.info(f"  Skipping: {data_args.skip_steps} steps ({skip_epochs} epochs)")
+
+    train_time = cur_step = 0
+    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
+    for epoch in epochs:
+        if training_args.do_train:
+            # ======================== Training ================================
+            train_start = time.time()
+
+            if epoch < skip_epochs:
+                logger.info(f"Skipping epoch {epoch + 1}")
+                # Create sampling rng
+                rng, input_rng = jax.random.split(rng)
+                continue
+
+            # Generate an epoch by randomly shuffling sampling indices from the train dataset and grouping by length
+            train_samples_idx = get_grouped_indices(vectorized_datasets[data_args.train_split_name], batch_size_per_update, input_rng)
+            train_batch_idx = generate_batch_splits(train_samples_idx, batch_size_per_update)
+
+            if data_args.skip_steps > cur_step:
+                logger.info(f"Skipping {data_args.skip_steps - (epoch * (num_train_samples // batch_size_per_update))} steps...")
+            # Gather the indices for creating the batch and do a training step
+            for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1), 1):
+                cur_step = epoch * (num_train_samples // batch_size_per_update) + step
+                if cur_step <= data_args.skip_steps:
+                    continue
+
+                samples = [vectorized_datasets[data_args.train_split_name][int(idx)] for idx in batch_idx]              
+                batch = data_collator(samples)
+                batch = shard(batch.data)
+                try:
+                    state, train_metric = p_train_step(state, batch)
+                except TypeError as e:
+                    logger.warning("Encountered following error: \n", e)
+
+
+                if cur_step % training_args.logging_steps == 0:
+                    # Save metrics
+                    train_metric = unreplicate(train_metric)
+                    train_time += time.time() - train_start
+                    # need to upcast all device arrays to fp32 for wandb logging (jnp.bfloat16 not supported) -> do this here OR in train_step
+                    write_wandb_log(to_fp32(train_metric), cur_step, prefix=data_args.train_split_name)
+                    # we won't log to tensorboard for now (it is fiddly logging param and grad norms on a layer-by-layer basis)
+                    # if has_tensorboard and jax.process_index() == 0:
+                    # write_train_metric(summary_writer, train_metrics, train_time, cur_step)
+
+                    epochs.write(
+                        f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']}, Gradient Norm: {train_metric['grad_norm']})"
+                    )
+                    p_train_step.clear_cache()
+
+                if cur_step % total_train_steps == 0:
+                    break
+
+                if training_args.eval_steps and cur_step % training_args.eval_steps == 0:
+                    run_evaluation(cur_step)
+
+                if cur_step % training_args.save_steps == 0:
+                    save_checkpoint(cur_step)
+
+            if training_args.eval_steps == 0 and (epoch + 1) != num_epochs:
+                # run evaluation at the end of the epoch if eval steps are not specified
+                run_evaluation(cur_step)
+                save_checkpoint(cur_step)
+
+    if training_args.do_train:
+        save_checkpoint(cur_step)
+
+    cur_step = max_steps if max_steps > 0 else cur_step  # set step to max steps so that eval happens in alignment with training
+
+    if training_args.do_eval:
+        run_evaluation(cur_step)
+
+    # TODO: collapse 'do_predict' into the run_evaluation function
+    if training_args.do_predict:
+        for split in [data_args.test_split_name]:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the test dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[split], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc=f"Predicting {split}...", position=2)):
+                samples = [vectorized_datasets[split][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({cur_step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, cur_step, prefix=split)
+            write_wandb_pred(pred_str, label_str, cur_step, prefix=split)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, cur_step, pred_str=pred_str)
+
+
+if __name__ == "__main__":
+    main()

wandb/run-20220810_145446-1k92sv35/files/config.yaml

@@ -0,0 +1,33 @@
+wandb_version: 1
+
+_wandb:
+  desc: null
+  value:
+    cli_version: 0.12.9
+    code_path: code/run_flax_speech_recognition_ctc.py
+    framework: huggingface
+    huggingface_version: 4.21.0
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    python_version: 3.8.10
+    start_time: 1660143286
+    t:
+      1:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      2:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      3:
+      - 13
+      4: 3.8.10
+      5: 0.12.9
+      6: 4.21.0
+      8:
+      - 5

wandb/run-20220810_145446-1k92sv35/files/diff.patch

@@ -0,0 +1,132 @@
+diff --git a/run.recover.sh b/run.recover.sh
+index 77ad3fd..6891af1 100755
+--- a/run.recover.sh
++++ b/run.recover.sh
+@@ -10,10 +10,9 @@ WANDB_ENTITY=NbAiLab WANDB_PROJECT=wav2vec2 python run_flax_speech_recognition_c
+         --num_train_epochs="40" \
+         --per_device_train_batch_size="2" \
+         --per_device_eval_batch_size="2" \
+-        --gradient_accumulation_steps="1" \
+-        --precision="full_mixed" \
++        --gradient_accumulation_steps="2" \
++        --precision="half_mixed" \
+         --matmul_precision="bfloat16" \
+-        --multisteps \
+         --learning_rate="6.394633237505332e-05" \
+         --skip_steps="275000" \
+         --warmup_steps="2000" \
+diff --git a/run_flax_speech_recognition_ctc.py b/run_flax_speech_recognition_ctc.py
+index a330879..4d0f5fc 100644
+--- a/run_flax_speech_recognition_ctc.py
++++ b/run_flax_speech_recognition_ctc.py
+@@ -415,12 +415,12 @@ class MixedPrecisionTrainState(struct.PyTreeNode):
+         )
+ 
+     @classmethod
+-    def create(cls, *, apply_fn, params, tx, to_dtype, **kwargs):
++    def create(cls, *, apply_fn, params, tx, to_dtype, step=0, **kwargs):
+         """Creates a new instance with `step=0` and initialized `opt_state`."""
+         # downcast optimizer state to bf16 if mixed-precision training
+         opt_state = tx.init(to_dtype(params)) if tx is not None else None
+         return cls(
+-            step=0,
++            step=step,
+             apply_fn=apply_fn,
+             params=params,
+             tx=tx,
+@@ -1339,6 +1339,7 @@ def main():
+ 
+     # Setup train state
+     state = MixedPrecisionTrainState.create(
++        step=data_args.skip_steps,
+         apply_fn=model.__call__,
+         get_attention_mask_fn=model._get_feature_vector_attention_mask,
+         params=model.params,
+@@ -1520,11 +1521,10 @@ def main():
+ 
+             if epoch < skip_epochs:
+                 logger.info(f"Skipping epoch {epoch + 1}")
++                # Create sampling rng
++                rng, input_rng = jax.random.split(rng)
+                 continue
+ 
+-            # Create sampling rng
+-            rng, input_rng = jax.random.split(rng)
+-
+             # Generate an epoch by randomly shuffling sampling indices from the train dataset and grouping by length
+             train_samples_idx = get_grouped_indices(vectorized_datasets[data_args.train_split_name], batch_size_per_update, input_rng)
+             train_batch_idx = generate_batch_splits(train_samples_idx, batch_size_per_update)
+@@ -1559,6 +1559,7 @@ def main():
+                     epochs.write(
+                         f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']}, Gradient Norm: {train_metric['grad_norm']})"
+                     )
++                    p_train_step.clear_cache()
+ 
+                 if cur_step % total_train_steps == 0:
+                     break
+diff --git a/special_tokens_map.json b/special_tokens_map.json
+index 218961f..cc1961e 100644
+--- a/special_tokens_map.json
++++ b/special_tokens_map.json
+@@ -399,6 +399,34 @@
+       "rstrip": false,
+       "single_word": false
+     },
++    {
++      "content": "</s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "<s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "</s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "<s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
+     {
+       "content": "</s>",
+       "lstrip": false,
+diff --git a/wandb/debug-internal.log b/wandb/debug-internal.log
+index 23926ef..aef858a 120000
+--- a/wandb/debug-internal.log
++++ b/wandb/debug-internal.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug-internal.log
+
++run-20220810_145446-1k92sv35/logs/debug-internal.log
+
+diff --git a/wandb/debug.log b/wandb/debug.log
+index 279853d..0d5686d 120000
+--- a/wandb/debug.log
++++ b/wandb/debug.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug.log
+
++run-20220810_145446-1k92sv35/logs/debug.log
+
+diff --git a/wandb/latest-run b/wandb/latest-run
+index f069a7a..3128ad6 120000
+--- a/wandb/latest-run
++++ b/wandb/latest-run
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4
+
++run-20220810_145446-1k92sv35
+

wandb/run-20220810_145446-1k92sv35/files/output.log

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

wandb/run-20220810_145446-1k92sv35/files/requirements.txt

@@ -0,0 +1,158 @@
+absl-py==1.0.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+appdirs==1.4.4
+astunparse==1.6.3
+async-timeout==4.0.2
+attrs==21.4.0
+audioread==2.1.9
+backcall==0.2.0
+cachetools==4.2.4
+certifi==2021.10.8
+cffi==1.15.1
+charset-normalizer==2.0.10
+chex==0.1.3
+click==8.0.3
+cloud-tpu-client==0.10
+cloud-tpu-profiler==2.4.0
+clu==0.0.6
+colorama==0.4.5
+commonmark==0.9.1
+configparser==5.2.0
+contextlib2==21.6.0
+cycler==0.11.0
+datasets==2.4.0
+decorator==5.1.0
+dill==0.3.4
+dm-tree==0.1.6
+docker-pycreds==0.4.0
+etils==0.6.0
+exceptiongroup==1.0.0rc8
+filelock==3.4.2
+flatbuffers==2.0
+flax==0.5.3
+fonttools==4.28.5
+frozenlist==1.2.0
+fsspec==2021.11.1
+future==0.18.2
+gast==0.4.0
+gitdb==4.0.9
+gitpython==3.1.26
+google-api-core==1.31.5
+google-api-python-client==1.8.0
+google-auth-httplib2==0.1.0
+google-auth-oauthlib==0.4.6
+google-auth==2.3.3
+google-pasta==0.2.0
+googleapis-common-protos==1.54.0
+grpcio==1.43.0
+h5py==3.6.0
+httplib2==0.20.2
+huggingface-hub==0.2.1
+hypothesis==6.53.0
+idna==3.3
+importlib-metadata==4.10.0
+importlib-resources==5.4.0
+ipython==7.31.0
+jax==0.3.15
+jaxlib==0.3.15
+jedi==0.18.1
+jiwer==2.3.0
+joblib==1.1.0
+keras-preprocessing==1.1.2
+keras==2.7.0
+kiwisolver==1.3.2
+libclang==12.0.0
+librosa==0.9.2
+libtpu-nightly==0.1.dev20220722
+llvmlite==0.39.0
+markdown==3.3.6
+matplotlib-inline==0.1.3
+matplotlib==3.5.1
+ml-collections==0.1.0
+msgpack==1.0.3
+multidict==5.2.0
+multiprocess==0.70.12.2
+numba==0.56.0
+numpy==1.22.0
+oauth2client==4.1.3
+oauthlib==3.1.1
+opt-einsum==3.3.0
+optax==0.1.3
+packaging==21.3
+pandas==1.3.5
+parso==0.8.3
+pathtools==0.1.2
+pexpect==4.8.0
+pickleshare==0.7.5
+pillow==9.0.0
+pip==22.2.2
+pkg-resources==0.0.0
+pooch==1.6.0
+promise==2.3
+prompt-toolkit==3.0.24
+protobuf==3.19.1
+psutil==5.9.0
+ptyprocess==0.7.0
+pyarrow==6.0.1
+pyasn1-modules==0.2.8
+pyasn1==0.4.8
+pycparser==2.21
+pyctcdecode==0.4.0
+pygments==2.11.1
+pygtrie==2.5.0
+pyparsing==3.0.6
+python-dateutil==2.8.2
+python-levenshtein==0.12.2
+pytz==2021.3
+pyyaml==6.0
+regex==2021.11.10
+requests-oauthlib==1.3.0
+requests==2.27.0
+resampy==0.3.1
+responses==0.18.0
+rich==11.2.0
+rsa==4.8
+sacremoses==0.0.46
+scikit-learn==1.1.1
+scipy==1.7.3
+sentry-sdk==1.5.2
+setuptools==44.0.0
+shortuuid==1.0.8
+six==1.16.0
+smmap==5.0.0
+sortedcontainers==2.4.0
+soundfile==0.10.3.post1
+sox==1.4.1
+subprocess32==3.5.4
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.0
+tensorboard==2.7.0
+tensorflow-cpu==2.7.0
+tensorflow-datasets==4.4.0
+tensorflow-estimator==2.7.0
+tensorflow-io-gcs-filesystem==0.23.1
+tensorflow-metadata==1.5.0
+tensorflow==2.7.0
+tensorstore==0.1.21
+termcolor==1.1.0
+threadpoolctl==3.1.0
+tokenizers==0.11.2
+toolz==0.11.2
+torch==1.12.0
+torchaudio==0.12.0+cpu
+tqdm==4.62.3
+traitlets==5.1.1
+transformers==4.21.0
+typing-extensions==4.3.0
+uritemplate==3.0.1
+urllib3==1.26.7
+wandb==0.12.9
+wcwidth==0.2.5
+werkzeug==2.0.2
+wheel==0.37.1
+wrapt==1.13.3
+xxhash==2.0.2
+yarl==1.7.2
+yaspin==2.1.0
+zipp==3.7.0

wandb/run-20220810_145446-1k92sv35/files/wandb-metadata.json

@@ -0,0 +1,69 @@
+{
+    "os": "Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29",
+    "python": "3.8.10",
+    "heartbeatAt": "2022-08-10T14:54:50.575340",
+    "startedAt": "2022-08-10T14:54:46.729335",
+    "docker": null,
+    "cpu_count": 96,
+    "cuda": null,
+    "args": [
+        "--model_name_or_path=./",
+        "--hub_model_id=NbAiLab/wav2vec2-1b-npsc-nst-tpu",
+        "--tokenizer_name=./",
+        "--output_dir=./",
+        "--overwrite_output_dir",
+        "--num_train_epochs=40",
+        "--per_device_train_batch_size=2",
+        "--per_device_eval_batch_size=2",
+        "--gradient_accumulation_steps=2",
+        "--precision=half_mixed",
+        "--matmul_precision=bfloat16",
+        "--learning_rate=6.394633237505332e-05",
+        "--skip_steps=275000",
+        "--warmup_steps=2000",
+        "--length_column_name=input_length",
+        "--evaluation_strategy=steps",
+        "--text_column_name=text",
+        "--save_steps=5000",
+        "--eval_steps=5000",
+        "--logging_steps=100",
+        "--layerdrop=0.041",
+        "--attention_dropout=0.094",
+        "--activation_dropout=0.055",
+        "--hidden_dropout=0.047",
+        "--save_total_limit=5",
+        "--freeze_feature_encoder",
+        "--feat_proj_dropout=0.04",
+        "--mask_time_prob=0.082",
+        "--mask_time_length=10",
+        "--mask_feature_prob=0.25",
+        "--mask_feature_length=64",
+        "--gradient_checkpointing",
+        "--min_duration_in_seconds=0.5",
+        "--max_duration_in_seconds=30.0",
+        "--use_auth_token",
+        "--seed=42",
+        "--group_by_length",
+        "--do_train",
+        "--do_eval",
+        "--push_to_hub",
+        "--preprocessing_num_workers=32",
+        "--ctc_zero_infinity",
+        "--do_lower_case",
+        "--wandb_project=wav2vec2",
+        "--wandb_name=wav2vec2-1b-npsc-nst-tpu (cont.)",
+        "--remove_punctuation"
+    ],
+    "state": "running",
+    "program": "run_flax_speech_recognition_ctc.py",
+    "codePath": "run_flax_speech_recognition_ctc.py",
+    "git": {
+        "remote": "https://huggingface.co/NbAiLab/wav2vec2-1b-npsc-nst-tpu",
+        "commit": "f624ac4bfedfbf56891676d7c5f2e37b4c8e0745"
+    },
+    "email": "versae@gmail.com",
+    "root": "/data/wav2vec2-1b-npsc-nst-tpu",
+    "host": "t1v-n-eedfb410-w-0",
+    "username": "javierr",
+    "executable": "/data/flax/bin/python"
+}

wandb/run-20220810_145446-1k92sv35/files/wandb-summary.json

@@ -0,0 +1 @@
+{"_wandb": {"runtime": 727}}

wandb/run-20220810_145446-1k92sv35/logs/debug-internal.log

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

wandb/run-20220810_145446-1k92sv35/logs/debug.log

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

wandb/run-20220810_145446-1k92sv35/run-1k92sv35.wandb

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

wandb/run-20220810_151736-2jo5la5b/files/code/run_flax_speech_recognition_ctc.py

@@ -0,0 +1,1632 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 The HuggingFace 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 connectionist temporal classification (CTC) 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 re
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import datasets
+import numpy as np
+from datasets import DatasetDict, load_dataset, load_metric
+from tqdm import tqdm
+
+import flax
+import jax
+import jax.numpy as jnp
+import optax
+import transformers
+import wandb as wandb
+from flax import core, jax_utils, struct, traverse_util
+from flax.jax_utils import unreplicate, pad_shard_unpad
+from flax.training.common_utils import get_metrics, shard, shard_prng_key
+from huggingface_hub import Repository
+from models import Wav2Vec2Config, FlaxWav2Vec2ForCTC
+from optax._src import linear_algebra
+from transformers import (
+    AutoFeatureExtractor,
+    AutoProcessor,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainingArguments,
+    is_tensorboard_available,
+    set_seed,
+)
+from transformers.file_utils import get_full_repo_name
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.17.0.dev0")
+
+require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/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)."
+        },
+    )
+    freeze_feature_encoder: bool = field(
+        default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}
+    )
+    attention_dropout: float = field(
+        default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}
+    )
+    activation_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The hidden activation dropout probability in the embeddings, encoder, and pooler."
+        },
+    )
+    hidden_dropout: float = field(
+        default=0.1,
+        metadata={
+            "help": "The dropout probability for all fully connected layers in the embeddings, encoder, and pooler."
+        },
+    )
+    feat_proj_dropout: float = field(
+        default=0.0,
+        metadata={
+            "help": "The feat proj dropout probability for feature encoder representations."
+        },
+    )
+    final_dropout: float = field(
+        default=0.0,
+        metadata={"help": "The dropout probability for the final projection layer."},
+    )
+    mask_time_prob: float = field(
+        default=0.1,
+        metadata={
+            "help": "The spec aug dropout probability for feature encoder representations."
+        },
+    )
+    mask_time_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the time axis."},
+    )
+    mask_feature_prob: float = field(
+        default=0.0,
+        metadata={
+            "help": "Probability of each feature vector along the feature axis to be chosen as the start of the vector"
+                    "span to be masked. Approximately ``mask_feature_prob * sequence_length // mask_feature_length`` feature bins will be masked along the time axis."
+        },
+    )
+    mask_feature_length: int = field(
+        default=10,
+        metadata={"help": "Length of vector span to mask along the feature axis."},
+    )
+    layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})
+    ctc_loss_reduction: Optional[str] = field(
+        default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}
+    )
+    ctc_zero_infinity: Optional[bool] = field(
+        default=False, metadata={"help": "If True, will try yo aboud the CTC loss goinf to infinity."}
+    )
+
+
+@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."
+        },
+    )
+    max_test_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of test 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": "Truncate audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_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: Optional[int] = field(
+        default=512,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    min_label_length: Optional[int] = field(
+        default=2,
+        metadata={
+            "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
+            "than this will be filtered."
+        },
+    )
+    pad_input_to_multiple_of: Optional[int] = field(
+        default=32000,
+        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."
+        },
+    )
+    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."
+        },
+    )
+    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 training data set split to use (via the datasets library). Defaults to 'train'"
+        },
+    )
+    do_lower_case: bool = field(
+        default=True,
+        metadata={"help": "Whether the target text should be lower cased."},
+    )
+    wandb_project: str = field(
+        default="flax-speech-recognition-ctc",
+        metadata={"help": "The name of the wandb project."},
+    )
+    wandb_name: str = field(
+        default=None,
+        metadata={"help": "The name of the wandb run."},
+    )
+    wandb_job_type: str = field(
+        default="CTC",
+        metadata={"help": "The name of the wandb job type."},
+    )
+    test_split_name: str = field(
+        default="test",
+        metadata={"help": "The name of the test data set split to use (via the datasets library). Defaults to 'test'"},
+    )
+    remove_punctuation: bool = field(
+        default=False, metadata={"help": "Whether or not to remove punctuation during training."}
+    )
+    skip_steps: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": "Skip this number of steps. Useful to continue training"
+        },
+    )
+
+
+# @flax.struct.dataclass
+@dataclass
+class FlaxTrainingArguments(TrainingArguments):
+    precision: str = field(
+        default="full",
+        metadata={
+            "help": "Whether to enable mixed-precision training. If true, the optimizer is stored in half-precision (bfloat16) and computations are executed in half-precision"
+            "**Note that this only specifies the dtype of the computation and optimizer state. It does not influence the dtype of model parameters.**"
+        },
+    )
+    matmul_precision: str = field(
+        default="default",
+        metadata={
+            "help": "Default floating-point precision of internal computations used in TPU matrix multiplications and convolutions. "
+            "This configuration option controls the default precision for JAX operations that take an optional precision argument (e.g. `lax.conv_general_dilated` and `lax.dot`). "
+            "This configuration option does not change the behaviours of such calls with explicit precision arguments; "
+            "it only changes the behaviors of calls with no such argument provided. "
+            "One of `['highest', 'float32', 'high', 'bfloat16_3x', 'default', 'bfloat16', 'fastest', None]`."
+        },
+    )
+    multisteps: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to use Optax MultiSteps for gradient accumulation. If `False` (default) and `gradient_accumulation_steps > 1`, "
+            "a custom gradient accumulation implementation will be employed."
+        },
+    )
+
+
+def to_fp32(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+
+
+def to_bf16(t):
+    return jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16) if x.dtype == jnp.float32 else x, t)
+
+
+class MixedPrecisionTrainState(struct.PyTreeNode):
+    """Train state for use with a single Optax optimizer.
+    Adapted from flax train_state https://github.com/google/flax/blob/main/flax/training/train_state.py
+
+    Synopsis::
+
+        state = TrainState.create(
+            apply_fn=model.apply,
+            params=variables['params'],
+            tx=tx)
+        grad_fn = jax.grad(make_loss_fn(state.apply_fn))
+        for batch in data:
+          grads = grad_fn(state.params, batch)
+          state = state.apply_gradients(grads=grads)
+
+    Args:
+      step: Counter starts at 0 and is incremented by every call to
+        `.apply_gradients()`.
+      apply_fn: Usually set to `model.apply()`. Kept in this dataclass for
+        convenience to have a shorter params list for the `train_step()` function
+        in your training loop.
+      params: The parameters to be updated by `tx` and used by `apply_fn`.
+      tx: An Optax gradient transformation.
+      opt_state: The state for `tx`.
+      dropout_rng: PRNG key for stochastic operations.
+      bf16: Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training.
+    """
+
+    step: int
+    apply_fn: Callable = struct.field(pytree_node=False)
+    get_attention_mask_fn: Callable = struct.field(pytree_node=False)
+    params: core.FrozenDict[str, Any]
+    tx: optax.GradientTransformation = struct.field(pytree_node=False)
+    opt_state: optax.OptState
+    dropout_rng: jnp.ndarray
+    max_grad_norm: Optional[float] = 1.0
+
+    def apply_gradients(self, *, grads, to_dtype, **kwargs):
+        """Updates `step`, `params`, `opt_state` and `**kwargs` in return value.
+
+        Note that internally this function calls `.tx.update()` followed by a call
+        to `optax.apply_updates()` to update `params` and `opt_state`.
+
+        Args:
+          grads: Gradients that have the same pytree structure as `.params`.
+          **kwargs: Additional dataclass attributes that should be `.replace()`-ed.
+
+        Returns:
+          An updated instance of `self` with `step` incremented by one, `params`
+          and `opt_state` updated by applying `grads`, and additional attributes
+          replaced as specified by `kwargs`.
+        """
+
+        # clip gradients by global l2 norm
+        casted_max_grad_norm = to_dtype(self.max_grad_norm)
+        g_norm = linear_algebra.global_norm(grads)
+        g_norm = jnp.maximum(casted_max_grad_norm, g_norm)
+        grads = jax.tree_util.tree_map(lambda t: (t / g_norm) * casted_max_grad_norm, grads)
+
+        # perform update step in fp32 and subsequently downcast optimizer states if mixed precision training
+        # grads and opt_state in bf16 (need to upcast), params in fp32 (leave as is)
+        updates, new_opt_state = self.tx.update(to_fp32(grads), to_fp32(self.opt_state), self.params)
+
+        new_params = optax.apply_updates(self.params, updates)
+        return self.replace(
+            step=self.step + 1,
+            params=new_params,
+            opt_state=to_dtype(new_opt_state),
+            **kwargs,
+        )
+
+    @classmethod
+    def create(cls, *, apply_fn, params, tx, to_dtype, step=0, **kwargs):
+        """Creates a new instance with `step=0` and initialized `opt_state`."""
+        # downcast optimizer state to bf16 if mixed-precision training
+        opt_state = tx.init(to_dtype(params)) if tx is not None else None
+        return cls(
+            step=step,
+            apply_fn=apply_fn,
+            params=params,
+            tx=tx,
+            opt_state=opt_state,
+            **kwargs,
+        )
+
+    def replicate(self):
+        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
+
+
+@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).
+        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
+    input_padding: Union[bool, str] = "longest"
+    label_padding: Union[bool, str] = "max_length"
+    pad_input_to_multiple_of: Optional[int] = None
+    pad_to_multiple_of_label: Optional[int] = None
+    max_input_length: Optional[float] = None
+    max_label_length: Optional[float] = 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
+        input_features = [{"input_values": feature["input_values"]} for feature in features]
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+
+        # reformat list to dict and set to pytorch format
+        batch = self.processor.feature_extractor.pad(
+            input_features,
+            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_label_length,
+            padding=self.label_padding,
+            pad_to_multiple_of=self.pad_to_multiple_of_label,
+            return_tensors="np",
+        )
+
+        labels = labels_batch["input_ids"]
+        labels = np.ma.array(labels, mask=np.not_equal(labels_batch.attention_mask, 1))
+        labels = labels.filled(fill_value=-100)
+
+        batch["labels"] = labels
+
+        return batch
+
+
+def get_grouped_indices(
+    dataset, batch_size: int, rng: Optional[List[int]] = None, mega_batch_mult: Optional[int] = None
+) -> np.array:
+    """
+    Adapted from the `get_length_grouped_indices` function in the PyTorch Trainer utils file (https://github.com/huggingface/transformers/blob/main/src/transformers/trainer_pt_utils.py#L486)
+    Function that returns a list of indices in which each slice of `batch_size` consecutive indices correspond to elements of similar
+    lengths. To do this, the indices are:
+
+    - randomly permuted (if a JAX rng is specified)
+    - grouped in mega-batches of size `mega_batch_mult * batch_size`
+    - sorted by length in each mega-batch
+
+    The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of
+    maximum length placed first, so that an OOM happens sooner rather than later.
+    """
+    lengths = dataset["input_length"]
+
+    # Default for mega_batch_mult: 50 or the number to get 4 megabatches, whichever is smaller.
+    if mega_batch_mult is None:
+        mega_batch_mult = min(len(lengths) // (batch_size * 4), 50)
+        # Just in case, for tiny datasets
+        if mega_batch_mult == 0:
+            mega_batch_mult = 1
+
+    # We need to use JAX for the random permutation as the PRNG key will be set based on the seed outside of the sampler.
+    num_samples = len(lengths)
+    indices = jax.random.permutation(rng, np.arange(num_samples)) if rng is not None else np.arange(num_samples)
+
+    megabatch_size = mega_batch_mult * batch_size
+    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
+    megabatches = [list(sorted(megabatch, key=lambda i: lengths[i], reverse=True)) for megabatch in megabatches]
+
+    # The rest is to get the biggest batch first.
+    # Since each megabatch is sorted by descending length, the longest element is the first
+    megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches]
+    max_idx = np.argmax(megabatch_maximums).item()
+    # Switch to put the longest batch in first position
+    # (note that this is different to the PT grouped sampler in which we only put the longest element in the first position, and not its batch)
+    megabatches[0], megabatches[max_idx] = megabatches[max_idx], megabatches[0]
+
+    megabatches = np.array([i for megabatch in megabatches for i in megabatch])
+
+    return megabatches
+
+
+def generate_batch_splits(samples_idx: np.ndarray, batch_size: int, drop_last=True) -> np.ndarray:
+    """Generate batches of data for a specified batch size from sample indices. If the dataset size is not divisible by
+    the batch size and `drop_last` is `True`, the last incomplete batch is dropped. Else, it is returned."""
+    num_samples = len(samples_idx)
+    if drop_last:
+        samples_to_remove = num_samples % batch_size
+        if samples_to_remove != 0:
+            samples_idx = samples_idx[:-samples_to_remove]
+        sections_split = num_samples // batch_size
+        samples_idx = samples_idx.reshape((sections_split, batch_size))
+    else:
+        sections_split = math.ceil(num_samples / batch_size)
+        samples_idx = np.array_split(samples_idx, sections_split)
+    return samples_idx
+
+
+def write_train_metric(summary_writer, train_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)
+
+
+def write_eval_metric(summary_writer, eval_metrics, step, pred_str=None):
+    for metric_name, value in eval_metrics.items():
+        summary_writer.scalar(f"eval_{metric_name}", value, step)
+
+    if pred_str is not None:
+        # write output actual predictions for debugging
+        summary_writer.text("eval_predictions", "\n".join(pred_str), step)
+
+
+def write_wandb_log(metrics, step, prefix=None):
+    if jax.process_index() == 0:
+        log_metrics = {}
+        for k, v in metrics.items():
+            if "layer" in k:
+                log_metrics[f"{k}/"] = v
+            elif prefix is not None:
+                log_metrics[f"{prefix}/{k}"] = v
+            else:
+                log_metrics[k] = v
+        wandb.log(log_metrics, step)
+
+
+def write_wandb_pred(pred_str, label_str, step, num_log=50, prefix="eval"):
+    if jax.process_index() == 0:
+        # convert str data to a wandb compatible format
+        str_data = [[label_str[i], pred_str[i]] for i in range(len(pred_str))]
+        # we'll log the first 50 predictions for each epoch
+        wandb.log(
+            {
+                f"{prefix}/step_{int(step / 1000)}k": wandb.Table(
+                    columns=["label_str", "pred_str"], data=str_data[:num_log]
+                )
+            },
+            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 ctc_loss(
+    logits,
+    logits_attention_mask,
+    labels,
+    blank_id,
+    loss_reduction="mean",
+    output_emission_dict=False,
+    log_epsilon=-100000.0,
+):
+    """Computes CTC loss.
+    This function performs forward computation over an FSA with `N * 2` states
+    where `N` is the max number of labels. The states are split into two groups:
+    Phi states and emission states. a phi-state accepts repetition of
+    phi (blank)-symbols and transits to emission state when the correct label is
+    observed. An emission state accepts repetition of the label and transits to
+    the next phi states at any time (so called epsilon-transition).
+    Below, `B` denotes the batch size, `T` denotes the time steps in `logits`,
+    and `N` denotes the time steps in `labels`.
+    Args:
+      logits: (B, T, K)-array containing log-probabilities of each class.
+      logitpaddings: (B, T)-array. Padding indicators for `logits`.
+      labels: (B, N)-array containing reference integer labels.
+      labelpaddings: (B, N)-array. Padding indicators for `labels`. Currently,
+        `labels` must be right-padded, i.e. each row of `labelpaddings` must be
+        repetition of zeroes, followed by repetition of ones.
+      blank_id: Id for blank token.
+      loss_reduction: one of "mean", "sum", "default"
+        - "none": no reduction is applied.
+        - "mean": output loss will be divided by target lengths and then the
+          mean over the batch is taken.
+        - "sum": output loss are summed over batch
+      output_emission_dict: whether to output additional information about the emission probs
+    Returns:
+      A pair of `(per_seq_loss, aux)`.
+      per_seq_loss:
+        (B,)-array containing loss values for each sequence in the batch.
+      aux: Dictionary containing interim variables used for computing losses.
+        aux['logalpha_phi']: (T, B, N+1)-array. Log-forward-probabilities of each
+          phi-state corresponding to the n-th label.
+        aux['logalpha_emit']: (T, B, N)-array. Log-forward-probabilities of each
+          emission-state corresponding to the n-th label.
+        aux['logprobs_phi']: (T, B, 1)-array. Probability of the phi-symbol
+          corresponding to each time frame.
+        aux['logprobs_emit']: (T, B, N)-array. Probability of the n-th label
+          corresponding to each time frame.
+    """
+    # label paddings are indicated by -100
+    labelpaddings = labels < 0
+    # logit paddings are the inverse of attention_mask
+    logitpaddings = ~logits_attention_mask
+
+    # Copied from https://github.com/tensorflow/lingvo/blob/master/lingvo/jax/layers/ctc_objectives.py
+    batchsize, unused_maxinputlen, num_classes = logits.shape
+    batchsize_, maxlabellen = labels.shape
+
+    logprobs = jax.nn.log_softmax(logits)
+    labellens = maxlabellen - jnp.sum(labelpaddings, axis=1).astype(jnp.int32)
+
+    # repeat[b, n] == 1.0 when label[b, n] == label[b, n+1].
+    repeat = (labels[:, :-1] == labels[:, 1:]).astype(jnp.float32)
+    repeat = jnp.pad(repeat, ((0, 0), (0, 1)))
+
+    logprobs_phi = logprobs[:, :, blank_id : blank_id + 1]  # [B, T, 1]
+    logprobs_phi = jnp.transpose(logprobs_phi, (1, 0, 2))  # [T, B, 1]
+
+    one_hot = jax.nn.one_hot(labels, num_classes=num_classes)  # [B, N, K]
+    logprobs_emit = jnp.einsum("btk,bnk->btn", logprobs, one_hot)
+    logprobs_emit = jnp.transpose(logprobs_emit, (1, 0, 2))  # [T, B, N]
+
+    logalpha_phi_init = jnp.ones((batchsize, maxlabellen + 1)) * log_epsilon  # [B, N]
+    logalpha_phi_init = logalpha_phi_init.at[:, 0].set(0.0)
+    logalpha_emit_init = jnp.ones((batchsize, maxlabellen)) * log_epsilon  # [B, N]
+
+    def loop_body(prev, x):
+        prev_phi, prev_emit = prev
+        # emit-to-phi epsilon transition, except if the next label is repetition
+        prev_phi_orig = prev_phi
+        prev_phi = prev_phi.at[:, 1:].set(jnp.logaddexp(prev_phi[:, 1:], prev_emit + log_epsilon * repeat))
+
+        logprob_emit, logprob_phi, pad = x
+
+        # phi-to-emit transition
+        next_emit = jnp.logaddexp(prev_phi[:, :-1] + logprob_emit, prev_emit + logprob_emit)
+        # self-loop transition
+        next_phi = prev_phi + logprob_phi
+        # emit-to-phi blank transition only when the next label is repetition
+        next_phi = next_phi.at[:, 1:].set(
+            jnp.logaddexp(next_phi[:, 1:], prev_emit + logprob_phi + log_epsilon * (1.0 - repeat))
+        )
+
+        pad = pad.reshape((batchsize, 1))
+        next_emit = pad * prev_emit + (1.0 - pad) * next_emit
+        next_phi = pad * prev_phi_orig + (1.0 - pad) * next_phi
+
+        return (next_phi, next_emit), (next_phi, next_emit)
+
+    xs = (logprobs_emit, logprobs_phi, logitpaddings.transpose((1, 0)))
+    _, (logalpha_phi, logalpha_emit) = jax.lax.scan(loop_body, (logalpha_phi_init, logalpha_emit_init), xs)
+
+    # last row needs to be updated with the last epsilon transition
+    logalpha_phi_last = logalpha_phi[-1].at[:, 1:].set(jnp.logaddexp(logalpha_phi[-1, :, 1:], logalpha_emit[-1]))
+    logalpha_phi = logalpha_phi.at[-1].set(logalpha_phi_last)
+
+    # extract per_seq_loss
+    one_hot = jax.nn.one_hot(labellens, num_classes=maxlabellen + 1)  # [B, N+1]
+    per_seq_loss = -jnp.einsum("bn,bn->b", logalpha_phi_last, one_hot)
+
+    if loss_reduction == "mean":
+        target_lengths = labelpaddings.shape[-1] - labelpaddings.sum(axis=-1)
+        loss = (per_seq_loss / target_lengths).mean()
+    elif loss_reduction == "sum":
+        loss = per_seq_loss.sum()
+    else:
+        loss = per_seq_loss
+
+    if not output_emission_dict:
+        return loss
+
+    return loss, {
+        "logalpha_phi": logalpha_phi,
+        "logalpha_emit": logalpha_emit,
+        "logprobs_phi": logprobs_phi,
+        "logprobs_emit": logprobs_emit,
+    }
+
+
+def make_dataset(data_args, seed=42):
+    # Pre-processing dataset
+    import re
+
+    def map_nst(entry):
+        text = entry["text"].lower()
+        text = text.replace("(...vær stille under dette opptaket...)", "")
+        text = re.sub('[áàâ]', 'a', text)
+        text = re.sub('[ä]', 'æ', text)
+        text = re.sub('[éèëê]', 'e', text)
+        text = re.sub('[íìïî]', 'i', text)
+        text = re.sub('[óòöô]', 'o', text)
+        text = re.sub('[ö]', 'ø', text)
+        text = re.sub('[ç]', 'c', text)
+        text = re.sub('[úùüû]', 'u', text)
+        # text = re.sub('\\(?=(Punktum|Komma|Utropstegn|Spørsmålstegn))', ' ', text)
+        text = re.sub('\s+', ' ', text)
+        return {"text": text}
+
+    def filter_nst(entry):
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.match(entry["type"], "pIW|CA"):
+            return False  # Spelling out words
+        return True
+
+    def filter_npsc(entry):
+        # False if there are digits in the text
+        if not ((len(entry["text"]) <= len(entry["audio"]["array"]) // 320) and (len(entry["text"].strip()) >= 3)):
+            return False  # Too short
+        if re.search("\d", entry["text"]):
+            return False
+        return True
+
+    def map_npsc(entry):
+        batch = {"text": entry["text"].lower()}
+        batch["text"] = re.sub('[áàâ]', 'a', batch["text"])
+        batch["text"] = re.sub('[ä]', 'æ', batch["text"])
+        batch["text"] = re.sub('[éèëê]', 'e', batch["text"])
+        batch["text"] = re.sub('[íìïî]', 'i', batch["text"])
+        batch["text"] = re.sub('[óòöô]', 'o', batch["text"])
+        batch["text"] = re.sub('[ö]', 'ø', batch["text"])
+        batch["text"] = re.sub('[ç]', 'c', batch["text"])
+        batch["text"] = re.sub('[úùüû]', 'u', batch["text"])
+        batch["text"] = re.sub('\s', ' ', batch["text"])
+        batch["text"] = re.sub('<ee>', 'eee', batch["text"])
+        batch["text"] = re.sub('<qq>', 'qqq', batch["text"])
+        batch["text"] = re.sub('<mm>', 'mmm', batch["text"])
+        batch["text"] = re.sub('<inaudible>', 'xxx', batch["text"])
+        # batch["text"] = re.sub('<inaudible>', '?', batch["text"])
+        if "<" in batch["text"]:
+            raise ValueError(batch["text"])
+        return batch
+
+    nst = datasets.load_dataset("NbAiLab/NST", "no-close")
+    npsc = datasets.load_dataset("NbAiLab/NPSC", "16K_mp3")
+    # TODO NST_hesitate
+
+    split = len(npsc[data_args.train_split_name]) / (len(npsc[data_args.train_split_name]) + len(npsc[data_args.eval_split_name]))  # Use same train/val ratio as NPSC
+    nst_train = nst[data_args.train_split_name].train_test_split(train_size=split, seed=seed)
+    nst[data_args.train_split_name] = nst_train["train"]
+    nst[data_args.eval_split_name] = nst_train["test"]
+
+    nst = nst.filter(filter_nst).map(
+        map_nst,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NST",
+    ).shuffle(seed=seed)
+    npsc = npsc.filter(filter_npsc).map(
+        map_npsc,
+        num_proc=data_args.preprocessing_num_workers,
+        desc="filtering NPSC",
+    ).shuffle(seed=seed)
+
+    npsc_base = npsc.remove_columns([col for col in npsc[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+    nst_base = nst.remove_columns([col for col in nst[data_args.train_split_name].column_names if col not in ["text", "audio"]])
+
+    combined = {}
+    for split in data_args.train_split_name, data_args.eval_split_name, data_args.test_split_name:
+        probs = np.array([len(nst_base[split]), len(npsc_base[split])])  # Weight by number of examples
+        probs = (probs / probs.sum()).tolist()
+        comb = datasets.interleave_datasets([nst_base[split], npsc_base[split]], probabilities=probs, seed=seed)
+        combined[split] = comb
+
+    return datasets.DatasetDict(**combined)
+
+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, FlaxTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # 2. 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()
+
+    # Set up wandb run
+    if jax.process_index() == 0:
+        wandb.init(project=data_args.wandb_project, name=data_args.wandb_name, job_type=data_args.wandb_job_type)
+
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set the default TPU matmul precision and display the number of devices
+    jax.config.update("jax_default_matmul_precision", training_args.matmul_precision)
+    logger.info(f"JAX devices: {jax.device_count()}, matmul precision: {training_args.matmul_precision}")
+
+    # 4. Load dataset
+
+    set_seed(training_args.seed)
+    raw_datasets = make_dataset(data_args, seed=training_args.seed)
+   
+    # raw_datasets = DatasetDict()
+
+    # if training_args.do_train:
+    #     raw_datasets[data_args.train_split_name] = 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[data_args.eval_split_name] = 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 training_args.do_predict:
+    #     test_split = data_args.test_split_name.split("+")
+    #     for split in test_split:
+    #         raw_datasets[split] = load_dataset(
+    #             data_args.dataset_name,
+    #             data_args.dataset_config_name,
+    #             split=split,
+    #             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 and not training_args.do_predict:
+        raise ValueError(
+            "Cannot not train, not do evaluation and not do prediction. At least one of "
+            "training, evaluation or prediction has to be done."
+        )
+
+    # if not training, there is no need to run multiple epochs
+    if not training_args.do_train:
+        training_args.num_train_epochs = 1
+
+    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
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    config = Wav2Vec2Config.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,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    # update config according to training args, model args, and tokenizer attributes
+    config.update(
+        {
+            "feat_proj_dropout": model_args.feat_proj_dropout,
+            "attention_dropout": model_args.attention_dropout,
+            "hidden_dropout": model_args.hidden_dropout,
+            "final_dropout": model_args.final_dropout,
+            "mask_time_prob": model_args.mask_time_prob,
+            "mask_time_length": model_args.mask_time_length,
+            "mask_feature_prob": model_args.mask_feature_prob,
+            "mask_feature_length": model_args.mask_feature_length,
+            "gradient_checkpointing": training_args.gradient_checkpointing,
+            "layerdrop": model_args.layerdrop,
+            "ctc_loss_reduction": model_args.ctc_loss_reduction,
+            "ctc_zero_infinity": model_args.ctc_zero_infinity,
+            "pad_token_id": tokenizer.pad_token_id,
+            "vocab_size": tokenizer.vocab_size,  # len(tokenizer),
+            "activation_dropout": model_args.activation_dropout,
+        }
+    )
+
+    if tokenizer.do_lower_case and data_args.dataset_name != "librispeech_asr":
+        raise ValueError(
+            "Setting the tokenizer attribute `do_lower_case` to `True` converts all input strings to "
+            "uppercase prior to tokenization. This should only be done when the tokenizer is built on an uppercased corpus,"
+            "i.e. for the dataset `librispeech_asr` only. If your dataset is not `librispeech_asr`, the tokenizer is mostly likely "
+            "built on an lowercased corpus. In this case, set `tokenizer.do_lower_case` to ``False`."
+        )
+
+    if training_args.precision == "full_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = True
+    elif training_args.precision == "half_mixed":
+        dtype = jnp.bfloat16
+        training_args.mixed_precision = False
+    else:
+        dtype = jnp.float32
+        training_args.mixed_precision = False
+
+    try:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+        )
+    except:
+        model = FlaxWav2Vec2ForCTC.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            dtype=dtype,
+            cache_dir=model_args.cache_dir,
+            revision=model_args.model_revision,
+            use_auth_token=True if model_args.use_auth_token else None,
+            from_pt=True,
+        )
+
+    # 6. Resample speech dataset ALWAYS
+    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_target_length = data_args.max_label_length
+    min_target_length = data_args.min_label_length
+    pad_input_to_multiple_of = data_args.pad_input_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
+    dataset_name = data_args.dataset_name
+    chars_to_ignore = ', ? . ! - ; : " “ % ‘ ” ?'.split(" ")
+    chars_to_ignore_regex = f'[{"".join(chars_to_ignore)}]'
+    # gigaspeech_punctuation = {" <comma>": ",", " <period>": ".", " <questionmark>": "?", " <exclamationpoint>": "!"}
+    # gigaspeech_disfluencies = ["<other>", "<sil>"]
+    # swb_disfluencies = ["[noise]", "[laughter]", "[silence]", "<a_aside>", "<b_aside>", "<e_aside>", "[laughter-",
+    #                 "[vocalized-noise]", "_1"]
+    # swb_punctuations = ["{", "}", "[", "]-", "]"]
+    # earnings_disfluencies = ["<crosstalk>", "<affirmative>", "<inaudible>", "inaudible", "<laugh>", "<unk>"]
+    ignore_segments = ["ignore_time_segment_in_scoring", "<noise>", "<music>", "[noise]", "[laughter]", "[silence]",
+                       "[vocalized-noise]", "<crosstalk>", "<affirmative>", "<inaudible>", "<laugh>", "<other>", "<sil>", ""]
+
+    if training_args.do_train and data_args.max_train_samples is not None:
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].select(range(data_args.max_train_samples))
+
+    if training_args.do_eval and data_args.max_eval_samples is not None:
+        raw_datasets[data_args.eval_split_name] = raw_datasets[data_args.eval_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_predict and data_args.max_test_samples is not None:
+        raw_datasets[data_args.test_split_name] = raw_datasets[data_args.test_split_name].select(range(data_args.max_eval_samples))
+
+    if training_args.do_train and data_args.remove_punctuation:
+
+        def remove_punctuation(batch):
+            batch[text_column_name] = (
+                re.sub(chars_to_ignore_regex, "", batch[text_column_name]).replace("'", "").replace('"', "")
+            )
+
+        raw_datasets[data_args.train_split_name] = raw_datasets[data_args.train_split_name].map(
+            remove_punctuation,
+            num_proc=data_args.preprocessing_num_workers,
+            desc="removing punctuation from train split",
+        )
+
+    # filter data where the targets are ignored in scoring
+    def is_target_labels(input_str):
+        return input_str.lower() not in ignore_segments
+
+    raw_datasets = raw_datasets.filter(
+            is_target_labels,
+            num_proc=num_workers,
+            input_columns=[text_column_name],
+            desc="filtering data where the targets are ignored in scoring",
+        )
+
+    def prepare_dataset(batch):
+        # process audio
+        try:
+            sample = batch[audio_column_name]
+        except ValueError:
+            sample = {"array": np.array([0.]), "sampling_rate": feature_extractor.sampling_rate}
+        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
+        # process audio length
+        batch[model_input_name] = inputs.input_values[0]
+        batch["input_length"] = len(batch["input_values"])
+
+        # process targets
+        input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
+
+        # if dataset_name == "google/xtreme_s":
+        #     # Finally, we tokenize the processed text
+        #     batch["labels"] = tokenizer(input_str).input_ids
+        #     batch["labels_length"] = len(batch["labels"])
+        #     return batch
+
+        # # Common Voice 9
+        # if input_str.startswith('"') and input_str.endswith('"'):
+        #     # we can remove trailing quotation marks as they do not affect the transcription
+        #     input_str = input_str[1:-1]
+        # # normalize quotation marks
+        # input_str = re.sub(r'["“”]', '"', input_str)
+        # # normalize apostrophes
+        # input_str = re.sub(r"[’']", "'", input_str)
+        # # normalize hyphens
+        # input_str = re.sub(r"[—–]", "-", input_str)
+        # # replace double quotation marks with single
+        # input_str = input_str.replace('""', '"')
+        # if dataset_name == "mozilla-foundation/common_voice_9_0" and len(input_str):
+        #     # for CV9, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # TEDLIUM-3
+        # # delete the <unk> token from the text and replace spaced apostrophes with un-spaced
+        # input_str = input_str.replace("<unk>", "").replace(" '", "'")
+
+        # # GigaSpeech
+        # for disfluency in gigaspeech_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # convert spelled out punctuation to symbolic form
+        # for punctuation, replacement in gigaspeech_punctuation.items():
+        #     input_str = input_str.replace(punctuation, replacement)
+        # if dataset_name == "speechcolab/gigaspeech" and len(input_str):
+        #     # for GS, we'll normalize the text to always finish with punctuation
+        #     if input_str[-1] not in [".", "?", "!"]:
+        #         input_str = input_str + "."
+
+        # # SWB
+        # for disfluency in swb_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # remove parenthesised text (test data only)
+        # input_str = re.sub("[\(].*?[\)]", "", input_str)
+        # for punctuation in swb_punctuations:
+        #     input_str = input_str.replace(punctuation, "")
+        # # replace anomalous words with their correct transcriptions
+        # split_str = input_str.split("/")
+        # if len(split_str) > 1:
+        #     input_str = " ".join(
+        #         [" ".join([" ".join(i.split(" ")[:-1]) for i in split_str])] + [split_str[-1].split(" ")[-1]])
+
+        # # Earnings 22
+        # for disfluency in earnings_disfluencies:
+        #     input_str = input_str.replace(disfluency, "")
+        # # replace mal-formatted ellipsis
+        # input_str = input_str.replace("…", ".")
+
+        # JIWER compliance
+        # remove multiple spaces
+        input_str = re.sub(r"\s\s+", " ", input_str)
+        # strip trailing spaces
+        input_str = input_str.strip()
+
+        # Finally, we tokenize the processed text
+        batch["labels"] = tokenizer(input_str).input_ids
+        batch["labels_length"] = len(batch["labels"])
+        return batch
+
+    vectorized_datasets = raw_datasets.map(
+        prepare_dataset,
+        remove_columns=next(iter(raw_datasets.values())).column_names,
+        num_proc=num_workers,
+        desc="preprocess dataset",
+    )
+
+    # filter data with inputs shorter than min_input_length or longer than max_input_length
+    def is_audio_in_length_range(length):
+        return length > min_input_length and length < max_input_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_audio_in_length_range,
+        num_proc=num_workers,
+        input_columns=["input_length"],
+    )
+
+    # filter data with targets shorter than min_target_length or longer than max_target_length
+    def is_labels_in_length_range(length):
+        return length > min_target_length  # and length < max_target_length
+
+    vectorized_datasets = vectorized_datasets.filter(
+        is_labels_in_length_range,
+        num_proc=num_workers,
+        input_columns=["labels_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 Metrics
+    wer_metric = load_metric("wer")
+    cer_metric = load_metric("cer")
+
+    def compute_metrics(pred_ids: List[List[int]], label_ids: List[List[int]]):
+        padded_ids = np.where(np.asarray(label_ids) == -100, tokenizer.pad_token_id, np.asarray(label_ids))
+
+        pred_str = tokenizer.batch_decode(pred_ids)
+        # we do not want to group tokens when computing the metrics
+        label_str = tokenizer.batch_decode(padded_ids, group_tokens=False)
+
+        wer = wer_metric.compute(predictions=pred_str, references=label_str)
+        cer = cer_metric.compute(predictions=pred_str, references=label_str)
+
+        return {"wer": wer, "cer": cer}, pred_str, label_str
+
+    # 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,
+        input_padding="longest",
+        pad_input_to_multiple_of=pad_input_to_multiple_of,
+        max_label_length=data_args.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."
+        )
+
+    # 10. Handle the repository creation
+    if training_args.push_to_hub:
+        with open(os.path.join(training_args.output_dir, ".gitattributes"), "r+") as f:
+            git_lfs_extensions = f.read()
+            if "*.wandb" not in git_lfs_extensions:
+                f.write("*.wandb filter=lfs diff=lfs merge=lfs -text")
+        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
+        repo = Repository(training_args.output_dir, clone_from=repo_name)
+
+    # 11. Initialize our training
+    rng = jax.random.PRNGKey(training_args.seed)
+    rng, dropout_rng = jax.random.split(rng)
+
+    # Store some constants
+    max_steps = int(training_args.max_steps)
+    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
+    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
+    batch_size_per_update = train_batch_size * gradient_accumulation_steps
+    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
+    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
+    to_dtype = to_bf16 if training_args.mixed_precision else to_fp32
+
+    if training_args.do_train:
+        num_train_samples = len(vectorized_datasets[data_args.train_split_name])
+        steps_per_epoch = num_train_samples // batch_size_per_update
+        if max_steps > 0:
+            num_epochs = -(training_args.max_steps // -steps_per_epoch)
+            total_train_steps = max_steps
+        else:
+            num_epochs = int(training_args.num_train_epochs)
+            total_train_steps = steps_per_epoch * num_epochs
+
+        # Create learning rate schedule
+        # Create learning rate schedule
+        linear_decay_lr_schedule_fn = create_learning_rate_fn(
+            total_train_steps,
+            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.
+        # Note that this mask is specifically adapted for FlaxWav2Vec2 and FlaxBart.
+        # For FlaxT5, one should correct the layer norm parameter naming
+        # accordingly - see `run_t5_mlm_flax.py` e.g.
+        def decay_mask_fn(params):
+            flat_params = traverse_util.flatten_dict(params)
+            layer_norm_params = [
+                (name, "scale")
+                for name in ["layer_norm", "self_attn_layer_norm", "layernorm_embedding", "final_layer_norm"]
+            ]
+            flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_params) for path in flat_params}
+            return traverse_util.unflatten_dict(flat_mask)
+
+        if training_args.adafactor:
+            # Create Adafactor optimizer
+            optim = optax.adafactor(
+                learning_rate=linear_decay_lr_schedule_fn,
+                dtype_momentum=jnp.bfloat16 if training_args.mixed_precision else jnp.float32,
+                weight_decay_rate=training_args.weight_decay,
+                weight_decay_mask=decay_mask_fn,
+            )
+        else:
+            # Create AdamW optimizer
+            optim = 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,
+            )
+
+        # Optax MultiSteps for gradient accumulation. We'll only call this optimizer transformation if gradient accumulation is required (i.e. gradient accumulation steps > 1)
+        if training_args.multisteps and gradient_accumulation_steps > 1:
+            optim = optax.MultiSteps(optim, gradient_accumulation_steps, use_grad_mean=False)
+    else:
+        num_epochs = 0
+        total_train_steps = 0
+        num_train_samples = 0
+        optim = None
+
+    # Setup train state
+    state = MixedPrecisionTrainState.create(
+        step=data_args.skip_steps,
+        apply_fn=model.__call__,
+        get_attention_mask_fn=model._get_feature_vector_attention_mask,
+        params=model.params,
+        tx=optim,
+        to_dtype=to_dtype,
+        dropout_rng=dropout_rng,
+        max_grad_norm=training_args.max_grad_norm,
+    )
+
+    # Replicate the train state on each device
+    state = state.replicate()
+    blank_id = model.config.pad_token_id
+
+    # Define gradient update step fn
+    def train_step(state, batch):
+        # only one single rng per grad step, with or without accumulation, as the graph should be identical over one effective training batch
+        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
+
+        def compute_loss(params, minibatch):
+            labels = minibatch.pop("labels")
+            logits = state.apply_fn(
+                **minibatch,
+                params=params,
+                dropout_rng=dropout_rng,
+                freeze_feature_encoder=model_args.freeze_feature_encoder,
+                train=True,
+            )[0]
+            logits_mask = state.get_attention_mask_fn(logits.shape[1], batch["attention_mask"])
+            loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+            return loss
+
+        grad_fn = jax.value_and_grad(compute_loss)
+
+        if gradient_accumulation_steps == 1 or training_args.multisteps:
+            loss, grad = grad_fn(to_dtype(state.params), batch)
+
+        # Custom gradient accumulation
+        else:
+            # add a first dimension over gradient_accumulation_steps for minibatch slices
+            batch = jax.tree_util.tree_map(
+                lambda x: x.reshape(
+                    gradient_accumulation_steps, training_args.per_device_train_batch_size, *x.shape[1::]
+                ),
+                batch,
+            )
+
+            def accum_minibatch_step(accum_grad, minibatch):
+                # compute loss, num labels and grad over minibatch and accumulate
+                loss, grad = grad_fn(to_dtype(state.params), minibatch)
+                return jax.tree_util.tree_map(jnp.add, accum_grad, grad), loss
+
+            # create an initial state for accumulating losses, num labels and gradients
+            init_grad = jax.tree_util.tree_map(jnp.zeros_like, to_dtype(state.params))
+            # loop accum minibatch step over the number of gradient accumulation steps
+            grad, loss = jax.lax.scan(accum_minibatch_step, init_grad, batch)
+
+        # update state
+        new_state = state.apply_gradients(
+            grads=grad,
+            dropout_rng=new_dropout_rng,
+            to_dtype=to_dtype,
+        )
+
+        # compute gradient norms over all layers and globally for detailed monitoring
+        layer_grad_norm = jax.tree_util.tree_map(jnp.linalg.norm, grad)
+        logs = {
+            "layer_grad_norm": layer_grad_norm,
+            "grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm)),
+        }
+
+        # compute parameter norms over all layers and globally for detailed monitoring
+        layer_param_norm = jax.tree_util.tree_map(jnp.linalg.norm, new_state.params)
+        logs["layer_param_norm"] = layer_param_norm
+        logs["param_norm"] = jnp.linalg.norm(jax.tree_util.tree_leaves(layer_param_norm))
+
+        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+        metrics.update(logs)
+
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+
+        return new_state, metrics
+
+    # Define eval fn
+    def eval_step(params, batch):
+        labels = batch.pop("labels")
+        logits = model(**batch, params=params, train=False)[0]
+
+        logits_mask = model._get_feature_vector_attention_mask(logits.shape[1], batch["attention_mask"])
+        loss = ctc_loss(logits, logits_mask, labels, blank_id, loss_reduction="mean")
+
+        pred_ids = jnp.argmax(logits, axis=-1)
+
+        # summarize metrics
+        metrics = {"loss": loss}
+        metrics = jax.lax.pmean(metrics, axis_name="batch")
+        # metrics = to_fp32(metrics)
+        return metrics, pred_ids
+
+    # Create parallel version of the train and eval step
+    if training_args.do_train:
+        p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
+
+    if training_args.do_eval:
+        p_eval_step = jax.pmap(eval_step, "batch")
+
+    def run_evaluation(step):
+        if training_args.do_eval:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the eval dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[data_args.eval_split_name], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
+                samples = [vectorized_datasets[data_args.eval_split_name][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, step, prefix="eval")
+            write_wandb_pred(pred_str, label_str, step)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, step, pred_str=pred_str)
+
+    def save_checkpoint(step):
+        # save 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"{wandb.run.id}: saving weights and logs of step {int(step / 1000)}k", blocking=False)
+
+    skip_epochs = data_args.skip_steps // (num_train_samples // batch_size_per_update)
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {num_train_samples}")
+    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"  Num gradient accumulation steps = {gradient_accumulation_steps}")
+    logger.info(f"  Total train batch size (w. parallel & distributed) = {batch_size_per_update}")
+    logger.info(f"  Total optimization steps = {total_train_steps}")
+    logger.info(f"  Gradient checkpointing: {config.gradient_checkpointing}")
+    logger.info(f"  Use scan: {config.use_scan}")
+    logger.info(f"  Fuse matmuls: {config.fuse_matmuls}")
+    logger.info(f"  Skipping: {data_args.skip_steps} steps ({skip_epochs} epochs)")
+
+    train_time = cur_step = 0
+    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
+    for epoch in epochs:
+        if training_args.do_train:
+            # ======================== Training ================================
+            train_start = time.time()
+
+            if epoch < skip_epochs:
+                logger.info(f"Skipping epoch {epoch + 1}")
+                # Create sampling rng
+                rng, input_rng = jax.random.split(rng)
+                continue
+
+            # Generate an epoch by randomly shuffling sampling indices from the train dataset and grouping by length
+            train_samples_idx = get_grouped_indices(vectorized_datasets[data_args.train_split_name], batch_size_per_update, input_rng)
+            train_batch_idx = generate_batch_splits(train_samples_idx, batch_size_per_update)
+
+            if data_args.skip_steps > cur_step:
+                logger.info(f"Skipping {data_args.skip_steps - (epoch * (num_train_samples // batch_size_per_update))} steps...")
+            # Gather the indices for creating the batch and do a training step
+            for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1), 1):
+                cur_step = epoch * (num_train_samples // batch_size_per_update) + step
+                if cur_step <= data_args.skip_steps:
+                    continue
+
+                samples = [vectorized_datasets[data_args.train_split_name][int(idx)] for idx in batch_idx]              
+                batch = data_collator(samples)
+                batch = shard(batch.data)
+                try:
+                    state, train_metric = p_train_step(state, batch)
+                except TypeError as e:
+                    logger.warning("Encountered following error: \n", e)
+
+
+                if cur_step % training_args.logging_steps == 0:
+                    # Save metrics
+                    train_metric = unreplicate(train_metric)
+                    train_time += time.time() - train_start
+                    # need to upcast all device arrays to fp32 for wandb logging (jnp.bfloat16 not supported) -> do this here OR in train_step
+                    write_wandb_log(to_fp32(train_metric), cur_step, prefix=data_args.train_split_name)
+                    # we won't log to tensorboard for now (it is fiddly logging param and grad norms on a layer-by-layer basis)
+                    # if has_tensorboard and jax.process_index() == 0:
+                    # write_train_metric(summary_writer, train_metrics, train_time, cur_step)
+
+                    epochs.write(
+                        f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']}, Gradient Norm: {train_metric['grad_norm']})"
+                    )
+                    p_train_step.clear_cache()
+
+                if cur_step % total_train_steps == 0:
+                    break
+
+                if training_args.eval_steps and cur_step % training_args.eval_steps == 0:
+                    run_evaluation(cur_step)
+
+                if cur_step % training_args.save_steps == 0:
+                    save_checkpoint(cur_step)
+
+            if training_args.eval_steps == 0 and (epoch + 1) != num_epochs:
+                # run evaluation at the end of the epoch if eval steps are not specified
+                run_evaluation(cur_step)
+                save_checkpoint(cur_step)
+
+    if training_args.do_train:
+        save_checkpoint(cur_step)
+
+    cur_step = max_steps if max_steps > 0 else cur_step  # set step to max steps so that eval happens in alignment with training
+
+    if training_args.do_eval:
+        run_evaluation(cur_step)
+
+    # TODO: collapse 'do_predict' into the run_evaluation function
+    if training_args.do_predict:
+        for split in [data_args.test_split_name]:
+            # ======================== Evaluating ==============================
+            eval_metrics = []
+            eval_preds = []
+            eval_labels = []
+
+            # Generate eval set by sequentially sampling indices from the test dataset and grouping by length
+            eval_samples_idx = get_grouped_indices(vectorized_datasets[split], eval_batch_size)
+            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size, drop_last=False)
+
+            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc=f"Predicting {split}...", position=2)):
+                samples = [vectorized_datasets[split][int(idx)] for idx in batch_idx]
+                batch = data_collator(samples)
+                labels = batch["labels"]
+
+                metrics, pred_ids = pad_shard_unpad(p_eval_step)(state.params, batch.data, min_device_batch=per_device_eval_batch_size)
+                eval_preds.extend(jax.device_get(pred_ids.reshape(-1, pred_ids.shape[-1])))
+                eval_metrics.append(metrics)
+
+                eval_labels.extend(labels)
+
+            # normalize eval metrics
+            eval_metrics = get_metrics(eval_metrics)
+            eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
+            eval_metrics = to_fp32(eval_metrics)
+
+            # always run compute metrics
+            error_rate_metric, pred_str, label_str = compute_metrics(eval_preds, eval_labels)
+            eval_metrics.update(error_rate_metric)
+            error_rate_desc = " ".join([f"Eval {key}: {value} |" for key, value in error_rate_metric.items()])
+
+            # Print metrics and update progress bar
+            desc = f"Step... ({cur_step}/{total_train_steps} | Eval Loss: {eval_metrics['loss']} | {error_rate_desc})"
+            epochs.write(desc)
+            epochs.desc = desc
+
+            # Save metrics
+            write_wandb_log(eval_metrics, cur_step, prefix=split)
+            write_wandb_pred(pred_str, label_str, cur_step, prefix=split)
+            # if has_tensorboard and jax.process_index() == 0:
+            # write_eval_metric(summary_writer, eval_metrics, cur_step, pred_str=pred_str)
+
+
+if __name__ == "__main__":
+    main()

wandb/run-20220810_151736-2jo5la5b/files/config.yaml

@@ -0,0 +1,33 @@
+wandb_version: 1
+
+_wandb:
+  desc: null
+  value:
+    cli_version: 0.12.9
+    code_path: code/run_flax_speech_recognition_ctc.py
+    framework: huggingface
+    huggingface_version: 4.21.0
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    python_version: 3.8.10
+    start_time: 1660144656
+    t:
+      1:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      2:
+      - 1
+      - 2
+      - 3
+      - 11
+      - 12
+      3:
+      - 13
+      4: 3.8.10
+      5: 0.12.9
+      6: 4.21.0
+      8:
+      - 5

wandb/run-20220810_151736-2jo5la5b/files/diff.patch

@@ -0,0 +1,144 @@
+diff --git a/run.recover.sh b/run.recover.sh
+index 77ad3fd..632a336 100755
+--- a/run.recover.sh
++++ b/run.recover.sh
+@@ -11,9 +11,8 @@ WANDB_ENTITY=NbAiLab WANDB_PROJECT=wav2vec2 python run_flax_speech_recognition_c
+         --per_device_train_batch_size="2" \
+         --per_device_eval_batch_size="2" \
+         --gradient_accumulation_steps="1" \
+-        --precision="full_mixed" \
++        --precision="half_mixed" \
+         --matmul_precision="bfloat16" \
+-        --multisteps \
+         --learning_rate="6.394633237505332e-05" \
+         --skip_steps="275000" \
+         --warmup_steps="2000" \
+diff --git a/run_flax_speech_recognition_ctc.py b/run_flax_speech_recognition_ctc.py
+index a330879..4d0f5fc 100644
+--- a/run_flax_speech_recognition_ctc.py
++++ b/run_flax_speech_recognition_ctc.py
+@@ -415,12 +415,12 @@ class MixedPrecisionTrainState(struct.PyTreeNode):
+         )
+ 
+     @classmethod
+-    def create(cls, *, apply_fn, params, tx, to_dtype, **kwargs):
++    def create(cls, *, apply_fn, params, tx, to_dtype, step=0, **kwargs):
+         """Creates a new instance with `step=0` and initialized `opt_state`."""
+         # downcast optimizer state to bf16 if mixed-precision training
+         opt_state = tx.init(to_dtype(params)) if tx is not None else None
+         return cls(
+-            step=0,
++            step=step,
+             apply_fn=apply_fn,
+             params=params,
+             tx=tx,
+@@ -1339,6 +1339,7 @@ def main():
+ 
+     # Setup train state
+     state = MixedPrecisionTrainState.create(
++        step=data_args.skip_steps,
+         apply_fn=model.__call__,
+         get_attention_mask_fn=model._get_feature_vector_attention_mask,
+         params=model.params,
+@@ -1520,11 +1521,10 @@ def main():
+ 
+             if epoch < skip_epochs:
+                 logger.info(f"Skipping epoch {epoch + 1}")
++                # Create sampling rng
++                rng, input_rng = jax.random.split(rng)
+                 continue
+ 
+-            # Create sampling rng
+-            rng, input_rng = jax.random.split(rng)
+-
+             # Generate an epoch by randomly shuffling sampling indices from the train dataset and grouping by length
+             train_samples_idx = get_grouped_indices(vectorized_datasets[data_args.train_split_name], batch_size_per_update, input_rng)
+             train_batch_idx = generate_batch_splits(train_samples_idx, batch_size_per_update)
+@@ -1559,6 +1559,7 @@ def main():
+                     epochs.write(
+                         f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']}, Gradient Norm: {train_metric['grad_norm']})"
+                     )
++                    p_train_step.clear_cache()
+ 
+                 if cur_step % total_train_steps == 0:
+                     break
+diff --git a/special_tokens_map.json b/special_tokens_map.json
+index 218961f..3c0d148 100644
+--- a/special_tokens_map.json
++++ b/special_tokens_map.json
+@@ -399,6 +399,48 @@
+       "rstrip": false,
+       "single_word": false
+     },
++    {
++      "content": "</s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "<s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "</s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "<s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "</s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
++    {
++      "content": "<s>",
++      "lstrip": false,
++      "normalized": true,
++      "rstrip": false,
++      "single_word": false
++    },
+     {
+       "content": "</s>",
+       "lstrip": false,
+diff --git a/wandb/debug-internal.log b/wandb/debug-internal.log
+index 23926ef..90a074d 120000
+--- a/wandb/debug-internal.log
++++ b/wandb/debug-internal.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug-internal.log
+
++run-20220810_151736-2jo5la5b/logs/debug-internal.log
+
+diff --git a/wandb/debug.log b/wandb/debug.log
+index 279853d..de899a6 120000
+--- a/wandb/debug.log
++++ b/wandb/debug.log
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4/logs/debug.log
+
++run-20220810_151736-2jo5la5b/logs/debug.log
+
+diff --git a/wandb/latest-run b/wandb/latest-run
+index f069a7a..0dfb7e0 120000
+--- a/wandb/latest-run
++++ b/wandb/latest-run
+@@ -1 +1 @@
+-run-20220805_230151-2y71vcu4
+
++run-20220810_151736-2jo5la5b
+

wandb/run-20220810_151736-2jo5la5b/files/output.log

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

wandb/run-20220810_151736-2jo5la5b/files/requirements.txt

@@ -0,0 +1,158 @@
+absl-py==1.0.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+appdirs==1.4.4
+astunparse==1.6.3
+async-timeout==4.0.2
+attrs==21.4.0
+audioread==2.1.9
+backcall==0.2.0
+cachetools==4.2.4
+certifi==2021.10.8
+cffi==1.15.1
+charset-normalizer==2.0.10
+chex==0.1.3
+click==8.0.3
+cloud-tpu-client==0.10
+cloud-tpu-profiler==2.4.0
+clu==0.0.6
+colorama==0.4.5
+commonmark==0.9.1
+configparser==5.2.0
+contextlib2==21.6.0
+cycler==0.11.0
+datasets==2.4.0
+decorator==5.1.0
+dill==0.3.4
+dm-tree==0.1.6
+docker-pycreds==0.4.0
+etils==0.6.0
+exceptiongroup==1.0.0rc8
+filelock==3.4.2
+flatbuffers==2.0
+flax==0.5.3
+fonttools==4.28.5
+frozenlist==1.2.0
+fsspec==2021.11.1
+future==0.18.2
+gast==0.4.0
+gitdb==4.0.9
+gitpython==3.1.26
+google-api-core==1.31.5
+google-api-python-client==1.8.0
+google-auth-httplib2==0.1.0
+google-auth-oauthlib==0.4.6
+google-auth==2.3.3
+google-pasta==0.2.0
+googleapis-common-protos==1.54.0
+grpcio==1.43.0
+h5py==3.6.0
+httplib2==0.20.2
+huggingface-hub==0.2.1
+hypothesis==6.53.0
+idna==3.3
+importlib-metadata==4.10.0
+importlib-resources==5.4.0
+ipython==7.31.0
+jax==0.3.15
+jaxlib==0.3.15
+jedi==0.18.1
+jiwer==2.3.0
+joblib==1.1.0
+keras-preprocessing==1.1.2
+keras==2.7.0
+kiwisolver==1.3.2
+libclang==12.0.0
+librosa==0.9.2
+libtpu-nightly==0.1.dev20220722
+llvmlite==0.39.0
+markdown==3.3.6
+matplotlib-inline==0.1.3
+matplotlib==3.5.1
+ml-collections==0.1.0
+msgpack==1.0.3
+multidict==5.2.0
+multiprocess==0.70.12.2
+numba==0.56.0
+numpy==1.22.0
+oauth2client==4.1.3
+oauthlib==3.1.1
+opt-einsum==3.3.0
+optax==0.1.3
+packaging==21.3
+pandas==1.3.5
+parso==0.8.3
+pathtools==0.1.2
+pexpect==4.8.0
+pickleshare==0.7.5
+pillow==9.0.0
+pip==22.2.2
+pkg-resources==0.0.0
+pooch==1.6.0
+promise==2.3
+prompt-toolkit==3.0.24
+protobuf==3.19.1
+psutil==5.9.0
+ptyprocess==0.7.0
+pyarrow==6.0.1
+pyasn1-modules==0.2.8
+pyasn1==0.4.8
+pycparser==2.21
+pyctcdecode==0.4.0
+pygments==2.11.1
+pygtrie==2.5.0
+pyparsing==3.0.6
+python-dateutil==2.8.2
+python-levenshtein==0.12.2
+pytz==2021.3
+pyyaml==6.0
+regex==2021.11.10
+requests-oauthlib==1.3.0
+requests==2.27.0
+resampy==0.3.1
+responses==0.18.0
+rich==11.2.0
+rsa==4.8
+sacremoses==0.0.46
+scikit-learn==1.1.1
+scipy==1.7.3
+sentry-sdk==1.5.2
+setuptools==44.0.0
+shortuuid==1.0.8
+six==1.16.0
+smmap==5.0.0
+sortedcontainers==2.4.0
+soundfile==0.10.3.post1
+sox==1.4.1
+subprocess32==3.5.4
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.0
+tensorboard==2.7.0
+tensorflow-cpu==2.7.0
+tensorflow-datasets==4.4.0
+tensorflow-estimator==2.7.0
+tensorflow-io-gcs-filesystem==0.23.1
+tensorflow-metadata==1.5.0
+tensorflow==2.7.0
+tensorstore==0.1.21
+termcolor==1.1.0
+threadpoolctl==3.1.0
+tokenizers==0.11.2
+toolz==0.11.2
+torch==1.12.0
+torchaudio==0.12.0+cpu
+tqdm==4.62.3
+traitlets==5.1.1
+transformers==4.21.0
+typing-extensions==4.3.0
+uritemplate==3.0.1
+urllib3==1.26.7
+wandb==0.12.9
+wcwidth==0.2.5
+werkzeug==2.0.2
+wheel==0.37.1
+wrapt==1.13.3
+xxhash==2.0.2
+yarl==1.7.2
+yaspin==2.1.0
+zipp==3.7.0