run_parler_tts_training.py 54.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#!/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.

Yoach Lacombe's avatar
Yoach Lacombe committed
17
""" Train Parler-TTS using 🤗 Accelerate"""
18
19
20
21

import logging
import os
import re
22
import sys
Yoach Lacombe's avatar
Yoach Lacombe committed
23
import time
24
from multiprocess import set_start_method
25
from datetime import timedelta
26
27

from tqdm import tqdm
Yoach Lacombe's avatar
Yoach Lacombe committed
28
from pathlib import Path
29
30

import torch
31
32
33
34
35
from torch.utils.data import DataLoader

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

36
from huggingface_hub import HfApi
37
38

import transformers
Yoach Lacombe's avatar
Yoach Lacombe committed
39
from transformers import AutoFeatureExtractor, AutoTokenizer, HfArgumentParser
40
from transformers.trainer_pt_utils import LengthGroupedSampler
Yoach Lacombe's avatar
Yoach Lacombe committed
41
from transformers.optimization import get_scheduler
Yoach Lacombe's avatar
Yoach Lacombe committed
42
from transformers.utils import send_example_telemetry
43

44

45
from accelerate import Accelerator, skip_first_batches
46
47
from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin
from accelerate.utils.memory import release_memory
48

Yoach Lacombe's avatar
Yoach Lacombe committed
49
50
from parler_tts import (
    ParlerTTSConfig,
51
    ParlerTTSForConditionalGeneration,
Yoach Lacombe's avatar
Yoach Lacombe committed
52
53
    build_delay_pattern_mask,
)
54

55
56
57
58
59
60
61
62
63
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,
)
Dan Lyth's avatar
Dan Lyth committed
64
65
from training.arguments import ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments
from training.data import load_multiple_datasets, DataCollatorParlerTTSWithPadding, DataCollatorEncodecWithPadding
66
from training.eval import clap_similarity, wer, si_sdr
67

68
logger = logging.getLogger(__name__)
69

Yoach Lacombe's avatar
Yoach Lacombe committed
70

71
72
73
74
75
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.

Yoach Lacombe's avatar
Yoach Lacombe committed
76
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
77
78
79
80
81
82
83
84
85
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
Yoach Lacombe's avatar
Yoach Lacombe committed
86
    send_example_telemetry("run_parler_tts", model_args, data_args)
Yoach Lacombe's avatar
Yoach Lacombe committed
87

Yoach Lacombe's avatar
Yoach Lacombe committed
88
89
    if training_args.dtype == "float16":
        mixed_precision = "fp16"
90
        torch_dtype = torch.float16
Yoach Lacombe's avatar
Yoach Lacombe committed
91
92
    elif training_args.dtype == "bfloat16":
        mixed_precision = "bf16"
93
        torch_dtype = torch.bfloat16
Yoach Lacombe's avatar
Yoach Lacombe committed
94
95
    else:
        mixed_precision = "no"
96
        torch_dtype = torch.float32
Yoach Lacombe's avatar
Yoach Lacombe committed
97
98
99
100
101
102
103
104
105

    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`"
        )
106
107

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

109
    ####### A. Preparation
110
    kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=120))]
Yoach Lacombe's avatar
Yoach Lacombe committed
111

Yoach Lacombe's avatar
Yoach Lacombe committed
112
113
114
115
116
    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,
117
        kwargs_handlers=kwargs_handlers,
Yoach Lacombe's avatar
Yoach Lacombe committed
118
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

    accelerator.init_trackers(
        project_name=data_args.wandb_project,
        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,
        },
139
        init_kwargs={"wandb": {"name": data_args.wandb_run_name}} if data_args.wandb_run_name else {},
Yoach Lacombe's avatar
Yoach Lacombe committed
140
141
    )

Yoach Lacombe's avatar
Yoach Lacombe committed
142
    # Detecting last checkpoint and eventually continue from last checkpoint
143
144
145
146
147
148
149
150
    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."
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
151
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
152
153
154
155
156
157
158
159
160
161
162
            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)],
    )
163
    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
164

Yoach Lacombe's avatar
Yoach Lacombe committed
165
    # Log a small summary on each proces
166
167
168
169
    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}"
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
170
171
172
173

    # 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()
174
        transformers.utils.logging.set_verbosity_info()
Yoach Lacombe's avatar
Yoach Lacombe committed
175
176
177
178
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

179
180
181
182
    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed before initializing model.
    set_seed(training_args.seed)
183
    num_workers = data_args.preprocessing_num_workers
Yoach Lacombe's avatar
Yoach Lacombe committed
184

185
186
187
    # 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.
Yoach Lacombe's avatar
Yoach Lacombe committed
188

189
190
191
192
193
194
195
196
    # 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
Yoach Lacombe's avatar
Yoach Lacombe committed
197

198
199
200
201
202
203
204
    # 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,
205
        padding_side=model_args.prompt_padding_side,
206
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
207

208
209
210
211
212
213
214
215
    # 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,
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
216

217
    if model_args.use_fast_tokenizer:
Yoach Lacombe's avatar
Yoach Lacombe committed
218
219
220
        logger.warning(
            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
        )
221
222
        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
223

224
    # 2. Now, let's load the dataset
Yoach Lacombe's avatar
Yoach Lacombe committed
225

226
227
    if data_args.save_to_disk is not None:
        os.makedirs(data_args.save_to_disk, exist_ok=True)
Yoach Lacombe's avatar
Yoach Lacombe committed
228

229
230
231
    # 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:
232
233
        with accelerator.local_main_process_first():
            vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
Yoach Lacombe's avatar
Yoach Lacombe committed
234
    else:
235
236
237
238
        raw_datasets = DatasetDict()

        columns_to_keep = {
            "target_audio_column_name": data_args.target_audio_column_name,
Yoach Lacombe's avatar
Yoach Lacombe committed
239
            "prompt_column_name": data_args.prompt_column_name,
240
241
        }
        if data_args.description_column_name is not None:
242
            columns_to_keep["description_column_name"] = data_args.description_column_name
Yoach Lacombe's avatar
Yoach Lacombe committed
243

244
245
246
247
248
249
250
251
252
253
254
255
256
        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(),
257
                prompt_column_name=data_args.prompt_column_name,
258
259
                audio_column_name=data_args.target_audio_column_name,
                sampling_rate=sampling_rate,
260
                logger=logger,
261
262
                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
263

264
265
266
267
268
269
            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)}."
Yoach Lacombe's avatar
Yoach Lacombe committed
270
                    )
271
272
273
274
275
276
277
278

            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,
Yoach Lacombe's avatar
Yoach Lacombe committed
279
280
281
                data_args.eval_dataset_config_name
                if data_args.eval_dataset_config_name
                else data_args.train_dataset_config_name,
282
283
284
285
286
287
                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(),
288
289
290
                prompt_column_name=data_args.prompt_column_name,
                audio_column_name=data_args.target_audio_column_name,
                sampling_rate=sampling_rate,
291
                logger=logger,
292
293
                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
            )
294

295
            if data_args.max_eval_samples is not None:
296
297
298
299
                with accelerator.local_main_process_first():
                    raw_datasets["eval"] = (
                        raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
                    )
300

301
    # 3. Next, let's load the config.
Yoach Lacombe's avatar
Yoach Lacombe committed
302
    config = ParlerTTSConfig.from_pretrained(
303
304
305
306
307
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
308

309
    # update pad token id and decoder_start_token_id
310
311
312
313
314
315
316
    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
        }
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
317
318
    config.update(
        {
Yoach Lacombe's avatar
Yoach Lacombe committed
319
            "pad_token_id": model_args.pad_token_id if model_args.pad_token_id is not None else config.pad_token_id,
Yoach Lacombe's avatar
Yoach Lacombe committed
320
321
            "decoder_start_token_id": model_args.decoder_start_token_id
            if model_args.decoder_start_token_id is not None
322
            else config.decoder_start_token_id,
Yoach Lacombe's avatar
Yoach Lacombe committed
323
324
325
        }
    )

Yoach Lacombe's avatar
Yoach Lacombe committed
326
    # create model
Yoach Lacombe's avatar
Yoach Lacombe committed
327
    model = ParlerTTSForConditionalGeneration.from_pretrained(
328
329
330
331
332
        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,
333
        attn_implementation=model_args.attn_implementation,
334
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
335

336
337
338
    # enable gradient checkpointing if necessary
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
Yoach Lacombe's avatar
Yoach Lacombe committed
339

340
    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
341
342
343
    # 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`
Yoach Lacombe's avatar
Yoach Lacombe committed
344

345
    # derive max & min input length for sample rate & max duration
346
347
348
    sampling_rate = feature_extractor.sampling_rate
    max_target_length = data_args.max_duration_in_seconds * sampling_rate
    min_target_length = data_args.min_duration_in_seconds * sampling_rate
349
350
351
352
    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]
Yoach Lacombe's avatar
Yoach Lacombe committed
353
354
    audio_encoder_pad_token_id = config.decoder.pad_token_id
    audio_encoder_eos_token_id = config.decoder.eos_token_id
Yoach Lacombe's avatar
Yoach Lacombe committed
355
356
357
    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
Yoach Lacombe's avatar
Yoach Lacombe committed
358
    bandwidth = model_args.bandwidth
359
    attn_implementation = model_args.attn_implementation
Yoach Lacombe's avatar
Yoach Lacombe committed
360

361
362
    # Freeze Encoders
    model.freeze_encoders(model_args.freeze_text_encoder)
Yoach Lacombe's avatar
Yoach Lacombe committed
363

364
    # Test all gather - used for warmout and avoiding timeout
365
    logger.debug(str(accelerator.process_index), main_process_only=False, in_order=True)
366
367
368
369
    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()
Yoach Lacombe's avatar
Yoach Lacombe committed
370
371

    if not dataset_was_precomputed:
372
        # Filter on text length
373
        if description_column_name is not None and data_args.max_text_length is not None:
374
            with accelerator.local_main_process_first():
375
376
377
378
379
380
                # 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],
                )
381

382
383
384
385
        # Preprocessing the dataset.
        # We need to tokenize the texts.
        def pass_through_processors(description, prompt):
            batch = {}
Yoach Lacombe's avatar
Yoach Lacombe committed
386

387
388
            batch["input_ids"] = description_tokenizer(description.strip())["input_ids"]
            batch["prompt_input_ids"] = prompt_tokenizer(prompt.strip())["input_ids"]
389
390

            return batch
Yoach Lacombe's avatar
Yoach Lacombe committed
391

392
        with accelerator.local_main_process_first():
393
            # this is a trick to avoid to rewrite the entire audio column which takes ages
394
            vectorized_datasets = raw_datasets.map(
395
396
                pass_through_processors,
                remove_columns=next(iter(raw_datasets.values())).column_names,
397
                input_columns=[description_column_name, prompt_column_name],
398
399
400
                num_proc=num_workers,
                desc="preprocess datasets",
            )
401

402
        # We use Accelerate to perform distributed inference
403
        # T5 doesn't support fp16
Yoach Lacombe's avatar
Yoach Lacombe committed
404
        autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
405
406

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

409
        logger.info("*** Encode target audio with encodec ***")
Yoach Lacombe's avatar
Yoach Lacombe committed
410

411
412
        # 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
413
414
415
416
        if training_args.torch_compile:
            audio_decoder = accelerator.prepare_model(model.audio_encoder, evaluation_mode=True)
        else:
            audio_decoder = model.audio_encoder
417

Yoach Lacombe's avatar
Yoach Lacombe committed
418
419
420
421
422
423
424
        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,
        )
425
426
427
428
429
430
431
432
433

        def apply_audio_decoder(batch):
            len_audio = batch.pop("len_audio")
            audio_decoder.to(batch["input_values"].device).eval()
            with torch.no_grad():
                labels = audio_decoder.encode(**batch, bandwidth=bandwidth)["audio_codes"]
            output = {}
            output["len_audio"] = len_audio
            # (1, bsz, codebooks, seq_len) -> (bsz, seq_len, codebooks)
Yoach Lacombe's avatar
Yoach Lacombe committed
434
            output["labels"] = labels.squeeze(0).transpose(1, 2)
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469

            # 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:]}
Yoach Lacombe's avatar
Yoach Lacombe committed
470
            return output
471

472
473
        for split in vectorized_datasets:
            data_loader = DataLoader(
474
                raw_datasets[split],
Yoach Lacombe's avatar
Yoach Lacombe committed
475
                batch_size=training_args.audio_encoder_per_device_batch_size,
476
477
478
                collate_fn=encoder_data_collator,
                num_workers=training_args.dataloader_num_workers,
                pin_memory=True,
479
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
480
            data_loader = accelerator.prepare(data_loader)
481
482
483
484
485
486
487
488
489
            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)
Yoach Lacombe's avatar
Yoach Lacombe committed
490

491
492
            all_generated_labels = []
            all_lens = []
493
494
495
496
497
498
            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)
Yoach Lacombe's avatar
Yoach Lacombe committed
499

500
501
502
503
504
                    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)]
Yoach Lacombe's avatar
Yoach Lacombe committed
505

506
507
                        all_generated_labels.extend(lab)
                        all_lens.extend(lens)
Yoach Lacombe's avatar
Yoach Lacombe committed
508

509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
                        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 = []
Yoach Lacombe's avatar
Yoach Lacombe committed
524

525
                accelerator.wait_for_everyone()
Yoach Lacombe's avatar
Yoach Lacombe committed
526

527
528
529
            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(
530
                    postprocess_dataset,
Yoach Lacombe's avatar
Yoach Lacombe committed
531
                    num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
532
                    input_columns=["labels"],
533
534
                    desc="Postprocessing labeling",
                )
535
536
537
538
539
540
541
542
543
544
545
546
547
548
                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.local_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)
549
550

        accelerator.free_memory()
551
        del generate_labels, all_lens
552

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

558
559
560
561
562
563
564
565
566
            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"],
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
567

568
569
570
571
572
573
574
575
        if description_column_name is not None and data_args.max_description_token_length is not None:
            with accelerator.local_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"],
                )
576

577
578
579
580
581
582
583
584
        if data_args.max_prompt_token_length is not None:
            with accelerator.local_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"],
                )
Yoach Lacombe's avatar
Yoach Lacombe committed
585

586
    if data_args.save_to_disk is not None and not dataset_was_precomputed:
587
        if accelerator.is_main_process:
Yoach Lacombe's avatar
Yoach Lacombe committed
588
589
590
591
            vectorized_datasets.save_to_disk(
                data_args.save_to_disk,
                num_proc=min(data_args.preprocessing_num_workers, len(vectorized_datasets["eval"]) - 1),
            )
592
        accelerator.wait_for_everyone()
593
        logger.info(f"Dataset saved at {data_args.save_to_disk}")
Yoach Lacombe's avatar
Yoach Lacombe committed
594

595
    audio_max_length = None
596
    if padding == "max_length":
597
        audio_max_length = max(vectorized_datasets["train"]["target_length"])
598
        with accelerator.local_main_process_first():
599
            max_sample = vectorized_datasets["train"].filter(
Yoach Lacombe's avatar
Yoach Lacombe committed
600
601
602
603
                lambda x: x == audio_max_length,
                num_proc=num_workers,
                input_columns=["target_length"],
            )
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
        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.local_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.local_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"],
            )
623

624
625
626
627
628
    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)}

629
        with accelerator.local_main_process_first():
630
631
632
633
634
635
            vectorized_datasets = vectorized_datasets.map(
                add_target_lengths,
                num_proc=num_workers,
                input_columns=["target_length", "prompt_input_ids", "input_ids"],
            )

636
637
638
639
640
    # 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
641
    if data_args.preprocessing_only and data_args.save_to_disk is None:
Yoach Lacombe's avatar
Yoach Lacombe committed
642
643
644
        raise ValueError(
            "`preprocessing_only=True` but `save_to_disk` is not set. The latter should indicates where to save the dataset locally."
        )
645
646
    elif data_args.preprocessing_only:
        logger.info(f"Data preprocessing finished. Files save at {data_args.save_to_disk}")
647
        return
Yoach Lacombe's avatar
Yoach Lacombe committed
648

649
    # 6. Next, we can prepare the training.
Yoach Lacombe's avatar
Yoach Lacombe committed
650

Yoach Lacombe's avatar
Yoach Lacombe committed
651
    # Let's use word CLAP similary and WER metrics as our evaluation metrics,
652
653
654
655
656
657
658
659
660
    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,
    ):
661
        results = {}
Yoach Lacombe's avatar
Yoach Lacombe committed
662
        input_ids = descriptions
663
        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
Yoach Lacombe's avatar
Yoach Lacombe committed
664
        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
665
        audios = [a.float().cpu().numpy() for a in audios]
Yoach Lacombe's avatar
Yoach Lacombe committed
666

667
668
669
670
671
672
673
674
675
        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)
676

677
        word_error, transcriptions, clean_word_error, noisy_word_error, percent_clean_samples = wer(
Yoach Lacombe's avatar
Yoach Lacombe committed
678
679
680
681
682
683
            model_args.asr_model_name_or_path,
            prompts,
            audios,
            device,
            training_args.per_device_eval_batch_size,
            sampling_rate,
684
685
            noise_level_to_compute_clean_wer,
            si_sdr_measures,
Yoach Lacombe's avatar
Yoach Lacombe committed
686
        )
Yoach Lacombe's avatar
Yoach Lacombe committed
687
        results["wer"] = word_error
688
689
690
691
        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
692

693
        return results, texts, prompts, audios, transcriptions, si_sdr_measures
Yoach Lacombe's avatar
Yoach Lacombe committed
694

Yoach Lacombe's avatar
Yoach Lacombe committed
695
696
697
698
699
700
    # Define Training Schedule
    # Store some constants
    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
    train_batch_size = per_device_train_batch_size * accelerator.num_processes
    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
Yoach Lacombe's avatar
Yoach Lacombe committed
701

Yoach Lacombe's avatar
Yoach Lacombe committed
702
703
704
705
706
707
708
709
710
711
712
713
    if training_args.max_steps < 0:
        num_epochs = int(training_args.num_train_epochs)
        steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps)
        total_train_steps = steps_per_epoch * num_epochs
    elif training_args.max_steps > 0:
        logger.info("max_steps is given, it will override any value given in num_train_epochs")
        total_train_steps = int(training_args.max_steps)
        # 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:
Yoach Lacombe's avatar
Yoach Lacombe committed
714
        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
Yoach Lacombe's avatar
Yoach Lacombe committed
715
716
717
        eval_steps = steps_per_epoch
    else:
        eval_steps = training_args.eval_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
718

719
    # T5 doesn't support fp16
Yoach Lacombe's avatar
Yoach Lacombe committed
720
721
    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))

Yoach Lacombe's avatar
Yoach Lacombe committed
722
723
724
725
726
727
    # 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,
728
        weight_decay=training_args.weight_decay,
Yoach Lacombe's avatar
Yoach Lacombe committed
729
    )
730

Yoach Lacombe's avatar
Yoach Lacombe committed
731
732
733
734
    # 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,
Yoach Lacombe's avatar
Yoach Lacombe committed
735
        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
Yoach Lacombe's avatar
Yoach Lacombe committed
736
737
        num_training_steps=total_train_steps * accelerator.num_processes,
    )
738
739

    # Instantiate custom data collator
Yoach Lacombe's avatar
Yoach Lacombe committed
740
    data_collator = DataCollatorParlerTTSWithPadding(
Yoach Lacombe's avatar
Yoach Lacombe committed
741
742
743
744
745
746
747
        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,
748
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
749

Yoach Lacombe's avatar
Yoach Lacombe committed
750
751
    # Prepare everything with accelerate
    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
Yoach Lacombe's avatar
Yoach Lacombe committed
752

Yoach Lacombe's avatar
Yoach Lacombe committed
753
754
    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {total_train_steps * train_batch_size * gradient_accumulation_steps}")
755
    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
Yoach Lacombe's avatar
Yoach Lacombe committed
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
    logger.info(
        f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
    )
    logger.info(f"  Total optimization steps = {total_train_steps}")

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

    checkpoint = None
    if training_args.resume_from_checkpoint is not None:
        checkpoint = training_args.resume_from_checkpoint
    elif last_checkpoint is not None:
        checkpoint = last_checkpoint
Yoach Lacombe's avatar
Yoach Lacombe committed
777

Yoach Lacombe's avatar
Yoach Lacombe committed
778
779
    if accelerator.is_main_process:
        if training_args.push_to_hub:
780
781
782
            api = HfApi(token=training_args.hub_token)

            # Create repo (repo_name from args or inferred)
Yoach Lacombe's avatar
Yoach Lacombe committed
783
784
785
            repo_name = training_args.hub_model_id
            if repo_name is None:
                repo_name = Path(training_args.output_dir).absolute().name
786
            repo_id = api.create_repo(repo_name, exist_ok=True).repo_id
Yoach Lacombe's avatar
Yoach Lacombe committed
787
788
789
790
791
792
793

            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()
Yoach Lacombe's avatar
Yoach Lacombe committed
794

Yoach Lacombe's avatar
Yoach Lacombe committed
795
796
    # Now save everything to be able to create a single processor later
    # make sure all processes wait until data is saved
797
798
799
800
801
802
803
804
805
806
807
808
809
810
    # 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)
Yoach Lacombe's avatar
Yoach Lacombe committed
811

812
813
814
        feature_extractor.save_pretrained(training_args.output_dir)
        config.save_pretrained(training_args.output_dir)
    accelerator.wait_for_everyone()
Yoach Lacombe's avatar
Yoach Lacombe committed
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830

    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):
831
832
            with accelerator.local_main_process_first():
                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
Yoach Lacombe's avatar
Yoach Lacombe committed
833

Yoach Lacombe's avatar
Yoach Lacombe committed
834
835
        if training_args.max_steps < 0:
            # we know exactly the number of steps per epoch, so can skip through the required number of batches
836
            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
837
838
839
840
841
        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
842
843
            with accelerator.local_main_process_first():
                vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
Yoach Lacombe's avatar
Yoach Lacombe committed
844
845
    else:
        resume_step = None
Yoach Lacombe's avatar
Yoach Lacombe committed
846

Yoach Lacombe's avatar
Yoach Lacombe committed
847
848
    gen_kwargs = {
        "do_sample": model_args.do_sample,
yoach@huggingface.co's avatar
yoach@huggingface.co committed
849
        "temperature": model_args.temperature,
Yoach Lacombe's avatar
Yoach Lacombe committed
850
        "max_length": model_args.max_length,
851
852
853
854
        # 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,
Yoach Lacombe's avatar
Yoach Lacombe committed
855
    }
Yoach Lacombe's avatar
Yoach Lacombe committed
856

Yoach Lacombe's avatar
Yoach Lacombe committed
857
858
859
    # Define gradient update step fn
    def train_step(
        batch,
860
861
        accelerator,
        autocast_kwargs,
Yoach Lacombe's avatar
Yoach Lacombe committed
862
    ):
863
        if mixed_precision == "fp16":
864
865
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
866
                if training_args.parallel_mode.value != "distributed":
Yoach Lacombe's avatar
Yoach Lacombe committed
867
868
869
                    encoder_outputs = model.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
870
                else:
Yoach Lacombe's avatar
Yoach Lacombe committed
871
872
873
                    encoder_outputs = model.module.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
                # 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
890
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
891

Yoach Lacombe's avatar
Yoach Lacombe committed
892
893
894
895
896
897
        outputs = model(**batch)
        # CE (data) loss
        ce_loss = outputs.loss

        metrics = {"loss": ce_loss}
        return ce_loss, metrics
Yoach Lacombe's avatar
Yoach Lacombe committed
898

Yoach Lacombe's avatar
Yoach Lacombe committed
899
    # Define eval fn
Yoach Lacombe's avatar
Yoach Lacombe committed
900
901
902
903
904
    def eval_step(
        batch,
        accelerator,
        autocast_kwargs,
    ):
Yoach Lacombe's avatar
Yoach Lacombe committed
905
906
        eval_model = model if not training_args.torch_compile else model._orig_mod

907
        if mixed_precision == "fp16":
908
909
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
                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
934
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
935
936

        with torch.no_grad():
Yoach Lacombe's avatar
Yoach Lacombe committed
937
            outputs = eval_model(**batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
938
939
940
941
942
        # CE (data) loss
        ce_loss = outputs.loss
        metrics = {"loss": ce_loss}
        return metrics

943
    def generate_step(batch, accelerator):
944
        batch.pop("decoder_attention_mask", None)
945
        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=True)
Yoach Lacombe's avatar
Yoach Lacombe committed
946
        if training_args.torch_compile:
947
            # if the model is compiled, we use the original model bc compile is not compatible with .generate
Yoach Lacombe's avatar
Yoach Lacombe committed
948
949
            eval_model = model._orig_mod

950
951
        # 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"))):
Yoach Lacombe's avatar
Yoach Lacombe committed
952
        output_audios = eval_model.generate(**batch, **gen_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
953
954
955
        output_audios = accelerator.pad_across_processes(output_audios, dim=1, pad_index=0)
        return output_audios

956
957
    model.train()

Yoach Lacombe's avatar
Yoach Lacombe committed
958
    for epoch in range(epochs_trained, num_epochs):
959
960
        with accelerator.local_main_process_first():
            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
Yoach Lacombe's avatar
Yoach Lacombe committed
961
962
963
        sampler = None
        if training_args.group_by_length:
            sampler = LengthGroupedSampler(train_batch_size, lengths=vectorized_datasets["train"]["target_length"])
Yoach Lacombe's avatar
Yoach Lacombe committed
964
965
966
967
        train_dataloader = DataLoader(
            vectorized_datasets["train"],
            collate_fn=data_collator,
            batch_size=per_device_train_batch_size,
Yoach Lacombe's avatar
Yoach Lacombe committed
968
            sampler=sampler,
Yoach Lacombe's avatar
Yoach Lacombe committed
969
970
971
972
973
974
975
976
977
            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
978
            logger.info(f"  Skip first {resume_step} batches")
Yoach Lacombe's avatar
Yoach Lacombe committed
979
980
            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
            resume_step = None
981
            accelerator.wait_for_everyone()
Yoach Lacombe's avatar
Yoach Lacombe committed
982
983
984

        for batch in train_dataloader:
            with accelerator.accumulate(model):
985
                loss, train_metric = train_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), training_args.max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            # Check if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                steps_trained_progress_bar.update(1)
                cur_step += 1

                if cur_step % training_args.logging_steps == 0:
                    steps_trained_progress_bar.write(
                        f"Step... ({cur_step} / {total_train_steps} | Loss:"
                        f" {train_metric['loss']}, Learning Rate:"
                        f" {lr_scheduler.get_last_lr()[0]})"
                    )
                    log_metric(
                        accelerator,
                        metrics=train_metric,
                        learning_rate=lr_scheduler.get_last_lr()[0],
                        train_time=train_time + time.time() - train_start,
                        step=cur_step,
                        epoch=epoch,
                        prefix="train",
                    )

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

                        if cur_step == total_train_steps:
                            # un-wrap student model for save
Yoach Lacombe's avatar
Yoach Lacombe committed
1028
1029
                            unwrapped_model = accelerator.unwrap_model(model)
                            unwrapped_model.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
1030
1031

                        if training_args.push_to_hub:
1032
1033
1034
                            api.upload_folder(
                                repo_id=repo_id,
                                folder_path=training_args.output_dir,
Yoach Lacombe's avatar
Yoach Lacombe committed
1035
                                commit_message=f"Saving train state of step {cur_step}",
1036
                                run_as_future=True,
Yoach Lacombe's avatar
Yoach Lacombe committed
1037
                            )
1038
                    accelerator.wait_for_everyone()
Yoach Lacombe's avatar
Yoach Lacombe committed
1039
1040
1041
1042

                if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
                    train_time += time.time() - train_start
                    # ======================== Evaluating ==============================
1043
                    model.eval()
Yoach Lacombe's avatar
Yoach Lacombe committed
1044
1045
1046
1047
1048
                    eval_metrics = []
                    eval_preds = []
                    eval_descriptions = []
                    eval_prompts = []
                    eval_start = time.time()
Yoach Lacombe's avatar
Yoach Lacombe committed
1049

Yoach Lacombe's avatar
Yoach Lacombe committed
1050
1051
                    # release training input batch
                    batch = release_memory(batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
1052
1053
1054
1055
1056

                    validation_dataloader = DataLoader(
                        vectorized_datasets["eval"],
                        collate_fn=data_collator,
                        batch_size=per_device_eval_batch_size,
1057
                        drop_last=False,
1058
                        num_workers=training_args.eval_dataloader_num_workers,
Yoach Lacombe's avatar
Yoach Lacombe committed
1059
1060
1061
1062
1063
1064
                        pin_memory=training_args.dataloader_pin_memory,
                    )
                    validation_dataloader = accelerator.prepare(validation_dataloader)

                    for batch in tqdm(
                        validation_dataloader,
1065
                        desc=f"Evaluating - Inference ...",
Yoach Lacombe's avatar
Yoach Lacombe committed
1066
1067
1068
1069
                        position=2,
                        disable=not accelerator.is_local_main_process,
                    ):
                        # Model forward
1070
                        eval_metric = eval_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
1071
1072
1073
                        eval_metric = accelerator.gather_for_metrics(eval_metric)
                        eval_metrics.append(eval_metric)

1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
                    if training_args.predict_with_generate:
                        validation_dataloader = DataLoader(
                            vectorized_datasets["eval"],
                            collate_fn=data_collator,
                            batch_size=per_device_eval_batch_size,
                            drop_last=False,
                            num_workers=training_args.dataloader_pin_memory,
                            pin_memory=training_args.dataloader_pin_memory,
                        )
                        validation_dataloader = accelerator.prepare(validation_dataloader)
Yoach Lacombe's avatar
Yoach Lacombe committed
1084
                        # generation
1085
                        for batch in tqdm(
Yoach Lacombe's avatar
Yoach Lacombe committed
1086
1087
1088
1089
1090
                            validation_dataloader,
                            desc=f"Evaluating - Generation ...",
                            position=2,
                            disable=not accelerator.is_local_main_process,
                        ):
1091
                            generated_audios = generate_step(batch, accelerator)
Yoach Lacombe's avatar
Yoach Lacombe committed
1092
                            # Gather all predictions and targets
Yoach Lacombe's avatar
Yoach Lacombe committed
1093
1094
1095
1096
1097
1098
                            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)
                            )
1099
1100
1101
                            eval_preds.extend(generated_audios.to("cpu"))
                            eval_descriptions.extend(input_ids.to("cpu"))
                            eval_prompts.extend(prompts.to("cpu"))
Yoach Lacombe's avatar
Yoach Lacombe committed
1102
1103
1104
1105

                    eval_time = time.time() - eval_start
                    # normalize eval metrics
                    eval_metrics = {
1106
                        key: torch.mean(torch.cat([d[key] for d in eval_metrics])).to("cpu") for key in eval_metrics[0]
Yoach Lacombe's avatar
Yoach Lacombe committed
1107
1108
1109
1110
1111
                    }

                    # compute metrics
                    metrics_desc = ""
                    if training_args.predict_with_generate:
1112
1113
1114
                        if accelerator.is_local_main_process:
                            (
                                metric_values,
1115
1116
1117
                                pred_descriptions,
                                pred_prompts,
                                audios,
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
                                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,
1128
                            )
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
                            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()
Yoach Lacombe's avatar
Yoach Lacombe committed
1144
1145

                    # Print metrics and update progress bar
1146
1147
1148
1149
1150
                    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})"
                        )
Yoach Lacombe's avatar
Yoach Lacombe committed
1151
1152
1153
1154
1155
1156
1157
1158
1159

                    log_metric(
                        accelerator,
                        metrics=eval_metrics,
                        train_time=eval_time,
                        step=cur_step,
                        epoch=epoch,
                        prefix="eval",
                    )
Yoach Lacombe's avatar
Yoach Lacombe committed
1160

1161
                    # release eval batch and relax metrics
1162
1163
1164
1165
1166
1167
1168
1169
                    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:
                        generated_audios, input_ids, prompts = release_memory(generated_audios, input_ids, prompts)

                    # train mode
                    model.train()
1170

Yoach Lacombe's avatar
Yoach Lacombe committed
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
                    # 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()
1183
1184
1185


if __name__ == "__main__":
Yoach Lacombe's avatar
Yoach Lacombe committed
1186
    main()