run_parler_tts_training.py

#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" Train Parler-TTS using 🤗 Accelerate"""

import logging
import os
import re
import sys
import time
import math
import shutil
import yaml
import random
import mlflow
import contextlib
import soundfile as sf
from multiprocess import set_start_method
from datetime import timedelta
import inspect
from tqdm import tqdm
from pathlib import Path

import torch
from torch.utils.data import DataLoader

import datasets
from datasets import DatasetDict, Dataset, IterableDataset, concatenate_datasets

from huggingface_hub import HfApi

import transformers
from transformers import AutoFeatureExtractor, AutoTokenizer, HfArgumentParser
from transformers.trainer_pt_utils import LengthGroupedSampler
from transformers.optimization import get_scheduler
from transformers.utils import send_example_telemetry


from accelerate import Accelerator, skip_first_batches
from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin, DistributedDataParallelKwargs
from accelerate.utils.memory import release_memory

from parler_tts import (
    ParlerTTSConfig,
    ParlerTTSForConditionalGeneration,
    build_delay_pattern_mask,
)

from training.utils import (
    get_last_checkpoint,
    rotate_checkpoints,
    log_pred,
    log_metric,
    load_all_codec_checkpoints,
    save_codec_checkpoint,
    get_last_codec_checkpoint_step,
)
from training.arguments import ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments
from training.data import load_multiple_datasets, DataCollatorParlerTTSWithPadding, DataCollatorEncodecWithPadding, DynamicSizeBatchSampler
from training.eval import clap_similarity, wer, si_sdr

logger = logging.getLogger(__name__)


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.
    
    args_yaml = None
    for i, arg in enumerate(sys.argv):
        if arg.startswith("--args_yaml"):
            args_yaml = sys.argv[i + 1] if i + 1 < len(sys.argv) else None
            del sys.argv[i:i + 2]
            break
        
    default_args = []
    if args_yaml is not None and os.path.isfile(args_yaml):
        with open(args_yaml) as f:
            default_configs = yaml.load(f, Loader=yaml.FullLoader)
        default_args = []
        for k, v in default_configs.items():
            if isinstance(v, str) and v.startswith("$"):
                v = os.environ.get(v[1:], None)
            default_args.append(f"--{k}")
            if isinstance(v, list):
                for sub_v in v:
                    sub_v = str(sub_v)
                    default_args.append(sub_v)
            else:
                v = str(v)
                default_args.append(v)
        default_args.extend(sys.argv[1:])

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
    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(args=default_args)

    # 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 Python/PyTorch versions.
    send_example_telemetry("run_parler_tts", model_args, data_args)

    if training_args.dtype == "float16":
        mixed_precision = "fp16"
        torch_dtype = torch.float16
    elif training_args.dtype == "bfloat16":
        mixed_precision = "bf16"
        torch_dtype = torch.bfloat16
    else:
        mixed_precision = "no"
        torch_dtype = torch.float32

    if data_args.pad_to_max_length and (
        data_args.max_duration_in_seconds is None
        or data_args.max_prompt_token_length is None
        or data_args.max_description_token_length is None
    ):
        raise ValueError(
            "`pad_to_max_length` is `True` but one of the following parameters has not been set: `max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`"
        )

    padding = "max_length" if data_args.pad_to_max_length else "longest"

    ####### A. Preparation
    kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=10000000)), DistributedDataParallelKwargs(find_unused_parameters=True)]
    accelerator = Accelerator(
        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
        mixed_precision=mixed_precision,
        log_with=training_args.report_to,
        project_dir=training_args.output_dir,
        even_batches=training_args.dynamic_batch_size_map is None,
        kwargs_handlers=kwargs_handlers,
    )

    init_kwargs = {}
    if "mlflow" in training_args.report_to:
        project_name = data_args.mlflow_project
        init_kwargs["mlflow"] = {"run_id": data_args.mlflow_run_id} if data_args.mlflow_run_id else {}
    elif "wandb" in training_args.report_to:
        project_name = data_args.wandb_project
        init_kwargs["wandb"] = {"name": data_args.wandb_run_name} if data_args.wandb_run_name else {}
    else:
        raise ValueError(f"Unsupported logging platform: {training_args.report_to}")
        
    accelerator.init_trackers(
        project_name=project_name,
        config={
            "learning_rate": training_args.learning_rate,
            "model_name_or_path": model_args.model_name_or_path,
            "num_train_epochs": training_args.num_train_epochs,
            "gradient_accumulation_steps": training_args.gradient_accumulation_steps,
            "per_device_train_batch_size": training_args.per_device_train_batch_size,
            "global_batch_size": training_args.per_device_train_batch_size * accelerator.num_processes,
            "mixed_precision": mixed_precision,
            "lr_scheduler_type": training_args.lr_scheduler_type,
            "warmup_steps": training_args.warmup_steps,
            "freeze_text_encoder": model_args.freeze_text_encoder,
            "max_duration_in_seconds": data_args.max_duration_in_seconds,
            "weight_decay": training_args.weight_decay,
            "adam_beta1": training_args.adam_beta1,
            "adam_beta2": training_args.adam_beta2,
            "temperature": model_args.temperature,
        },
        init_kwargs=init_kwargs,
    )

    # Detecting last checkpoint and eventually continue from last checkpoint
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)

    # Log a small summary on each proces
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
    )

    # Set the verbosity to info of the Transformers logger (on main process only)
    if accelerator.is_local_main_process:
        datasets.utils.logging.set_verbosity_warning()
        transformers.utils.logging.set_verbosity_info()
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed before initializing model.
    set_seed(training_args.seed)
    num_workers = data_args.preprocessing_num_workers

    # 1. First, lett's instantiate the feature extractor, tokenizers and model
    # Note for distributed training, the .from_pretrained methods guarantee that only
    # one local process can concurrently download model & vocab.

    # load feature extractor
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        model_args.feature_extractor_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
    sampling_rate = feature_extractor.sampling_rate

    # load prompt tokenizer
    prompt_tokenizer = AutoTokenizer.from_pretrained(
        model_args.prompt_tokenizer_name or model_args.description_tokenizer_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
        use_fast=model_args.use_fast_tokenizer,
        padding_side=model_args.prompt_padding_side,
    )

    # load description tokenizer
    description_tokenizer = AutoTokenizer.from_pretrained(
        model_args.description_tokenizer_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
        use_fast=model_args.use_fast_tokenizer,
    )

    if model_args.use_fast_tokenizer:
        logger.warning(
            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
        )
        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True

    # 2. Now, let's load the dataset

    if data_args.save_to_disk is not None:
        os.makedirs(data_args.save_to_disk, exist_ok=True)

    # assume that the dataset has been saved to `save_to_disk` if the latter is not empty
    dataset_was_precomputed = len(os.listdir(data_args.save_to_disk)) > 0
    if dataset_was_precomputed:
        logger.info("load from disk")
        with accelerator.main_process_first():
            vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
        vectorized_datasets["train"] = vectorized_datasets["train"].select(range(500000))
        vectorized_datasets["eval"] = vectorized_datasets["eval"].select(range(30))
        
    else:
        raw_datasets = DatasetDict()

        columns_to_keep = {
            "target_audio_column_name": data_args.target_audio_column_name,
            "prompt_column_name": data_args.prompt_column_name,
            "speaker": "speaker",
            "gender": "gender",
            "age": "age",
        }
        if data_args.description_column_name is not None:
            columns_to_keep["description_column_name"] = data_args.description_column_name

        print("streaming", data_args.streaming)
        if training_args.do_train:
            raw_datasets["train"] = load_multiple_datasets(
                accelerator,
                data_args.train_dataset_name,
                data_args.train_dataset_config_name,
                metadata_dataset_names=data_args.train_metadata_dataset_name,
                splits=data_args.train_split_name,
                dataset_samples=data_args.train_dataset_samples,
                seed=training_args.seed,
                cache_dir=model_args.cache_dir,
                num_proc=data_args.preprocessing_num_workers,
                id_column_name=data_args.id_column_name,
                columns_to_keep=columns_to_keep.values(),
                prompt_column_name=data_args.prompt_column_name,
                audio_column_name=data_args.target_audio_column_name,
                sampling_rate=sampling_rate,
                logger=logger,
                streaming=data_args.streaming, # TODO(SG): optionally enable streaming mode
                sharded=data_args.sharded,
            )

            for key in columns_to_keep:
                if columns_to_keep[key] not in raw_datasets["train"].column_names:
                    raise ValueError(
                        f"--{key} '{columns_to_keep[key]}' not found in dataset '{data_args.train_dataset_name}'."
                        f" Make sure to set `--{key}` to the correct audio column - one of"
                        f" {', '.join(raw_datasets['train'].column_names)}."
                    )

            if data_args.max_train_samples is not None:
                raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))

        if training_args.do_eval:
            raw_datasets["eval"] = load_multiple_datasets(
                accelerator,
                data_args.eval_dataset_name if data_args.eval_dataset_name else data_args.train_dataset_name,
                data_args.eval_dataset_config_name
                if data_args.eval_dataset_config_name
                else data_args.train_dataset_config_name,
                metadata_dataset_names=data_args.eval_metadata_dataset_name,
                splits=data_args.eval_split_name,
                cache_dir=model_args.cache_dir,
                num_proc=data_args.preprocessing_num_workers,
                id_column_name=data_args.id_column_name,
                columns_to_keep=columns_to_keep.values(),
                prompt_column_name=data_args.prompt_column_name,
                audio_column_name=data_args.target_audio_column_name,
                sampling_rate=sampling_rate,
                logger=logger,
                use_auth_token=data_args.use_auth_token,
                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
            )

            if data_args.max_eval_samples is not None:
                with accelerator.main_process_first():
                    raw_datasets["eval"] = (
                        raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
                    )

    # 3. Next, let's load the config.
    config = ParlerTTSConfig.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
 
    if training_args.codebook_weights is not None and len(training_args.codebook_weights) != config.decoder.num_codebooks:
        raise ValueError(f"`codebook_weights` has length {len(training_args.codebook_weights)} when it should be of length {config.decoder.num_codebooks}.")

    # update pad token id and decoder_start_token_id
    config.decoder.update(
        {
            "cross_attention_implementation_strategy": model_args.cross_attention_implementation_strategy
            if model_args.cross_attention_implementation_strategy is not None
            else None,
            "codebook_weights": training_args.codebook_weights if training_args.codebook_weights is not None else config.decoder.codebook_weights
        }
    )
    config.update(
        {
            "pad_token_id": model_args.pad_token_id if model_args.pad_token_id is not None else config.pad_token_id,
            "decoder_start_token_id": model_args.decoder_start_token_id
            if model_args.decoder_start_token_id is not None
            else config.decoder_start_token_id,
        }
    )

    # create model
    model = ParlerTTSForConditionalGeneration.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        config=config,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
        attn_implementation={"decoder": model_args.attn_implementation, "text_encoder": "eager"},
    )
    
    # print("** model **")
    # print(model)

    # enable gradient checkpointing if necessary
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,
    # so that we just need to set the correct target sampling rate and normalize the input
    # via the `feature_extractor`

    # derive max & min input length for sample rate & max duration
    sampling_rate = feature_extractor.sampling_rate
    max_target_length = int(data_args.max_duration_in_seconds * sampling_rate)
    min_target_length = int(data_args.min_duration_in_seconds * sampling_rate)
    target_audio_column_name = data_args.target_audio_column_name
    description_column_name = data_args.description_column_name
    prompt_column_name = data_args.prompt_column_name
    feature_extractor_input_name = feature_extractor.model_input_names[0]
    audio_encoder_pad_token_id = config.decoder.pad_token_id
    audio_encoder_eos_token_id = config.decoder.eos_token_id
    audio_encoder_bos_token_id = model.generation_config.decoder_start_token_id
    max_length = model.generation_config.max_length
    num_codebooks = model.decoder.config.num_codebooks
    bandwidth = model_args.bandwidth
    attn_implementation = model_args.attn_implementation

    # Freeze Encoders
    model.freeze_encoders(model_args.freeze_text_encoder)

    # Test all gather - used for warmout and avoiding timeout
    logger.debug(str(accelerator.process_index), main_process_only=False, in_order=True)
    test_tensor = torch.tensor([accelerator.process_index], device=accelerator.device)
    gathered_tensor = accelerator.gather(test_tensor)
    print("gathered_tensor", gathered_tensor)
    accelerator.wait_for_everyone()

    if not dataset_was_precomputed:
        # Filter on text length
        if description_column_name is not None and data_args.max_text_length is not None:
            with accelerator.main_process_first():
                # filter description that is shorter than max_text_length
                raw_datasets = raw_datasets.filter(
                    lambda x: len(x) < data_args.max_text_length,
                    num_proc=num_workers,
                    input_columns=[description_column_name],
                )

        # Preprocessing the dataset.
        # We need to tokenize the texts.
        def pass_through_processors(prompt, speaker=None, gender=None, age=None, description=None):
            batch = {}
            
            # print(f"speaker: {speaker}, gender: {gender}, age: {age}")

            if description is not None:
                batch["input_ids"] = description_tokenizer(description.strip())["input_ids"]
                
            prefix = ""
            if gender is not None:
                prefix += gender
            if age is not None:
                prefix += ' ' + age
            prefix = prefix.strip()
            
            if prefix != "":
                batch["prompt_input_ids"] = prompt_tokenizer(f"{prefix} {prompt.strip()}")["input_ids"]
            else:
                batch["prompt_input_ids"] = prompt_tokenizer(f"{prompt.strip()}")["input_ids"]
            
            # print('with meta', prompt_tokenizer(f"{speaker} {gender} {age} {prompt.strip()}")["input_ids"])
            # print('original', prompt_tokenizer(prompt.strip())["input_ids"])

            return batch

        with accelerator.main_process_first():
            # this is a trick to avoid to rewrite the entire audio column which takes ages
            input_columns = [description_column_name, prompt_column_name, 'speaker', 'gender', 'age'] if description_column_name is not None else [prompt_column_name, 'speaker', 'gender', 'age']
            vectorized_datasets = raw_datasets.map(
                pass_through_processors,
                remove_columns=next(iter(raw_datasets.values())).column_names,
                input_columns=input_columns,
                num_proc=num_workers,
                desc="preprocess datasets",
            )

        # We use Accelerate to perform distributed inference
        # T5 doesn't support fp16
        autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))

        # Now we encode the audio labels with encodec.
        ####### B. Encode audio

        logger.info("*** Encode target audio with encodec ***")

        # no need to prepare audio_decoder because used for inference without mixed precision
        # see: https://huggingface.co/docs/accelerate/main/en/package_reference/accelerator#accelerate.Accelerator.prepare
        if training_args.torch_compile:
            audio_decoder = accelerator.prepare_model(model.audio_encoder, evaluation_mode=True)
        else:
            audio_decoder = model.audio_encoder

        encoder_data_collator = DataCollatorEncodecWithPadding(
            feature_extractor,
            audio_column_name=target_audio_column_name,
            feature_extractor_input_name=feature_extractor_input_name,
            max_length=max_target_length,
            padding=padding,
        )
        encoder_signature = set(inspect.signature(audio_decoder.forward).parameters)

        def apply_audio_decoder(batch):
            len_audio = batch.pop("len_audio")
            audio_decoder.to(batch["input_values"].device).eval()
            if bandwidth is not None:
                batch["bandwidth"] = bandwidth
            elif "num_quantizers" in encoder_signature:
                batch["num_quantizers"] = num_codebooks
            elif "num_codebooks" in encoder_signature:
                batch["num_codebooks"] = num_codebooks
            elif "n_quantizers" in encoder_signature:
                batch["n_quantizers"] = num_codebooks

            with torch.no_grad():
                labels = audio_decoder.encode(**batch)["audio_codes"]
            output = {}
            output["len_audio"] = len_audio
            # (1, bsz, codebooks, seq_len) -> (bsz, seq_len, codebooks)
            output["labels"] = labels.squeeze(0).transpose(1, 2)

            # if `pad_to_max_length`, the maximum corresponding audio length of the current batch is max_duration*sampling_rate
            max_length = len_audio.max() if padding != "max_length" else max_target_length
            output["ratio"] = torch.ones_like(len_audio) * labels.shape[-1] / max_length
            return output

        # (1, codebooks, seq_len) where seq_len=1
        bos_labels = torch.ones((1, num_codebooks, 1)) * audio_encoder_bos_token_id

        def postprocess_dataset(labels):
            # (1, codebooks, seq_len)
            labels = torch.tensor(labels).unsqueeze(0)
            # add bos
            labels = torch.cat([bos_labels, labels], dim=-1)

            labels, delay_pattern_mask = build_delay_pattern_mask(
                labels,
                bos_token_id=audio_encoder_bos_token_id,
                pad_token_id=audio_encoder_eos_token_id,
                max_length=labels.shape[-1] + num_codebooks,
                num_codebooks=num_codebooks,
            )

            # the first ids of the delay pattern mask are precisely labels, we use the rest of the labels mask
            # to take care of EOS
            # we want labels to look like this:
            #  - [B, a, b, E, E, E, E]
            #  - [B, B, c, d, E, E, E]
            #  - [B, B, B, e, f, E, E]
            #  - [B, B, B, B, g, h, E]
            labels = torch.where(delay_pattern_mask == -1, audio_encoder_eos_token_id, delay_pattern_mask)

            # the first timestamp is associated to a row full of BOS, let's get rid of it
            # we also remove the last timestampts (full of PAD)
            output = {"labels": labels[:, 1:]}
            return output

        for split in vectorized_datasets:
            data_loader = DataLoader(
                raw_datasets[split],
                batch_size=training_args.audio_encoder_per_device_batch_size,
                collate_fn=encoder_data_collator,
                num_workers=training_args.dataloader_num_workers,
                pin_memory=True,
            )
            data_loader = accelerator.prepare(data_loader)
            total_inference_steps = len(data_loader)

            start_step = get_last_codec_checkpoint_step(os.path.join(data_args.temporary_save_to_disk, split))
            accelerator.wait_for_everyone()
            if start_step > 0:
                logger.info(f"Resuming {split} from step {start_step}")
                # efficiently skip the first n batches
                start_step += 1
                data_loader = skip_first_batches(data_loader, start_step)

            all_generated_labels = []
            all_lens = []
            if start_step < total_inference_steps:
                for i, batch in enumerate(tqdm(data_loader, disable=not accelerator.is_local_main_process)):
                    cur_step = start_step + i
                    generate_labels = apply_audio_decoder(batch)
                    generate_labels = accelerator.pad_across_processes(generate_labels, dim=1, pad_index=0)
                    generate_labels = accelerator.gather_for_metrics(generate_labels)
                    accelerator.wait_for_everyone()

                    if accelerator.is_main_process:
                        lab = generate_labels["labels"].cpu().transpose(1, 2).to(torch.int16)
                        rat = generate_labels["ratio"].cpu().squeeze(1)
                        lens = generate_labels["len_audio"].cpu().squeeze(1)
                        lab = [l[:, : int(ratio * length)] for (l, ratio, length) in zip(lab, rat, lens)]

                        all_generated_labels.extend(lab)
                        all_lens.extend(lens)

                        if ((cur_step + 1) % data_args.save_codec_steps == 0) or (
                            cur_step == total_inference_steps - 1
                        ):
                            tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
                            tmp_labels = tmp_labels.map(
                                postprocess_dataset,
                                num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
                                input_columns=["labels"],
                                desc="Postprocessing labeling",
                            )
                            save_codec_checkpoint(
                                os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step
                            )
                            all_generated_labels = []
                            all_lens = []

                accelerator.wait_for_everyone()

            if accelerator.is_main_process and len(all_generated_labels) > 0:
                tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
                tmp_labels = tmp_labels.map(
                    postprocess_dataset,
                    num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
                    input_columns=["labels"],
                    desc="Postprocessing labeling",
                )
                save_codec_checkpoint(os.path.join(data_args.temporary_save_to_disk, split), tmp_labels, cur_step)
                all_generated_labels = []
                all_lens = []
            accelerator.wait_for_everyone()

            del all_generated_labels
            accelerator.wait_for_everyone()

            with accelerator.main_process_first():
                tmp_labels = load_all_codec_checkpoints(os.path.join(data_args.temporary_save_to_disk, split)).select(
                    range(len(vectorized_datasets[split]))
                )
                logger.info(f"Concatenating {split}: {tmp_labels} with {vectorized_datasets[split]}")
                vectorized_datasets[split] = concatenate_datasets([vectorized_datasets[split], tmp_labels], axis=1)

        accelerator.free_memory()
        del generate_labels, all_lens

        with accelerator.main_process_first():
            # NOTE: filtering is done at the end because in the `datasets` library, caching audio files is done after most operations
            # caching audio files is time and disk-space consuming, so we want to avoid it at all costs, especially for large (>1Kh) audio datasets.
            # That's also why we avoid to concat the processed datasets (vectorized_datasets) with the audio column present in raw_datasets.

            def is_audio_in_length_range(length):
                return length > min_target_length and length < max_target_length

            # filter data that is shorter than min_target_length
            vectorized_datasets = vectorized_datasets.filter(
                is_audio_in_length_range,
                num_proc=num_workers,
                input_columns=["target_length"],
            )

        if description_column_name is not None and data_args.max_description_token_length is not None:
            with accelerator.main_process_first():
                # filter description that is shorter than max_text_length
                vectorized_datasets = vectorized_datasets.filter(
                    lambda x: len(x) < data_args.max_description_token_length,
                    num_proc=num_workers,
                    input_columns=["input_ids"],
                )

        if data_args.max_prompt_token_length is not None:
            with accelerator.main_process_first():
                # filter description that is shorter than max_text_length
                vectorized_datasets = vectorized_datasets.filter(
                    lambda x: len(x) < data_args.max_prompt_token_length,
                    num_proc=num_workers,
                    input_columns=["prompt_input_ids"],
                )

    if data_args.save_to_disk is not None and not dataset_was_precomputed:
        if accelerator.is_main_process:
            vectorized_datasets.save_to_disk(
                data_args.save_to_disk,
                num_proc=min(data_args.preprocessing_num_workers, len(vectorized_datasets["eval"]) - 1),
            )
        accelerator.wait_for_everyone()
        logger.info(f"Dataset saved at {data_args.save_to_disk}")

    audio_max_length = None
    if padding == "max_length":
        audio_max_length = max(vectorized_datasets["train"]["target_length"])
        with accelerator.main_process_first():
            max_sample = vectorized_datasets["train"].filter(
                lambda x: x == audio_max_length,
                num_proc=num_workers,
                input_columns=["target_length"],
            )
        audio_max_length = max([len(l[0]) for l in max_sample["labels"]])

    if description_column_name is not None and data_args.max_description_token_length is not None:
        with accelerator.main_process_first():
            # filter description that is shorter than max_text_length
            vectorized_datasets = vectorized_datasets.filter(
                lambda x: len(x) < data_args.max_description_token_length,
                num_proc=num_workers,
                input_columns=["input_ids"],
            )

    if data_args.max_prompt_token_length is not None:
        with accelerator.main_process_first():
            # filter description that is shorter than max_text_length
            vectorized_datasets = vectorized_datasets.filter(
                lambda x: len(x) < data_args.max_prompt_token_length,
                num_proc=num_workers,
                input_columns=["prompt_input_ids"],
            )

    if training_args.group_by_length:
        # apply a simple heuristic to take into account audio and text lengths
        def add_target_lengths(target_length, prompt, description=[]):
            return {"target_length": target_length + len(prompt) + len(description)}

        with accelerator.main_process_first():
            input_columns=["target_length", "prompt_input_ids", "input_ids"] if description_column_name is not None else ["target_length", "prompt_input_ids"]
            vectorized_datasets = vectorized_datasets.map(
                add_target_lengths,
                num_proc=num_workers,
                input_columns=input_columns,
            )

    # 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 and data_args.save_to_disk is None:
        raise ValueError(
            "`preprocessing_only=True` but `save_to_disk` is not set. The latter should indicates where to save the dataset locally."
        )
    elif data_args.preprocessing_only:
        logger.info(f"Data preprocessing finished. Files save at {data_args.save_to_disk}")
        return

    # 6. Next, we can prepare the training.

    # Let's use word CLAP similary and WER metrics as our evaluation metrics,
    def compute_metrics(
        audios,
        descriptions,
        prompts,
        device="cpu",
        compute_clap_similarity_metric=False,
        compute_noise_level_metric=False,
        noise_level_to_compute_clean_wer=None,
    ):
        results = {}
        input_ids = descriptions
        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
        audios = [a.float().cpu().numpy() for a in audios]

        if compute_clap_similarity_metric:
            clap_score = clap_similarity(
                model_args.clap_model_name_or_path, texts, audios, device, input_sampling_rate=sampling_rate
            )
            results["clap"] = clap_score

        si_sdr_measures = None
        if compute_noise_level_metric:
            si_sdr_measures = si_sdr(audios, device, input_sampling_rate=sampling_rate)

        word_error, transcriptions, clean_word_error, noisy_word_error, percent_clean_samples = wer(
            model_args.asr_model_name_or_path,
            prompts,
            audios,
            device,
            training_args.per_device_eval_batch_size,
            sampling_rate,
            noise_level_to_compute_clean_wer,
            si_sdr_measures,
        )
        results["wer"] = word_error
        if clean_word_error is not None:
            results["clean_wer"] = clean_word_error
            results["noisy_word_error"] = noisy_word_error
            results["percent_clean_samples"] = percent_clean_samples

        return results, texts, prompts, audios, transcriptions, si_sdr_measures

    # Dummy sampler for counting the number of batches
    if training_args.dynamic_batch_size_map is not None:
        sampler = DynamicSizeBatchSampler(
            lengths=vectorized_datasets["train"]["target_length"],
            batch_size_map=training_args.dynamic_batch_size_map,
        )

    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
    train_batch_size = per_device_train_batch_size * accelerator.num_processes
    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)

    if training_args.max_steps < 0:
        num_epochs = int(training_args.num_train_epochs)
        if training_args.dynamic_batch_size_map is not None:
            steps_per_epoch = len(sampler) // gradient_accumulation_steps
        else:
            steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
        total_train_steps = steps_per_epoch * num_epochs
    elif training_args.max_steps > 0:
        logger.info("max_steps is given, it will override any value given in num_train_epochs")
        total_train_steps = int(training_args.max_steps)
        # Setting a very large number of epochs so we go as many times as necessary over the iterator.
        num_epochs = sys.maxsize
        steps_per_epoch = total_train_steps

    if training_args.eval_steps is None:
        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
        eval_steps = steps_per_epoch
    else:
        eval_steps = training_args.eval_steps
        
    if training_args.eval_generation_steps is None:
        eval_generation_steps = eval_steps
    else:
        eval_generation_steps = training_args.eval_generation_steps

    # T5 doesn't support fp16
    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))

    # Define optimizer, LR scheduler, collator
    optimizer = torch.optim.AdamW(
        params=model.parameters(),
        lr=training_args.learning_rate,
        betas=(training_args.adam_beta1, training_args.adam_beta2),
        eps=training_args.adam_epsilon,
        weight_decay=training_args.weight_decay,
    )

    # LR scheduler gets stepped by `num_processes` each time -> account for this in warmup / total steps
    lr_scheduler = get_scheduler(
        name=training_args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
        num_training_steps=total_train_steps * accelerator.num_processes,
    )

    # Instantiate custom data collator
    data_collator = DataCollatorParlerTTSWithPadding(
        prompt_tokenizer=prompt_tokenizer,
        description_tokenizer=description_tokenizer,
        pad_to_multiple_of=data_args.pad_to_multiple_of,
        padding=padding,
        prompt_max_length=data_args.max_prompt_token_length,
        description_max_length=data_args.max_description_token_length,
        audio_max_length=audio_max_length,
    )

    dummy_data_loader = DataLoader(
        vectorized_datasets["train"],
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        sampler=None,
        shuffle=not training_args.group_by_length,
        num_workers=training_args.dataloader_num_workers,
        pin_memory=training_args.dataloader_pin_memory,
    )
    # Prepare everything with accelerate
    # accelerator get the batch size from the dummy_data_loader
    model, optimizer, lr_scheduler, _ = accelerator.prepare(model, optimizer, lr_scheduler, dummy_data_loader)

    num_examples = len(vectorized_datasets["train"])
    logger.info("***** Running training *****")
    logger.info(f"  Training examples = {len(vectorized_datasets['train'])}")
    logger.info(f"  Eval examples = {len(vectorized_datasets['eval'])}")
    logger.info(f"  Num examples = {num_examples}")
    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
    if training_args.dynamic_batch_size_map is not None:
        logger.info(f"  Num batches (w. dynamic batch size) = {len(sampler)}")
    else:
        logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
        logger.info(
            f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
        )
    logger.info(f"  Total optimization steps = {total_train_steps}")

    # ======================== Training ================================
    train_time = 0
    train_start = time.time()
    steps_trained_progress_bar = tqdm(
        range(total_train_steps), desc="Train steps ... ", position=0, disable=not accelerator.is_local_main_process
    )
    continue_training = True
    epochs_trained = 0
    cur_step = 0

    checkpoint = None
    if training_args.resume_from_checkpoint is not None:
        checkpoint = training_args.resume_from_checkpoint
    elif last_checkpoint is not None:
        checkpoint = last_checkpoint

    if accelerator.is_main_process:
        if training_args.push_to_hub:
            api = HfApi(token=training_args.hub_token)

            # Create repo (repo_name from args or inferred)
            repo_name = training_args.hub_model_id
            if repo_name is None:
                repo_name = Path(training_args.output_dir).absolute().name
            repo_id = api.create_repo(repo_name, exist_ok=True).repo_id

            with open(os.path.join(training_args.output_dir, ".gitignore"), "w+") as gitignore:
                if "wandb" not in gitignore:
                    gitignore.write("wandb\n")
        elif training_args.output_dir is not None:
            os.makedirs(training_args.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    # Now save everything to be able to create a single processor later
    # make sure all processes wait until data is saved
    # only the main process saves them
    if accelerator.is_main_process:
        # save feature extractor, tokenizer and config
        if (
            model_args.prompt_tokenizer_name is None
            and model_args.description_tokenizer_name
            or (model_args.prompt_tokenizer_name == model_args.description_tokenizer_name)
        ):
            prompt_tokenizer.save_pretrained(training_args.output_dir)
        else:
            logger.warning(
                f"Prompt tokenizer ('{model_args.prompt_tokenizer_name}') and description tokenizer ('{model_args.description_tokenizer_name}') are not the same. Saving only the prompt tokenizer."
            )
            prompt_tokenizer.save_pretrained(training_args.output_dir)

        feature_extractor.save_pretrained(training_args.output_dir)
        config.save_pretrained(training_args.output_dir)
    accelerator.wait_for_everyone()

    if checkpoint is not None:
        accelerator.load_state(checkpoint)
        # Find num steps and epoch from saved state string pattern
        pattern = r"checkpoint-(\d+)-epoch-(\d+)"
        match = re.search(pattern, checkpoint)
        cur_step = int(match.group(1))
        epochs_trained = int(match.group(2))

        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
        logger.info(f"  Continuing training from epoch {epochs_trained}")
        logger.info(f"  Continuing training from global step {cur_step}")

        steps_trained_progress_bar.update(cur_step)

        for epoch in range(0, epochs_trained):
            with accelerator.main_process_first():
                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)

        if training_args.max_steps < 0:
            # we know exactly the number of steps per epoch, so can skip through the required number of batches
            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
        else:
            # Currently we don't know how many steps we've taken in the current epoch
            # So we just shuffle the dataset one extra time and start from a fresh epoch
            # This is "good enough" for our purposes but not fully correct
            resume_step = None
            with accelerator.main_process_first():
                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
    else:
        resume_step = None

    gen_kwargs = {
        "do_sample": model_args.do_sample,
        "temperature": model_args.temperature,
        "max_length": model_args.max_length,
        # Because of the delayed pattern mask, generation might stop earlier because of unexpected behaviour
        # on the first tokens of the codebooks that are delayed.
        # This fix the issue.
        "min_new_tokens": num_codebooks + 1,
    }

    # Define gradient update step fn
    def train_step(
        batch,
        accelerator,
        autocast_kwargs,
        num_items_in_batch,
        gradient_accumulation_steps,
    ):
        if mixed_precision == "fp16":
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
                if training_args.parallel_mode.value != "distributed":
                    encoder_outputs = model.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
                else:
                    encoder_outputs = model.module.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
                # we optionnally project last_hidden_state to avoid recomputing every time
                encoder_hidden_states = encoder_outputs.last_hidden_state
                if (
                    config.text_encoder.hidden_size != config.decoder.hidden_size
                    and config.decoder.cross_attention_hidden_size is None
                ):
                    encoder_hidden_states = (
                        model.enc_to_dec_proj(encoder_hidden_states)
                        if training_args.parallel_mode.value != "distributed"
                        else model.module.enc_to_dec_proj(encoder_hidden_states)
                    )

                if batch.get("attention_mask", None) is not None:
                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]

                encoder_outputs.last_hidden_state = encoder_hidden_states
                batch["encoder_outputs"] = encoder_outputs

        outputs = model(**batch, loss_reduction="sum")
        # CE (data) loss
        ce_loss = (outputs.loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch

        metrics = {"loss": ce_loss}
        
        # per CE loss
        per_codebook_losses = outputs.per_codebook_losses
        metrics.update({f"codebook_{i}_loss": ((l  * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch) for (i,l) in enumerate(per_codebook_losses)})
        return ce_loss, metrics

    # Define eval fn
    def eval_step(
        batch,
        accelerator,
        autocast_kwargs,
    ):
        eval_model = model if not training_args.torch_compile else model._orig_mod

        if mixed_precision == "fp16":
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
                if training_args.parallel_mode.value != "distributed":
                    encoder_outputs = model.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
                else:
                    encoder_outputs = model.module.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
                # we optionnally project last_hidden_state to avoid recomputing every time
                encoder_hidden_states = encoder_outputs.last_hidden_state
                if (
                    config.text_encoder.hidden_size != config.decoder.hidden_size
                    and config.decoder.cross_attention_hidden_size is None
                ):
                    encoder_hidden_states = (
                        model.enc_to_dec_proj(encoder_hidden_states)
                        if training_args.parallel_mode.value != "distributed"
                        else model.module.enc_to_dec_proj(encoder_hidden_states)
                    )

                if batch.get("attention_mask", None) is not None:
                    encoder_hidden_states = encoder_hidden_states * batch.get("attention_mask", None)[..., None]

                encoder_outputs.last_hidden_state = encoder_hidden_states
                batch["encoder_outputs"] = encoder_outputs

        with torch.no_grad():
            outputs = eval_model(**batch)
        # CE (data) loss
        ce_loss = outputs.loss
        metrics = {"loss": ce_loss}
        
        # per CE loss
        per_codebook_losses = outputs.per_codebook_losses
        metrics.update({f"codebook_{i}_loss": l for (i,l) in enumerate(per_codebook_losses)})
        return metrics

    def generate_step(batch, accelerator):
        batch.pop("decoder_attention_mask", None)
        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=True)
        if training_args.torch_compile:
            # if the model is compiled, we use the original model bc compile is not compatible with .generate
            eval_model = model._orig_mod

        # since we've might have loaded the weights in fp32, we have to autocast to ensure FA2 weights are in half-precision.
        # with accelerator.autocast(autocast_handler=AutocastKwargs(enabled=(attn_implementation=="flash_attention_2"))):
        output_audios = eval_model.generate(**batch, **gen_kwargs)
        output_audios = accelerator.pad_across_processes(output_audios, dim=1, pad_index=0)
        return output_audios

    model.train()

    total_batched_samples = resume_step if resume_step is not None else 0
    for epoch in range(epochs_trained, num_epochs):
        with accelerator.main_process_first():
            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
        sampler = None
        if training_args.dynamic_batch_size_map is not None:
            sampler = DynamicSizeBatchSampler(
                lengths=vectorized_datasets["train"]["target_length"],
                batch_size_map=training_args.dynamic_batch_size_map,
            )
            train_dataloader = DataLoader(
                vectorized_datasets["train"],
                collate_fn=data_collator,
                batch_sampler=sampler,
                shuffle=False,
                num_workers=training_args.dataloader_num_workers,
                pin_memory=training_args.dataloader_pin_memory,
            )
        else:
            if training_args.group_by_length:
                sampler = LengthGroupedSampler(train_batch_size, lengths=vectorized_datasets["train"]["target_length"])
            train_dataloader = DataLoader(
                vectorized_datasets["train"],
                collate_fn=data_collator,
                batch_size=per_device_train_batch_size,
                sampler=sampler,
                shuffle=not training_args.group_by_length,
                num_workers=training_args.dataloader_num_workers,
                pin_memory=training_args.dataloader_pin_memory,
            )
        train_dataloader = accelerator.prepare(train_dataloader)
        if hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDataset):
            train_dataloader.dataset.set_epoch(epoch)

        if resume_step is not None:
            # Skip the first N batches in the dataloader when resuming from a checkpoint
            logger.info(f"  Skip first {resume_step} batches")
            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
            resume_step = None
            accelerator.wait_for_everyone()

        # We chunkify the epoch iterator into gradient accumulation steps `n` batches
        train_iterator = iter(train_dataloader)
        num_steps_in_epoch = len(train_dataloader)
        remainder = num_steps_in_epoch % gradient_accumulation_steps
        remainder = remainder if remainder != 0 else gradient_accumulation_steps
        total_updates = math.ceil(num_steps_in_epoch / gradient_accumulation_steps)
        
        update_step = -1
        for _ in range(total_updates):
            update_step += 1
            
            # preload the total batch per step
            batch_samples = []
            num_batches_in_step = gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
            for _ in range(num_batches_in_step):
                batch_samples += [next(train_iterator)]
                
            # get num items in batch - if different than BOS and than -100
            num_items_in_batch = sum([(batch["labels"].ne(audio_encoder_bos_token_id) | batch["labels"].ne(-100) | batch["labels"].ne(audio_encoder_eos_token_id)).sum((0,1))[0] for batch in batch_samples])
            num_items_in_batch = accelerator.gather(num_items_in_batch).sum().item()
            accelerator.wait_for_everyone()
            
            # losses = []
            for i,batch in enumerate(batch_samples):
                total_batched_samples += 1
                # ctx = model.no_sync if (i < len(batch_samples) - 1 and accelerator.num_processes > 1) else contextlib.nullcontext
                
                with contextlib.nullcontext():
                    loss, train_metric = train_step(batch, accelerator, autocast_kwargs, num_items_in_batch, gradient_accumulation_steps)
                    accelerator.backward(loss)
                    
                # with ctx():
                #     loss, train_metric = train_step(batch, accelerator, autocast_kwargs, num_items_in_batch, gradient_accumulation_steps)
                #     accelerator.backward(loss)
                    # losses.append(loss.detach())
            
            grad_norm = accelerator.clip_grad_norm_(model.parameters(), training_args.max_grad_norm)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

            # The accelerator has performed an optimization step behind the scenes
            steps_trained_progress_bar.update(1)
            cur_step += 1

            # losses = accelerator.gather(sum(losses)).sum().item() / (accelerator.num_processes * gradient_accumulation_steps)
            
            if cur_step % training_args.logging_steps == 0:
                steps_trained_progress_bar.write(
                    f"Step... ({cur_step} / {total_train_steps} | Loss:"
                    f" {train_metric['loss']}, Learning Rate:"
                    f" {lr_scheduler.get_last_lr()[0]})"
                )
                train_metric["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm
                log_metric(
                    accelerator,
                    metrics=train_metric,
                    learning_rate=lr_scheduler.get_last_lr()[0],
                    train_time=train_time + time.time() - train_start,
                    step=cur_step,
                    epoch=epoch,
                    prefix="train",
                )
                accelerator.wait_for_everyone()

            # save checkpoint and weights after each save_steps and at the end of training
            if (cur_step % training_args.save_steps == 0) or cur_step == total_train_steps:
                intermediate_dir = os.path.join(training_args.output_dir, f"checkpoint-{cur_step}-epoch-{epoch}")
                # safe_serialization=False to avoid shared tensors saving issue (TODO(YL): it's a temporary fix)
                # https://github.com/huggingface/transformers/issues/27293#issuecomment-1872560074
                accelerator.save_state(output_dir=intermediate_dir, safe_serialization=False)
                accelerator.wait_for_everyone()
                if accelerator.is_main_process:
                    rotate_checkpoints(
                        training_args.save_total_limit, output_dir=training_args.output_dir, logger=logger
                    )

                    if cur_step == total_train_steps:
                        # un-wrap student model for save
                        unwrapped_model = accelerator.unwrap_model(model)
                        unwrapped_model.save_pretrained(training_args.output_dir)

                    if training_args.push_to_hub:
                        api.upload_folder(
                            repo_id=repo_id,
                            folder_path=training_args.output_dir,
                            commit_message=f"Saving train state of step {cur_step}",
                            run_as_future=True,
                        )
                accelerator.wait_for_everyone()

            if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
                train_time += time.time() - train_start
                # ======================== Evaluating ==============================
                model.eval()
                eval_metrics = []
                eval_preds = []
                eval_descriptions = []
                eval_prompts = []
                eval_start = time.time()

                # release training input batch
                batch = release_memory(batch)

                validation_dataloader = DataLoader(
                    vectorized_datasets["eval"],
                    collate_fn=data_collator,
                    batch_size=per_device_eval_batch_size,
                    drop_last=False,
                    num_workers=training_args.eval_dataloader_num_workers,
                    pin_memory=training_args.dataloader_pin_memory,
                )
                validation_dataloader = accelerator.prepare(validation_dataloader)

                for batch in tqdm(
                    validation_dataloader,
                    desc=f"Evaluating - Inference ...",
                    position=2,
                    disable=not accelerator.is_local_main_process,
                ):
                    # Model forward
                    eval_metric = eval_step(batch, accelerator, autocast_kwargs)
                    eval_metric = accelerator.gather_for_metrics(eval_metric)
                    accelerator.wait_for_everyone()
                    eval_metric = {key: val.unsqueeze(0) if val.ndim == 0 else val for (key,val) in eval_metric.items()}
                    eval_metrics.append(eval_metric)

                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
                    validation_dataloader = DataLoader(
                        vectorized_datasets["eval"],
                        collate_fn=data_collator,
                        batch_size=per_device_eval_batch_size,
                        drop_last=False,
                        num_workers=training_args.eval_dataloader_num_workers,
                        pin_memory=training_args.dataloader_pin_memory,
                    )
                    validation_dataloader = accelerator.prepare(validation_dataloader)
                    # generation
                    for batch in tqdm(
                        validation_dataloader,
                        desc=f"Evaluating - Generation ...",
                        position=2,
                        disable=not accelerator.is_local_main_process,
                    ):
                        generated_audios = generate_step(batch, accelerator)
                        if config.use_text_encoder:   
                            # Gather all predictions and targets
                            generated_audios, input_ids, prompts = accelerator.pad_across_processes(
                                (generated_audios, batch["input_ids"], batch["prompt_input_ids"]), dim=1, pad_index=0
                            )
                            generated_audios, input_ids, prompts = accelerator.gather_for_metrics(
                                (generated_audios, input_ids, prompts)
                            )
                            accelerator.wait_for_everyone()
                            eval_preds.extend(generated_audios.to("cpu"))
                            eval_descriptions.extend(input_ids.to("cpu"))
                            eval_prompts.extend(prompts.to("cpu"))
                        else:
                            # Gather all predictions and targets
                            generated_audios, prompts = accelerator.pad_across_processes(
                                (generated_audios, batch["prompt_input_ids"]), dim=1, pad_index=0
                            )
                            generated_audios, prompts = accelerator.gather_for_metrics(
                                (generated_audios, prompts)
                            )
                            accelerator.wait_for_everyone()
                            eval_preds.extend(generated_audios.to("cpu"))
                            eval_prompts.extend(prompts.to("cpu"))
                            
                    # Assuming tracker is from accelerator.get_tracker("mlflow", unwrap=True)
                    if accelerator.is_main_process:
                        os.makedirs('./tmp', exist_ok=True)
                        random_audio_idxs = random.sample(list(range(len(eval_preds))), 10)
                        for idx in random_audio_idxs:
                            if torch_dtype == torch.bfloat16:
                                audio = eval_preds[idx].to(torch.float32).cpu().numpy()  # Convert to float32 before converting to NumPy
                            else:
                                audio = eval_preds[idx].cpu().numpy()  # Direct conversion for other types
                            prompt = prompt_tokenizer.decode(eval_prompts[idx], skip_special_tokens=True)
        
                            sample_rate = 8000
                            sf.write(f"tmp/{idx}_audio.wav", audio, sample_rate)
                            print(prompt)
                            with open(f"tmp/{idx}_prompt.txt", "w", encoding="utf-8") as f:
                                f.write(prompt)

                            # Log artifacts using MLflow directly
                            mlflow.log_artifact(f"tmp/{idx}_audio.wav", artifact_path=f"{cur_step}_steps_audio_samples")
                            mlflow.log_artifact(f"tmp/{idx}_prompt.txt", artifact_path=f"{cur_step}_steps_audio_prompts")
                        shutil.rmtree('./tmp', ignore_errors=True)
                    accelerator.wait_for_everyone()

                eval_time = time.time() - eval_start
                # normalize eval metrics
                eval_metrics = {
                    key: torch.mean(torch.cat([d[key] for d in eval_metrics])).to("cpu") for key in eval_metrics[0]
                }

                # compute metrics
                metrics_desc = ""
                # if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
                #     if accelerator.is_local_main_process:
                #         (
                #             metric_values,
                #             pred_descriptions,
                #             pred_prompts,
                #             audios,
                #             transcriptions,
                #             si_sdr_measures,
                #         ) = compute_metrics(
                #             eval_preds,
                #             eval_descriptions,
                #             eval_prompts,
                #             accelerator.device,
                #             training_args.compute_clap_similarity_metric,
                #             training_args.compute_noise_level_metric,
                #             training_args.noise_level_to_compute_clean_wer,
                #         )
                #         eval_metrics.update(metric_values)
                #         metrics_desc = " ".join([f"Eval {key}: {value} |" for key, value in metric_values.items()])
                #         if "wandb" in training_args.report_to:
                #             log_pred(
                #                 accelerator,
                #                 pred_descriptions,
                #                 pred_prompts,
                #                 transcriptions,
                #                 audios,
                #                 si_sdr_measures,
                #                 sampling_rate=sampling_rate,
                #                 step=cur_step,
                #                 prefix="eval",
                #             )
                #     accelerator.wait_for_everyone()

                # Print metrics and update progress bar
                if accelerator.is_local_main_process:
                    steps_trained_progress_bar.write(
                        f"Eval results for step ({cur_step} / {total_train_steps} | Eval Loss: {eval_metrics['loss']} |"
                        f" {metrics_desc})"
                    )

                log_metric(
                    accelerator,
                    metrics=eval_metrics,
                    train_time=eval_time,
                    step=cur_step,
                    epoch=epoch,
                    prefix="eval",
                )
                accelerator.wait_for_everyone()

                # release eval batch and relax metrics
                eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric = release_memory(
                    eval_metrics, eval_preds, eval_descriptions, eval_prompts, batch, eval_metric
                )
                if training_args.predict_with_generate and (cur_step % eval_generation_steps == 0 or cur_step == total_train_steps):
                    if config.use_text_encoder: 
                        generated_audios, input_ids, prompts = release_memory(generated_audios, input_ids, prompts)
                    else:
                        generated_audios, prompts = release_memory(generated_audios, prompts)

                accelerator.wait_for_everyone()
                # train mode
                model.train()

                # flush the train metrics
                train_start = time.time()

            # break condition
            if cur_step == total_train_steps:
                continue_training = False
                break

        if not continue_training:
            break

    accelerator.end_training()


if __name__ == "__main__":
    logger.info(f"PDSH_RCMD_TYPE: {os.getenv('PDSH_RCMD_TYPE')}")
    main()