test

run_flax_speech_recognition_seq2seq_streaming_v3.py

@@ -35,7 +35,7 @@ import jax.numpy as jnp
 import numpy as np
 import optax
 import torch
-from datasets import Dataset,DatasetDict, IterableDatasetDict, interleave_datasets, load_dataset
+from datasets import Dataset, DatasetDict, IterableDatasetDict, interleave_datasets, load_dataset
 from torch.utils.data import IterableDataset
 from flax import jax_utils, traverse_util
 from flax.jax_utils import pad_shard_unpad, unreplicate
@@ -66,7 +66,8 @@ from transformers.utils.versions import require_version
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
 check_min_version("4.27.0.dev0")
 
-require_version("datasets>=1.18.2", "To fix: pip install -r examples/flax/speech-recogintion/requirements.txt")
+require_version("datasets>=1.18.2",
+                "To fix: pip install -r examples/flax/speech-recogintion/requirements.txt")
 
 logger = logging.getLogger(__name__)
 
@@ -78,7 +79,8 @@ class ModelArguments:
     """
 
     model_name_or_path: str = field(
-        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+        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"}
@@ -91,15 +93,18 @@ class ModelArguments:
     )
     cache_dir: Optional[str] = field(
         default=None,
-        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+        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."},
+        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)."},
+        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,
@@ -142,7 +147,8 @@ class DataTrainingArguments:
     )
     text_column: Optional[str] = field(
         default=None,
-        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
+        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"}
@@ -170,23 +176,28 @@ class DataTrainingArguments:
     )
     audio_column_name: str = field(
         default="audio",
-        metadata={"help": "The name of the dataset column containing the audio data. Defaults to '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'"},
+        metadata={
+            "help": "The name of the dataset column containing the text data. Defaults to 'text'"},
     )
     max_duration_in_seconds: float = field(
         default=30.0,
-        metadata={"help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"},
+        metadata={
+            "help": "Filter audio files that are longer than `max_duration_in_seconds` seconds"},
     )
     min_duration_in_seconds: float = field(
         default=0.0,
-        metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"},
+        metadata={
+            "help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"},
     )
     max_label_length: float = field(
         default=128,
-        metadata={"help": "Truncate transcriptions that are longer `max_eval_length` tokens."},
+        metadata={
+            "help": "Truncate transcriptions that are longer `max_eval_length` tokens."},
     )
     pad_input_to_multiple_of: Optional[int] = field(
         default=None,
@@ -229,11 +240,13 @@ class DataTrainingArguments:
     )
     do_remove_punctuation: bool = field(
         default=False,
-        metadata={"help": "Whether the target text should be striped of punctuation."},
+        metadata={
+            "help": "Whether the target text should be striped of punctuation."},
     )
     do_normalize_eval: bool = field(
         default=True,
-        metadata={"help": "Whether to normalise the references and predictions in the eval WER calculation."},
+        metadata={
+            "help": "Whether to normalise the references and predictions in the eval WER calculation."},
     )
     language: str = field(
         default=None,
@@ -246,9 +259,11 @@ class DataTrainingArguments:
     )
     task: str = field(
         default="transcribe",
-        metadata={"help": "Task, either `transcribe` for speech recognition or `translate` for speech translation."},
+        metadata={
+            "help": "Task, either `transcribe` for speech recognition or `translate` for speech translation."},
     )
-    num_train_steps: int = field(default=50000, metadata={"help": "The number of training steps."})
+    num_train_steps: int = field(default=50000, metadata={
+                                 "help": "The number of training steps."})
     # num_eval_samples: int = field(default=50000, metadata={"help": "The number of samples to be used for evaluation"})
     shuffle_buffer_size: Optional[int] = field(
         default=500,
@@ -261,9 +276,11 @@ class DataTrainingArguments:
     )
     streaming: bool = field(
         default=True,
-        metadata={"help": "Whether to use streaming mode to load and pre-process the data."},
+        metadata={
+            "help": "Whether to use streaming mode to load and pre-process the data."},
     )
 
+
 def shift_tokens_right(label_ids: np.array, decoder_start_token_id: int) -> np.ndarray:
     """
     Shift label ids one token to the right.
@@ -348,17 +365,19 @@ class FlaxDataCollatorSpeechSeq2SeqWithPadding:
             labels = labels[:, 1:]
             labels_batch.attention_mask = labels_batch.attention_mask[:, 1:]
 
-        decoder_input_ids = shift_tokens_right(labels, self.decoder_start_token_id)
+        decoder_input_ids = shift_tokens_right(
+            labels, self.decoder_start_token_id)
 
         # replace padding with -100 to ignore correctly when computing the loss
-        labels = np.ma.array(labels, mask=np.not_equal(labels_batch.attention_mask, 1))
+        labels = np.ma.array(labels, mask=np.not_equal(
+            labels_batch.attention_mask, 1))
         labels = labels.filled(fill_value=-100)
 
         batch["labels"] = labels
         batch["decoder_input_ids"] = decoder_input_ids
 
         return batch
-        
+
 
 def load_maybe_streaming_dataset(dataset_name, dataset_config_name, split="train", streaming=True, **kwargs):
     """
@@ -369,7 +388,8 @@ def load_maybe_streaming_dataset(dataset_name, dataset_config_name, split="train
     if "+" in split:
         # load multiple splits separated by the `+` symbol with streaming mode
         dataset_splits = [
-            load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=streaming, **kwargs)
+            load_dataset(dataset_name, dataset_config_name,
+                         split=split_name, streaming=streaming, **kwargs)
             for split_name in split.split("+")
         ]
         # interleave multiple splits to form one dataset
@@ -377,7 +397,8 @@ def load_maybe_streaming_dataset(dataset_name, dataset_config_name, split="train
         return interleaved_dataset
     else:
         # load a single split *with* streaming mode
-        dataset = load_dataset(dataset_name, dataset_config_name, split=split, streaming=streaming, **kwargs)
+        dataset = load_dataset(
+            dataset_name, dataset_config_name, split=split, streaming=streaming, **kwargs)
         return dataset
 
 
@@ -394,7 +415,8 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
 
     if drop_last:
         steps_per_epoch = len(dataset) // batch_size
-        batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
+        # Skip incomplete batch.
+        batch_idx = batch_idx[: steps_per_epoch * batch_size]
         batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
     else:
         steps_per_epoch = math.ceil(len(dataset) / batch_size)
@@ -429,11 +451,13 @@ 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)
+    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])
+    schedule_fn = optax.join_schedules(
+        schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
     return schedule_fn
 
 
@@ -442,18 +466,21 @@ def main():
     # See all possible arguments in src/transformers/training_args.py
     # or by passing the --help flag to this script.
     # We now keep distinct sets of args, for a cleaner separation of concerns.
-    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
+    parser = HfArgumentParser(
+        (ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
 
     if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
         # If we pass only one argument to the script and it's the path to a json file,
         # let's parse it to get our arguments.
-        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+        model_args, data_args, training_args = parser.parse_json_file(
+            json_file=os.path.abspath(sys.argv[1]))
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
     # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
     # information sent is the one passed as arguments along with your JAX/Flax versions.
-    send_example_telemetry("run_speech_recognition_seq2seq", model_args, data_args, framework="flax")
+    send_example_telemetry("run_speech_recognition_seq2seq",
+                           model_args, data_args, framework="flax")
 
     # 2. Setup logging
     # Make one log on every process with the configuration for debugging.
@@ -464,7 +491,8 @@ def main():
     )
     # 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)
+    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()
@@ -490,16 +518,18 @@ def main():
     if training_args.push_to_hub:
         if training_args.hub_model_id is None:
             repo_name = get_full_repo_name(
-                Path(training_args.output_dir).absolute().name, token=training_args.hub_token
+                Path(training_args.output_dir).absolute(
+                ).name, token=training_args.hub_token
             )
         else:
             repo_name = training_args.hub_model_id
         create_repo(repo_name, exist_ok=True, token=training_args.hub_token)
-        repo = Repository(training_args.output_dir, clone_from=repo_name, token=training_args.hub_token)
+        repo = Repository(training_args.output_dir,
+                          clone_from=repo_name, token=training_args.hub_token)
 
     # 3. Load dataset
     raw_datasets = IterableDatasetDict() if data_args.streaming else DatasetDict()
-    
+
     if training_args.do_train:
         raw_datasets["train"] = load_maybe_streaming_dataset(
             data_args.dataset_name,
@@ -519,13 +549,14 @@ def main():
             streaming=data_args.streaming,
             use_auth_token=True if model_args.use_auth_token else None,
         )
-    
+
     if not training_args.do_train and not training_args.do_eval:
         raise ValueError(
             "Cannot not train and not do evaluation. At least one of training or evaluation has to be performed."
         )
 
-    raw_datasets_features = list(next(iter(raw_datasets.values())).features.keys())
+    raw_datasets_features = list(
+        next(iter(raw_datasets.values())).features.keys())
 
     if data_args.audio_column_name not in raw_datasets_features:
         raise ValueError(
@@ -572,21 +603,26 @@ def main():
     )
 
     if model.config.decoder_start_token_id is None:
-        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
+        raise ValueError(
+            "Make sure that `config.decoder_start_token_id` is correctly defined")
 
     # 6. Resample speech dataset: `datasets` takes care of automatically loading and resampling the audio,
     # so we just need to set the correct target sampling rate.
-    dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate
-    
+    dataset_sampling_rate = next(
+        iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate
+
     if dataset_sampling_rate != feature_extractor.sampling_rate:
         raw_datasets = raw_datasets.cast_column(
-            data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
+            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_input_length = int(
+        data_args.max_duration_in_seconds * feature_extractor.sampling_rate)
+    min_input_length = int(
+        data_args.min_duration_in_seconds * feature_extractor.sampling_rate)
     max_label_length = (
         data_args.max_label_length if data_args.max_label_length is not None else model.config.max_length
     )
@@ -602,18 +638,21 @@ def main():
 
     if data_args.language is not None:
         # We only need to set the task id when the language is specified (i.e. in a multilingual setting)
-        tokenizer.set_prefix_tokens(language=data_args.language, task=data_args.task)
+        tokenizer.set_prefix_tokens(
+            language=data_args.language, task=data_args.task)
 
     def prepare_dataset(batch):
         # process audio
         sample = batch[audio_column_name]
-        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
+        inputs = feature_extractor(
+            sample["array"], sampling_rate=sample["sampling_rate"])
         # process audio length
         batch[model_input_name] = inputs.get(model_input_name)[0]
         batch["input_length"] = len(sample["array"])
 
         # process targets
-        input_str = batch[text_column_name].lower() if do_lower_case else batch[text_column_name]
+        input_str = batch[text_column_name].lower(
+        ) if do_lower_case else batch[text_column_name]
         if do_remove_punctuation:
             input_str = normalizer(input_str).strip()
         batch["labels"] = tokenizer(input_str).input_ids
@@ -624,7 +663,7 @@ def main():
             prepare_dataset,
             remove_columns=raw_datasets_features,
         ).with_format("torch")
-  
+
     # filter training data with inputs longer than max_input_length
     def is_audio_in_length_range(length):
         return min_input_length < length < max_input_length
@@ -634,14 +673,13 @@ def main():
             is_audio_in_length_range,
             input_columns=["input_length"],
         )
-        
+
     if training_args.do_eval:
         vectorized_datasets["eval"] = vectorized_datasets["eval"].filter(
             is_audio_in_length_range,
             input_columns=["input_length"],
         )
 
-
     # 8. Load Metric
     metric = evaluate.load("wer")
     do_normalize_eval = data_args.do_normalize_eval
@@ -660,8 +698,10 @@ def main():
             pred_str = [normalizer(pred) for pred in pred_str]
             label_str = [normalizer(label) for label in label_str]
             # filtering step to only evaluate the samples that correspond to non-zero references:
-            pred_str = [pred_str[i] for i in range(len(pred_str)) if len(label_str[i]) > 0]
-            label_str = [label_str[i] for i in range(len(label_str)) if len(label_str[i]) > 0]
+            pred_str = [pred_str[i]
+                        for i in range(len(pred_str)) if len(label_str[i]) > 0]
+            label_str = [label_str[i]
+                         for i in range(len(label_str)) if len(label_str[i]) > 0]
 
         wer = 100 * metric.compute(predictions=pred_str, references=label_str)
 
@@ -690,7 +730,8 @@ def main():
         try:
             from flax.metrics.tensorboard import SummaryWriter
 
-            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
+            summary_writer = SummaryWriter(
+                log_dir=Path(training_args.output_dir))
         except ImportError as ie:
             has_tensorboard = False
             logger.warning(
@@ -708,10 +749,10 @@ def main():
 
     # Store some constant
     #num_epochs = int(training_args.num_train_epochs)
-    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
-    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
-
-    
+    train_batch_size = int(
+        training_args.per_device_train_batch_size) * jax.device_count()
+    eval_batch_size = int(
+        training_args.per_device_eval_batch_size) * jax.device_count()
 
     # Create learning rate schedule
     linear_decay_lr_schedule_fn = create_learning_rate_fn(
@@ -736,7 +777,8 @@ def main():
                 if layer_norm_name in "".join(layer).lower()
             ]
         )
-        flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_named_params) for path in flat_params}
+        flat_mask = {path: (path[-1] != "bias" and path[-2:]
+                            not in layer_norm_named_params) for path in flat_params}
         return traverse_util.unflatten_dict(flat_mask)
 
     # create adam optimizer
@@ -750,7 +792,8 @@ def main():
     )
 
     # Setup train state
-    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
+    state = TrainState.create(
+        apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
 
     # label smoothed cross entropy
     def loss_fn(logits, labels, label_smoothing_factor=0.0):
@@ -762,9 +805,11 @@ def main():
         confidence = 1.0 - label_smoothing_factor
         low_confidence = (1.0 - confidence) / (vocab_size - 1)
         normalizing_constant = -(
-            confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
+            confidence * jnp.log(confidence) + (vocab_size - 1) *
+            low_confidence * jnp.log(low_confidence + 1e-20)
         )
-        soft_labels = onehot(labels, vocab_size, on_value=confidence, off_value=low_confidence)
+        soft_labels = onehot(labels, vocab_size,
+                             on_value=confidence, off_value=low_confidence)
 
         loss = optax.softmax_cross_entropy(logits, soft_labels)
         loss = loss - normalizing_constant
@@ -782,7 +827,8 @@ def main():
 
         def compute_loss(params):
             labels = batch.pop("labels")
-            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+            logits = state.apply_fn(
+                **batch, params=params, dropout_rng=dropout_rng, train=True)[0]
             loss, num_labels = loss_fn(logits, labels, label_smoothing_factor)
             return loss, num_labels
 
@@ -797,9 +843,11 @@ def main():
         # true grad = total grad / total samples
         grad = jax.lax.psum(grad, "batch")
         grad = jax.tree_util.tree_map(lambda x: x / num_labels, grad)
-        new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
+        new_state = state.apply_gradients(
+            grads=grad, dropout_rng=new_dropout_rng)
 
-        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+        metrics = {"loss": loss,
+                   "learning_rate": linear_decay_lr_schedule_fn(state.step)}
         return new_state, metrics
 
     # Define eval fn
@@ -823,27 +871,32 @@ def main():
 
     def generate_step(params, batch):
         model.params = params
-        output_ids = model.generate(batch[model_input_name], attention_mask=batch.get("attention_mask"), **gen_kwargs)
+        output_ids = model.generate(batch[model_input_name], attention_mask=batch.get(
+            "attention_mask"), **gen_kwargs)
         return output_ids.sequences
 
     # Create parallel version of the train and eval step
     p_train_step = jax.pmap(
         partial(train_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch", donate_argnums=(0,)
     )
-    p_eval_step = jax.pmap(partial(eval_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch")
+    p_eval_step = jax.pmap(partial(
+        eval_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch")
     p_generate_step = jax.pmap(generate_step, "batch")
 
     # Replicate the train state on each device
     state = state.replicate()
 
     logger.info("***** Running training *****")
-    logger.info(f"  Num examples = {data_args.num_train_steps*train_batch_size}")
-    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
+    logger.info(
+        f"  Num examples = {data_args.num_train_steps*train_batch_size}")
+    logger.info(
+        f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
     logger.info(f"  Total optimization steps = {data_args.num_train_steps}")
 
     train_time = 0
-    
+
     # ======================== Training ================================
     train_start = time.time()
 
@@ -859,29 +912,32 @@ def main():
     num_workers = 0
     # This is not working
     # vectorized_datasets["train"] = vectorized_datasets["train"].shuffle()
-    train_data_loader = torch.utils.data.DataLoader( batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=True)
-    train_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(train_data_loader)
-    
-    
+    train_data_loader = torch.utils.data.DataLoader(
+        batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=True)
+    train_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(
+        train_data_loader)
+
     # train
-    for step in tqdm(range(data_args.num_train_steps), desc="Training...", position=1, leave=False):        
-        
+    for step in tqdm(range(data_args.num_train_steps), desc="Training...", position=1, leave=False):
+
         try:
             samples = next(train_data_iterator)
-            
+
         except StopIteration:
             epoch += 1
-            train_data_loader = torch.utils.data.DataLoader( batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=True)
-            train_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(train_data_loader)            
+            train_data_loader = torch.utils.data.DataLoader(
+                batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=True)
+            train_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(
+                train_data_loader)
             samples = next(train_data_iterator)
-            
+
             logger.info(
                 f"Completed epoch ({epoch} | Loss: {train_metric['loss']}, Learning Rate:"
                 f" {train_metric['learning_rate']})"
             )
-        
+
         # reshaped_samples = {key: [feature[key] for feature in samples] for key in samples[0].keys()}
-        #breakpoint()
+        # breakpoint()
         batch = data_collator(samples)
         batch = shard(batch.data)
         state, train_metric = p_train_step(state, batch)
@@ -896,8 +952,10 @@ def main():
             eval_labels = []
 
             #eval_loader = data_loader(input_rng, vectorized_datasets["eval"], eval_batch_size, drop_last=False)
-            eval_data_loader = torch.utils.data.DataLoader( batch_size=eval_batch_size, dataset=vectorized_datasets["eval"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=False)
-            eval_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(eval_data_loader)
+            eval_data_loader = torch.utils.data.DataLoader(
+                batch_size=eval_batch_size, dataset=vectorized_datasets["eval"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=False)
+            eval_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(
+                eval_data_loader)
 
             for _ in tqdm(range(training_args.eval_steps), desc="Evaluating...", position=2, leave=False):
                 # Model forward
@@ -912,10 +970,12 @@ def main():
 
             # generation
             if training_args.predict_with_generate:
-                generated_ids = pad_shard_unpad(p_generate_step)(state.params, batch.data)
-                eval_preds.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
+                generated_ids = pad_shard_unpad(
+                    p_generate_step)(state.params, batch.data)
+                eval_preds.extend(jax.device_get(
+                    generated_ids.reshape(-1, gen_kwargs["max_length"])))
                 eval_labels.extend(labels)
-        
+            breakpoint()
             # normalize eval metrics
             eval_metrics = get_metrics(eval_metrics)
             eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
@@ -925,7 +985,8 @@ def main():
             if training_args.predict_with_generate:
                 wer_metric = compute_metrics(eval_preds, eval_labels)
                 eval_metrics.update(wer_metric)
-                wer_desc = " ".join([f"Eval {key}: {value} |" for key, value in wer_metric.items()])
+                wer_desc = " ".join(
+                    [f"Eval {key}: {value} |" for key, value in wer_metric.items()])
 
             # Print metrics
             desc = f"Epoch... ({epoch} | Eval Loss: {eval_metrics['loss']} | {wer_desc})"
@@ -933,15 +994,18 @@ def main():
 
             # Save metrics
             if has_tensorboard and jax.process_index() == 0:
-                write_metric(summary_writer, train_metrics, eval_metrics, train_time, step)
+                write_metric(summary_writer, train_metrics,
+                             eval_metrics, train_time, step)
 
             # save checkpoint after each epoch and push checkpoint to the hub
             if jax.process_index() == 0:
-                params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
+                params = jax.device_get(
+                    jax.tree_util.tree_map(lambda x: x[0], state.params))
                 model.save_pretrained(training_args.output_dir, params=params)
                 tokenizer.save_pretrained(training_args.output_dir)
                 if training_args.push_to_hub:
-                    repo.push_to_hub(commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False)
+                    repo.push_to_hub(
+                        commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False)
 
 
 if __name__ == "__main__":

run_streaming.sh

@@ -2,6 +2,7 @@ python run_flax_speech_recognition_seq2seq_streaming_v3.py \
             --model_name_or_path openai/whisper-tiny.en \
             --dataset_name mozilla-foundation/common_voice_11_0 \
             --dataset_config es \
+            --language es \
 	        --text_column_name sentence \
             --train_split_name test\
             --eval_split_name test\