run_parler_tts_training.py 78.2 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
22

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

28

Yoach Lacombe's avatar
Yoach Lacombe committed
29
import evaluate
30
from tqdm import tqdm
Yoach Lacombe's avatar
Yoach Lacombe committed
31
from pathlib import Path
32
from dataclasses import dataclass, field
Yoach Lacombe's avatar
Yoach Lacombe committed
33
from typing import Dict, List, Optional, Union, Set
34
35
36
37

import datasets
import numpy as np
import torch
38
39
from torch.utils.data import DataLoader

40
41
from datasets import DatasetDict, load_dataset, Dataset, IterableDataset, interleave_datasets, concatenate_datasets

Yoach Lacombe's avatar
Yoach Lacombe committed
42
from huggingface_hub import Repository, create_repo
43
44
45
46
47
48
49
50
51
import transformers
from transformers import (
    AutoFeatureExtractor,
    AutoModel,
    AutoProcessor,
    AutoTokenizer,
    HfArgumentParser,
    Seq2SeqTrainingArguments,
)
Yoach Lacombe's avatar
Yoach Lacombe committed
52
from transformers.trainer_pt_utils import LengthGroupedSampler
Yoach Lacombe's avatar
Yoach Lacombe committed
53
54
from transformers import pipeline
from transformers.optimization import get_scheduler
55
56
57
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version
from transformers.integrations import is_wandb_available
Yoach Lacombe's avatar
Yoach Lacombe committed
58
from transformers import AutoModel
Yoach Lacombe's avatar
add DAC  
Yoach Lacombe committed
59

60
61

from accelerate import Accelerator
62
from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin
Yoach Lacombe's avatar
Yoach Lacombe committed
63
from accelerate.utils.memory import release_memory
64

Yoach Lacombe's avatar
Yoach Lacombe committed
65
66
67
68
69
from parler_tts import (
    ParlerTTSForConditionalGeneration,
    ParlerTTSConfig,
    build_delay_pattern_mask,
)
70

Yoach Lacombe's avatar
Yoach Lacombe committed
71
72
if is_wandb_available():
    from wandb import Audio
73
74
75
76
77
78
79
80
81
82
83
84
85

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.38.0.dev0")

require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt")


logger = logging.getLogger(__name__)


def list_field(default=None, metadata=None):
    return field(default_factory=lambda: default, metadata=metadata)

Yoach Lacombe's avatar
Yoach Lacombe committed
86

Yoach Lacombe's avatar
Yoach Lacombe committed
87
88
_RE_CHECKPOINT = re.compile(r"^checkpoint-(\d+)-epoch-(\d+)$")

Yoach Lacombe's avatar
Yoach Lacombe committed
89

Yoach Lacombe's avatar
Yoach Lacombe committed
90
91
92
93
94
95
96
97
98
99
100
def get_last_checkpoint(folder):
    content = os.listdir(folder)
    checkpoints = [
        path
        for path in content
        if _RE_CHECKPOINT.search(path) is not None and os.path.isdir(os.path.join(folder, path))
    ]
    if len(checkpoints) == 0:
        return
    return os.path.join(folder, max(checkpoints, key=lambda x: int(_RE_CHECKPOINT.search(x).groups()[0])))

Yoach Lacombe's avatar
Yoach Lacombe committed
101

Yoach Lacombe's avatar
Yoach Lacombe committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def sorted_checkpoints(output_dir=None, checkpoint_prefix="checkpoint") -> List[str]:
    """Helper function to sort saved checkpoints from oldest to newest."""
    ordering_and_checkpoint_path = []

    glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]

    for path in glob_checkpoints:
        regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
        if regex_match is not None and regex_match.groups() is not None:
            ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))

    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
    checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
    return checkpoints_sorted

Yoach Lacombe's avatar
Yoach Lacombe committed
117

Yoach Lacombe's avatar
Yoach Lacombe committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def rotate_checkpoints(save_total_limit=None, output_dir=None, checkpoint_prefix="checkpoint") -> None:
    """Helper function to delete old checkpoints."""
    if save_total_limit is None or save_total_limit <= 0:
        return
    # Check if we should delete older checkpoint(s)
    checkpoints_sorted = sorted_checkpoints(output_dir=output_dir, checkpoint_prefix=checkpoint_prefix)
    if len(checkpoints_sorted) <= save_total_limit:
        return

    number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
    checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
    for checkpoint in checkpoints_to_be_deleted:
        logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
        shutil.rmtree(checkpoint, ignore_errors=True)

Yoach Lacombe's avatar
Yoach Lacombe committed
133

Yoach Lacombe's avatar
Yoach Lacombe committed
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def log_metric(
    accelerator,
    metrics: Dict,
    train_time: float,
    step: int,
    epoch: int,
    learning_rate: float = None,
    prefix: str = "train",
):
    """Helper function to log all training/evaluation metrics with the correct prefixes and styling."""
    log_metrics = {}
    for k, v in metrics.items():
        log_metrics[f"{prefix}/{k}"] = v
    log_metrics[f"{prefix}/time"] = train_time
    log_metrics[f"{prefix}/epoch"] = epoch
    if learning_rate is not None:
        log_metrics[f"{prefix}/learning_rate"] = learning_rate
    accelerator.log(log_metrics, step=step)

Yoach Lacombe's avatar
Yoach Lacombe committed
153

Yoach Lacombe's avatar
Yoach Lacombe committed
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def log_pred(
    accelerator,
    pred_descriptions: List[str],
    pred_prompts: List[str],
    transcriptions: List[str],
    audios: List[torch.Tensor],
    sampling_rate: int,
    step: int,
    prefix: str = "eval",
    num_lines: int = 200000,
):
    """Helper function to log target/predicted transcriptions to weights and biases (wandb)."""
    if accelerator.is_main_process:
        wandb_tracker = accelerator.get_tracker("wandb")
        # pretty name for current step: step 50000 -> step 50k
        cur_step_pretty = f"{int(step // 1000)}k" if step > 1000 else step
        prefix_pretty = prefix.replace("/", "-")

        # convert str data to a wandb compatible format
        str_data = [[pred_descriptions[i], pred_prompts[i], transcriptions[i]] for i in range(len(pred_descriptions))]
        # log as a table with the appropriate headers
        wandb_tracker.log_table(
            table_name=f"predictions/{prefix_pretty}-step-{cur_step_pretty}",
            columns=["Target descriptions", "Target prompts", "Predicted transcriptions"],
            data=str_data[:num_lines],
            step=step,
            commit=False,
        )
Yoach Lacombe's avatar
Yoach Lacombe committed
182

Yoach Lacombe's avatar
Yoach Lacombe committed
183
        # wandb can only loads 100 audios per step
Yoach Lacombe's avatar
Yoach Lacombe committed
184
185
        wandb_tracker.log(
            {
Yoach Lacombe's avatar
Yoach Lacombe committed
186
187
188
189
190
191
                "Speech samples": [
                    Audio(
                        audio,
                        caption=f"{pred_prompts[i]} --- DESCRIPTION: {pred_descriptions[i]}",
                        sample_rate=sampling_rate,
                    )
Yoach Lacombe's avatar
Yoach Lacombe committed
192
193
194
195
196
197
                    for (i, audio) in enumerate(audios[: min(len(audios), 100)])
                ]
            },
            step=step,
        )

198
199
200
201
202
203

@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """
Yoach Lacombe's avatar
Yoach Lacombe committed
204

205
206
207
208
209
210
211
212
213
214
215
216
217
218
    # TODO: pretrain from scratch
    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    feature_extractor_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained feature extractor name or path if not the same as model_name"}
    )
    description_tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained description tokenizer name or path if not the same as model_name"}
    )
    prompt_tokenizer_name: Optional[str] = field(
Yoach Lacombe's avatar
Yoach Lacombe committed
219
220
        default=None,
        metadata={"help": "Pretrained prompt tokenizer name or path if not the same as description_tokenizer_name"},
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    pad_token_id: int = field(
        default=None,
        metadata={"help": "If specified, change the model pad token id."},
    )
    decoder_start_token_id: int = field(
        default=None,
        metadata={"help": "If specified, change the model decoder start token id."},
    )
    freeze_text_encoder: bool = field(
        default=False,
        metadata={"help": "Whether to freeze the text encoder."},
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
246
247
248
249
    do_sample: bool = field(
        default=False,
        metadata={"help": "Whether to do sampling or greedy decoding."},
    )
yoach@huggingface.co's avatar
yoach@huggingface.co committed
250
251
252
253
    temperature: float = field(
        default=0.4,
        metadata={"help": "Temperature if sampling."},
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
254
    max_length: int = field(
255
256
        default=2580,
        metadata={"help": "Generation max length."},
Yoach Lacombe's avatar
Yoach Lacombe committed
257
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
258
    bandwidth: float = field(
Yoach Lacombe's avatar
Yoach Lacombe committed
259
        default=6,  # TODO
Yoach Lacombe's avatar
Yoach Lacombe committed
260
261
        metadata={"help": "Audio encoder bandwidth."},
    )
262
263
264


@dataclass
Yoach Lacombe's avatar
Yoach Lacombe committed
265
class DataTrainingArguments:
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
    """
    Arguments pertaining to what data we are going to input our model for training and eval.

    Using `HfArgumentParser` we can turn this class
    into argparse arguments to be able to specify them on
    the command line.
    """

    train_dataset_name: str = field(
        default=None,
        metadata={
            "help": "The name of the training dataset to use (via the datasets library). Load and combine "
            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
        },
    )
    train_dataset_config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "The configuration name of the training dataset to use (via the datasets library). Load and combine "
            "multiple datasets by separating dataset configs by a '+' symbol."
        },
    )
    train_split_name: str = field(
        default="train",
        metadata={
            "help": ("The name of the training data set split to use (via the datasets library). Defaults to 'train'")
        },
    )
    train_dataset_samples: str = field(
        default=None,
        metadata={
            "help": "Number of samples in the training data. Load and combine "
            "multiple datasets by separating dataset samples by a '+' symbol."
        },
    )
    train_metadata_dataset_name: str = field(
        default=None,
        metadata={
            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
        },
    )
    eval_dataset_name: str = field(
        default=None,
        metadata={
            "help": "The name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset name if unspecified."
        },
    )
    eval_dataset_config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "The configuration name of the evaluation dataset to use (via the datasets library). Defaults to the training dataset config name if unspecified"
        },
    )
    eval_split_name: str = field(
        default="test",
        metadata={
            "help": "The name of the evaluation data set split to use (via the datasets library). Defaults to 'test'"
        },
    )
    eval_metadata_dataset_name: str = field(
        default=None,
        metadata={
            "help": "The name of the metadata training dataset to use (via the datasets library). Load and combine "
            "multiple datasets by separating dataset ids by a '+' symbol. For example, to load and combine "
            " librispeech and common voice, set `train_dataset_name='librispeech_asr+common_voice'`."
        },
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
336
    target_audio_column_name: str = field(  # TODO
337
338
339
        default="audio",
        metadata={"help": "The name of the dataset column containing the target audio data. Defaults to 'audio'"},
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
340
    description_column_name: str = field(  # TODO
341
342
343
        default=None,
        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'None'."},
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
344
    prompt_column_name: str = field(  # TODO
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
        default=None,
        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'None'."},
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of validation examples to this "
                "value if set."
            )
        },
    )
    max_duration_in_seconds: float = field(
        default=35.0,
        metadata={
            "help": (
377
378
                "Filter audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_seconds`."
                "Also, used to set maximum audio length if `pad_to_max_length=True`."
379
380
381
382
383
384
            )
        },
    )
    min_duration_in_seconds: float = field(
        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
    )
385
    max_text_length: int = field(
386
387
388
        default=500, metadata={"help": "If set, max description lengths in number of characters."}
    )
    max_prompt_token_length: int = field(
Yoach Lacombe's avatar
Yoach Lacombe committed
389
390
        default=None,
        metadata={
391
392
393
394
            "help": (
                "If set, filter samples with prompts that are longer than `max_prompt_token_length` tokens."
                "Also, used to set maximum prompt token length if `pad_to_max_length=True`."
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
395
        },
396
397
    )
    max_description_token_length: int = field(
Yoach Lacombe's avatar
Yoach Lacombe committed
398
399
        default=None,
        metadata={
400
401
402
403
            "help": (
                "If set, filter samples with descriptions that are longer than `max_description_token_length` tokens."
                "Also, used to set maximum desription token length if `pad_to_max_length=True`."
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
404
        },
405
406
    )
    pad_to_max_length: bool = field(
Yoach Lacombe's avatar
Yoach Lacombe committed
407
408
409
410
411
412
413
        default=False,
        metadata={
            "help": (
                "If `True`, pad audio, prompt and description to a maximum length set with respectively "
                "`max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`."
            )
        },
414
    )
415
416
417
418
419
420
421
    preprocessing_only: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to only do data preprocessing and skip training. This is especially useful when data"
                " preprocessing errors out in distributed training due to timeout. In this case, one should run the"
                " preprocessing in a non-distributed setup with `preprocessing_only=True` so that the cached datasets"
422
423
                " can consequently be loaded in distributed training."
                " In this training script, `save_to_disk` must be set to the path in which the dataset should be saved. "
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
            )
        },
    )
    token: str = field(
        default=None,
        metadata={
            "help": (
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
            )
        },
    )
    use_auth_token: bool = field(
        default=None,
        metadata={
            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
        },
    )
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
                "execute code present on the Hub on your local machine."
            )
        },
    )
    add_audio_samples_to_wandb: bool = field(
        default=False,
Yoach Lacombe's avatar
Yoach Lacombe committed
454
        metadata={"help": "If set and if `wandb` in args.report_to, will add generated audio samples to wandb logs."},
455
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
456
    id_column_name: str = field(default=None, metadata={"help": "id column name."})
Yoach Lacombe's avatar
Yoach Lacombe committed
457
    wandb_project: str = field(
Yoach Lacombe's avatar
Yoach Lacombe committed
458
        default="parler-speech",
Yoach Lacombe's avatar
Yoach Lacombe committed
459
460
        metadata={"help": "The name of the wandb project."},
    )
461
462
463
464
    save_to_disk: str = field(
        default=None,
        metadata={
            "help": "If set, will save the dataset to this path if this is an empyt folder. If not empty, will load the datasets from it."
Yoach Lacombe's avatar
Yoach Lacombe committed
465
        },
466
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
467
    temporary_save_to_disk: str = field(default=None, metadata={"help": "Temporarily save audio labels here."})
468
469
    pad_to_multiple_of: Optional[int] = field(
        default=2,
Yoach Lacombe's avatar
Yoach Lacombe committed
470
        metadata={"help": ("Pad to multiple of for tokenizers.")},
471
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
472
473


Yoach Lacombe's avatar
Yoach Lacombe committed
474
@dataclass
Yoach Lacombe's avatar
Yoach Lacombe committed
475
class ParlerTTSTrainingArguments(Seq2SeqTrainingArguments):
Yoach Lacombe's avatar
Yoach Lacombe committed
476
477
478
479
480
481
482
483
484
    dtype: Optional[str] = field(
        default="float32",
        metadata={
            "help": (
                "The data type (dtype) in which to run training. One of `float32` (full-precision), "
                "`float16` or `bfloat16` (both half-precision)."
            )
        },
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
485
486
    audio_encode_per_device_eval_batch_size: int = field(
        default=8,
Yoach Lacombe's avatar
Yoach Lacombe committed
487
        metadata={"help": ("TODO")},
Yoach Lacombe's avatar
Yoach Lacombe committed
488
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
489

Yoach Lacombe's avatar
Yoach Lacombe committed
490

491
492
493
@dataclass
class DataCollatorEncodecWithPadding:
    """
Yoach Lacombe's avatar
Yoach Lacombe committed
494
    Data collator that will dynamically pad the inputs received to the longest sequence in the batch or
495
    to `max_length` if `max_length` is set and `padding=max_length`.
496
497
498
    """

    feature_extractor: AutoFeatureExtractor
499
    audio_column_name: str
500
    feature_extractor_input_name: Optional[str] = "input_values"
501
    max_length: Optional[int] = None
502
    padding: Optional[str] = "longest"
503
504
505
506

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        # split inputs and labels since they have to be of different lengths and need
        # different padding methods
Yoach Lacombe's avatar
Yoach Lacombe committed
507
        audios = [feature[self.audio_column_name]["array"] for feature in features]
508
        len_audio = [len(audio) for audio in audios]
509
510

        batch = self.feature_extractor(audios, return_tensors="pt", padding=self.padding, max_length=self.max_length)
511
512
        batch["len_audio"] = torch.tensor(len_audio).unsqueeze(1)
        return batch
513

Yoach Lacombe's avatar
Yoach Lacombe committed
514

515
@dataclass
Yoach Lacombe's avatar
Yoach Lacombe committed
516
class DataCollatorParlerTTSWithPadding:
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
    """
    Data collator that will dynamically pad the inputs received.
    Args:
        prompt_tokenizer (:class:`~transformers.AutoTokenizer`)
            The prompt_tokenizer used for proccessing the data.
        description_tokenizer (:class:`~transformers.AutoTokenizer`)
            The description_tokenizer used for proccessing the data.
        audio_feature_extractor (:class:`~transformers.AutoFeatureExtractor`)
            The audio_feature_extractor used for proccessing the data.
        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
            among:
            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
              sequence if provided).
            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
              maximum acceptable input length for the model if that argument is not provided.
            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
              different lengths).
        pad_to_multiple_of (:obj:`int`, `optional`):
            If set will pad the sequence to a multiple of the provided value.
            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
            7.5 (Volta).
    """

    prompt_tokenizer: AutoTokenizer
    description_tokenizer: AutoTokenizer
    audio_feature_extractor: AutoFeatureExtractor
    feature_extractor_input_name: Optional[str] = "input_values"
    padding: Union[bool, str] = "longest"
    pad_to_multiple_of: Optional[int] = None
547
548
549
    prompt_max_length: Optional[int] = None
    description_max_length: Optional[int] = None
    audio_max_length: Optional[int] = None
550
551
552
553

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        # split inputs and labels since they have to be of different lengths and need
        # different padding methods
Yoach Lacombe's avatar
Yoach Lacombe committed
554
555

        labels = [torch.tensor(feature["labels"]).transpose(0, 1) for feature in features]
556
        # (bsz, seq_len, num_codebooks)
Yoach Lacombe's avatar
Yoach Lacombe committed
557
558
559
560
        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)
        if self.audio_max_length is not None and self.padding == "max_length":
            labels = torch.nn.functional.pad(labels, pad=(0, 0, 0, max(self.audio_max_length - labels.shape[1], 0)))

561
        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
562

Yoach Lacombe's avatar
Yoach Lacombe committed
563
564
565
566
567
568
569
570
571
572
573
        input_ids = self.description_tokenizer.pad(
            input_ids,
            return_tensors="pt",
            padding=self.padding,
            pad_to_multiple_of=self.pad_to_multiple_of,
            max_length=self.description_max_length,
        )

        batch = {"labels": labels, **input_ids}

        if self.audio_max_length is not None and self.padding == "max_length":
574
575
576
            # if we do torch.compile, we need to also specify the attention_mask
            decoder_attention_mask = torch.ones(labels.shape[:2], dtype=input_ids["attention_mask"].dtype)
            batch["decoder_attention_mask"] = decoder_attention_mask
Yoach Lacombe's avatar
Yoach Lacombe committed
577

578
        prompt_input_ids = [{"input_ids": feature["prompt_input_ids"]} for feature in features]
Yoach Lacombe's avatar
Yoach Lacombe committed
579
580
581
582
583
584
585
586
        prompt_input_ids = self.prompt_tokenizer.pad(
            prompt_input_ids,
            return_tensors="pt",
            padding=self.padding,
            pad_to_multiple_of=self.pad_to_multiple_of,
            max_length=self.prompt_max_length,
        )

587
588
589
        batch["prompt_input_ids"] = prompt_input_ids["input_ids"]
        if "attention_mask" in prompt_input_ids:
            batch["prompt_attention_mask"] = prompt_input_ids["attention_mask"]
Yoach Lacombe's avatar
Yoach Lacombe committed
590

591
        if self.feature_extractor_input_name in features[0]:
592
            # TODO (YL): verify that it works - IMPORTANT -> probably not working
Yoach Lacombe's avatar
Yoach Lacombe committed
593
594
595
            input_values = [
                {self.feature_extractor_input_name: feature[self.feature_extractor_input_name]} for feature in features
            ]
596
            input_values = self.feature_extractor.pad(input_values, return_tensors="pt")
Yoach Lacombe's avatar
Yoach Lacombe committed
597
598
599

            batch[self.feature_extractor_input_name : input_values]

600
        return batch
601

Yoach Lacombe's avatar
Yoach Lacombe committed
602

603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
def convert_dataset_str_to_list(
    dataset_names,
    dataset_config_names,
    metadata_dataset_names=None,
    splits=None,
    dataset_samples=None,
    default_split="train",
):
    if isinstance(dataset_names, str):
        dataset_names = dataset_names.split("+")
        dataset_config_names = dataset_config_names.split("+")
        splits = splits.split("+") if splits is not None else None
        dataset_samples = dataset_samples.split("+") if dataset_samples is not None else None
        metadata_dataset_names = metadata_dataset_names.split("+") if metadata_dataset_names is not None else None

    # basic checks to ensure we've got the right number of datasets/configs/splits/columns/probs
    if len(dataset_names) != len(dataset_config_names):
        raise ValueError(
            f"Ensure one config is passed for each dataset, got {len(dataset_names)} datasets and"
            f" {len(dataset_config_names)} configs."
        )

    if splits is not None and len(splits) != len(dataset_names):
        raise ValueError(
            f"Ensure one split is passed for each dataset, got {len(dataset_names)} datasets and {len(splits)} splits."
        )

    if metadata_dataset_names is not None and len(metadata_dataset_names) != len(dataset_names):
        raise ValueError(
            f"Ensure one metadata dataset is passed for each dataset, got {len(dataset_names)} datasets and {len(metadata_dataset_names)} metadata datasets."
        )

    if dataset_samples is not None:
        if len(dataset_samples) != len(dataset_names):
            raise ValueError(
                f"Ensure one sample is passed for each dataset, got {len(dataset_names)} datasets and "
                f"{len(dataset_samples)} samples."
            )
        dataset_samples = [float(ds_sample) for ds_sample in dataset_samples]
    else:
        dataset_samples = [None] * len(dataset_names)

    splits = splits if splits is not None else [default_split for _ in range(len(dataset_names))]

    dataset_names_dict = []
    for i, ds_name in enumerate(dataset_names):
        dataset_names_dict.append(
            {
                "name": ds_name,
                "config": dataset_config_names[i],
                "split": splits[i],
                "metadata_dataset_name": metadata_dataset_names[i],
                "samples": dataset_samples[i],
            }
        )
    return dataset_names_dict


def load_multiple_datasets(
662
    accelerator: Accelerator,
663
664
    dataset_names: Union[List, str],
    dataset_config_names: Union[List, str],
Yoach Lacombe's avatar
Yoach Lacombe committed
665
    metadata_dataset_names: Optional[str] = None,
666
667
668
669
670
671
672
    splits: Optional[Union[List, str]] = None,
    label_column_names: Optional[List] = None,
    stopping_strategy: Optional[str] = "first_exhausted",
    dataset_samples: Optional[Union[List, np.array]] = None,
    streaming: Optional[bool] = False,
    seed: Optional[int] = None,
    id_column_name: Optional[str] = None,
Yoach Lacombe's avatar
Yoach Lacombe committed
673
    columns_to_keep: Optional[Set[str]] = None,
674
    prompt_column_name: Optional[str] = None,
675
676
    sampling_rate: Optional[int] = None,
    audio_column_name: Optional[str] = None,
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
    **kwargs,
) -> Union[Dataset, IterableDataset]:
    dataset_names_dict = convert_dataset_str_to_list(
        dataset_names, dataset_config_names, metadata_dataset_names, splits, label_column_names, dataset_samples
    )

    if dataset_samples is not None:
        dataset_samples = [ds_dict["samples"] for ds_dict in dataset_names_dict]
        probabilities = np.array(dataset_samples) / np.sum(dataset_samples)
    else:
        probabilities = None

    all_datasets = []
    # iterate over the datasets we want to interleave
    for dataset_dict in tqdm(dataset_names_dict, desc="Combining datasets..."):
692
693
694
        with accelerator.main_process_first():
            dataset = load_dataset(
                dataset_dict["name"],
695
696
697
698
699
                dataset_dict["config"],
                split=dataset_dict["split"],
                streaming=streaming,
                **kwargs,
            )
700
            dataset_features = dataset.features.keys()
Yoach Lacombe's avatar
Yoach Lacombe committed
701

702
703
            if sampling_rate is not None and audio_column_name is not None:
                # resample target audio
Yoach Lacombe's avatar
Yoach Lacombe committed
704
705
                dataset = dataset.cast_column(audio_column_name, datasets.features.Audio(sampling_rate=sampling_rate))

706
707
            metadata_dataset_name = dataset_dict["metadata_dataset_name"]
            if metadata_dataset_name is not None:
Yoach Lacombe's avatar
Yoach Lacombe committed
708
709
710
                logger.info(
                    f'Merging {dataset_dict["name"]} - {dataset_dict["split"]} with {metadata_dataset_name} - {dataset_dict["split"]}'
                )
711
712
713
714
715
716
717
                metadata_dataset = load_dataset(
                    metadata_dataset_name,
                    dataset_dict["config"],
                    split=dataset_dict["split"],
                    streaming=streaming,
                    **kwargs,
                )
Yoach Lacombe's avatar
Yoach Lacombe committed
718

719
                # TODO(YL): I forgot to create unique ids for MLS english.
720
                # To iterate faster, I bypass the original id check and do another one. - Done once because assuming it won't change next time
Yoach Lacombe's avatar
Yoach Lacombe committed
721
                # if dataset_dict["name"] == "parler-tts/mls_eng_10k":
722
723
724
725
726
727
                #     def concat_ids(book_id, speaker_id, begin_time):
                #         return {"id": f"{book_id}_{speaker_id}_{str(begin_time).replace('.', '_')}"}
                #     dataset = dataset.map(concat_ids, input_columns=["book_id", "speaker_id", "begin_time"], num_proc=24)
                #     metadata_dataset = metadata_dataset.map(concat_ids, input_columns=["book_id", "speaker_id", "begin_time"], num_proc=24)
                #     metadata_dataset = metadata_dataset.rename_column(id_column_name, f"metadata_{id_column_name}")

Yoach Lacombe's avatar
Yoach Lacombe committed
728
                if dataset_dict["name"] != "parler-tts/mls_eng_10k":
729
730
731
732
                    if id_column_name is not None and id_column_name not in dataset.column_names:
                        raise ValueError(
                            f"id_column_name={id_column_name} but has not been found in the dataset columns"
                            f"- one of {', '.join(list(dataset.column_names))}."
Yoach Lacombe's avatar
Yoach Lacombe committed
733
                        )
734
735
736
737
                    if id_column_name is not None and id_column_name not in metadata_dataset.column_names:
                        raise ValueError(
                            f"id_column_name={id_column_name} but has not been found in the metadata dataset columns"
                            f"- one of {', '.join(list(metadata_dataset.column_names))}."
Yoach Lacombe's avatar
Yoach Lacombe committed
738
                        )
739
740
                    elif id_column_name is not None:
                        metadata_dataset = metadata_dataset.rename_column(id_column_name, f"metadata_{id_column_name}")
Yoach Lacombe's avatar
Yoach Lacombe committed
741

742
                metadata_columns_to_remove = set(metadata_dataset.column_names).intersection(set(dataset.column_names))
Yoach Lacombe's avatar
Yoach Lacombe committed
743

744
745
746
747
                if prompt_column_name is not None:
                    # We might have applied some transformations to the prompts (e.g  punctuation restoration)
                    # so we make sure to remove it from the original dataset
                    if prompt_column_name in dataset.column_names:
Yoach Lacombe's avatar
Yoach Lacombe committed
748
749
750
                        logger.info(
                            f"REMOVE {prompt_column_name} from dataset {dataset_dict['name']} - dataset_dict['split']"
                        )
751
752
                        dataset.remove_columns(prompt_column_name)

753
754
                metadata_columns_to_remove = set(metadata_dataset.column_names).intersection(set(dataset.column_names))
                metadata_dataset = metadata_dataset.remove_columns(metadata_columns_to_remove)
755

756
                dataset = concatenate_datasets([dataset, metadata_dataset], axis=1)
Yoach Lacombe's avatar
Yoach Lacombe committed
757

Yoach Lacombe's avatar
Yoach Lacombe committed
758
                if id_column_name is not None and dataset_dict["name"] != "parler-tts/mls_eng_10k":
Yoach Lacombe's avatar
Yoach Lacombe committed
759
760
761
762
763
764
765
766
767
768
769
770
                    if (
                        len(
                            dataset.filter(
                                lambda id1, id2: id1 != id2,
                                input_columns=[id_column_name, f"metadata_{id_column_name}"],
                            )
                        )
                        != 0
                    ):
                        raise ValueError(
                            f"Concatenate didn't work. Some ids don't correspond on dataset {dataset_dict['name']}"
                        )
771

772
                dataset_features = dataset.features.keys()
Yoach Lacombe's avatar
Yoach Lacombe committed
773

774
775
            if columns_to_keep is not None:
                dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
776
777
778
779
780
781
782
783
784
785
786
787
788
789
        all_datasets.append(dataset)

    if len(all_datasets) == 1:
        # we have a single dataset so just return it as is
        return all_datasets[0]

    if streaming:
        interleaved_dataset = interleave_datasets(
            all_datasets,
            stopping_strategy=stopping_strategy,
            probabilities=probabilities,
            seed=seed,
        )
    else:
790
791
        with accelerator.main_process_first():
            interleaved_dataset = concatenate_datasets(all_datasets)
792
793
794

    return interleaved_dataset

Yoach Lacombe's avatar
Yoach Lacombe committed
795

796
797
798
799
800
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
801
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
802
803
804
805
806
807
808
809
810
    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
811
    send_example_telemetry("run_parler_tts", model_args, data_args)
Yoach Lacombe's avatar
Yoach Lacombe committed
812

Yoach Lacombe's avatar
Yoach Lacombe committed
813
814
815
816
817
818
    if training_args.dtype == "float16":
        mixed_precision = "fp16"
    elif training_args.dtype == "bfloat16":
        mixed_precision = "bf16"
    else:
        mixed_precision = "no"
Yoach Lacombe's avatar
Yoach Lacombe committed
819
820
821
822
823
824
825
826
827

    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`"
        )
828
829

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

831
    ####### A. Preparation
832
833
834
    kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=60))]
    if training_args.torch_compile:
        # TODO(YL): add more compile modes?
Yoach Lacombe's avatar
Yoach Lacombe committed
835
836
        kwargs_handlers.append(TorchDynamoPlugin(backend="inductor", mode="default"))  # reduce-overhead

Yoach Lacombe's avatar
Yoach Lacombe committed
837
838
839
840
841
    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,
842
        kwargs_handlers=kwargs_handlers,
Yoach Lacombe's avatar
Yoach Lacombe committed
843
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865

    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,
        },
    )

Yoach Lacombe's avatar
Yoach Lacombe committed
866
    # Detecting last checkpoint and eventually continue from last checkpoint
867
868
869
870
871
872
873
874
    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
875
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
876
877
878
879
880
881
882
883
884
885
886
            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)],
    )
887
    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
888

Yoach Lacombe's avatar
Yoach Lacombe committed
889
    # Log a small summary on each proces
890
891
892
893
    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
894
895
896
897

    # 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()
898
        transformers.utils.logging.set_verbosity_info()
Yoach Lacombe's avatar
Yoach Lacombe committed
899
900
901
902
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

903
904
905
906
    logger.info("Training/evaluation parameters %s", training_args)

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

909
910
911
    # 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
912

913
914
915
916
917
918
919
920
    # 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
921

922
923
924
925
926
927
928
    # 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,
Yoach Lacombe's avatar
Yoach Lacombe committed
929
        padding_side="left",  # prompt has to be padded on the left bc it's preprend to codebooks hidden states
930
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
931

932
933
934
935
936
937
938
939
    # 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
940

941
    if model_args.use_fast_tokenizer:
Yoach Lacombe's avatar
Yoach Lacombe committed
942
943
944
        logger.warning(
            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
        )
945
946
        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
947

948
    # 2. Now, let's load the dataset
Yoach Lacombe's avatar
Yoach Lacombe committed
949

950
951
    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
952

953
954
955
956
    # 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:
        vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
Yoach Lacombe's avatar
Yoach Lacombe committed
957
    else:
958
959
960
961
        raw_datasets = DatasetDict()

        columns_to_keep = {
            "target_audio_column_name": data_args.target_audio_column_name,
Yoach Lacombe's avatar
Yoach Lacombe committed
962
            "prompt_column_name": data_args.prompt_column_name,
963
964
        }
        if data_args.description_column_name is not None:
965
            columns_to_keep["description_column_name"] = data_args.description_column_name
Yoach Lacombe's avatar
Yoach Lacombe committed
966

967
968
969
970
971
972
973
974
975
976
977
978
979
        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(),
980
                prompt_column_name=data_args.prompt_column_name,
981
982
                audio_column_name=data_args.target_audio_column_name,
                sampling_rate=sampling_rate,
983
984
                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
985

986
987
988
989
990
991
            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
992
                    )
993
994
995
996
997
998
999
1000

            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
1001
1002
1003
                data_args.eval_dataset_config_name
                if data_args.eval_dataset_config_name
                else data_args.train_dataset_config_name,
1004
1005
1006
1007
1008
1009
                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(),
1010
1011
1012
                prompt_column_name=data_args.prompt_column_name,
                audio_column_name=data_args.target_audio_column_name,
                sampling_rate=sampling_rate,
1013
1014
                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
            )
1015

1016
            if data_args.max_eval_samples is not None:
Yoach Lacombe's avatar
Yoach Lacombe committed
1017
1018
1019
                raw_datasets["eval"] = (
                    raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
                )
1020

1021
1022
    # 3. Next, let's load the config.
    # TODO(YL): add the option to create the config from scratch
Yoach Lacombe's avatar
Yoach Lacombe committed
1023
    config = ParlerTTSConfig.from_pretrained(
1024
1025
1026
1027
1028
        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
1029

1030
    # update pad token id and decoder_start_token_id
1031
    # TODO(YL): verify if this makes sense, maybe should do it for model.decoder
Yoach Lacombe's avatar
Yoach Lacombe committed
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
    config.update(
        {
            "pad_token_id": model_args.pad_token_id
            if model_args.pad_token_id is not None
            else model.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 model.config.decoder_start_token_id,
        }
    )

1043
    # create model + TODO(YL): not from_pretrained probably
Yoach Lacombe's avatar
Yoach Lacombe committed
1044
    model = ParlerTTSForConditionalGeneration.from_pretrained(
1045
1046
1047
1048
1049
1050
        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,
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
1051

1052
1053
1054
    # enable gradient checkpointing if necessary
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
Yoach Lacombe's avatar
Yoach Lacombe committed
1055

1056
    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
1057
1058
1059
    # 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
1060

1061
    # derive max & min input length for sample rate & max duration
1062
1063
1064
    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
1065
1066
1067
1068
    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
1069
1070
    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
1071
1072
1073
    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
1074
    bandwidth = model_args.bandwidth
Yoach Lacombe's avatar
Yoach Lacombe committed
1075

1076
1077
    # Freeze Encoders
    model.freeze_encoders(model_args.freeze_text_encoder)
Yoach Lacombe's avatar
Yoach Lacombe committed
1078

1079
1080
1081
1082
1083
1084
    # TODO: remove when releasing
    # Test all gather - used for warmout and avoiding timeout
    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
1085
1086

    if not dataset_was_precomputed:
1087
        # Filter on text length
1088
        if description_column_name is not None and data_args.max_text_length is not None:
1089
1090
1091
1092
1093
1094
1095
            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],
                )
1096

1097
1098
1099
1100
        # Preprocessing the dataset.
        # We need to tokenize the texts.
        def pass_through_processors(description, prompt):
            batch = {}
Yoach Lacombe's avatar
Yoach Lacombe committed
1101

1102
1103
1104
            batch["input_ids"] = description_tokenizer(description.strip())["input_ids"]
            # TODO: add possibility to train without description column
            batch["prompt_input_ids"] = prompt_tokenizer(prompt.strip())["input_ids"]
1105
1106

            return batch
Yoach Lacombe's avatar
Yoach Lacombe committed
1107

1108
        with accelerator.main_process_first():
1109
            # this is a trick to avoid to rewrite the entire audio column which takes ages
1110
            vectorized_datasets = raw_datasets.map(
1111
1112
                pass_through_processors,
                remove_columns=next(iter(raw_datasets.values())).column_names,
1113
                input_columns=[description_column_name, prompt_column_name],
1114
1115
1116
                num_proc=num_workers,
                desc="preprocess datasets",
            )
1117

1118
        # We use Accelerate to perform distributed inference
1119
        # T5 doesn't support fp16
Yoach Lacombe's avatar
Yoach Lacombe committed
1120
        autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))
1121
1122

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

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

1127
1128
        # 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
1129
1130
1131
1132
        if training_args.torch_compile:
            audio_decoder = accelerator.prepare_model(model.audio_encoder, evaluation_mode=True)
        else:
            audio_decoder = model.audio_encoder
1133

Yoach Lacombe's avatar
Yoach Lacombe committed
1134
1135
1136
1137
1138
1139
1140
        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,
        )
1141
1142
1143
1144
1145
1146
1147
1148
1149

        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
1150
1151
            output["labels"] = labels.squeeze(0).transpose(1, 2)
            output["ratio"] = torch.ones_like(len_audio) * labels.shape[-1] / len_audio.max()
Yoach Lacombe's avatar
Yoach Lacombe committed
1152
            return output
1153

1154
1155
        for split in vectorized_datasets:
            data_loader = DataLoader(
1156
                raw_datasets[split],
1157
1158
1159
1160
                batch_size=training_args.audio_encode_per_device_eval_batch_size,
                collate_fn=encoder_data_collator,
                num_workers=training_args.dataloader_num_workers,
                pin_memory=True,
1161
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
1162
1163
            data_loader = accelerator.prepare(data_loader)

1164
1165
1166
1167
1168
1169
            all_generated_labels = []
            all_lens = []
            for batch in tqdm(data_loader, disable=not accelerator.is_local_main_process):
                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
1170

1171
                if accelerator.is_main_process:
Yoach Lacombe's avatar
Yoach Lacombe committed
1172
                    lab = generate_labels["labels"].cpu().transpose(1, 2).to(torch.int16)
1173
1174
                    rat = generate_labels["ratio"].cpu().squeeze()
                    lens = generate_labels["len_audio"].cpu().squeeze()
Yoach Lacombe's avatar
Yoach Lacombe committed
1175
1176
                    lab = [l[:, : int(ratio * length)] for (l, ratio, length) in zip(lab, rat, lens)]

1177
1178
                    all_generated_labels.extend(lab)
                    all_lens.extend(lens)
Yoach Lacombe's avatar
Yoach Lacombe committed
1179

1180
1181
            # (1, codebooks, seq_len) where seq_len=1
            bos_labels = torch.ones((1, num_codebooks, 1)) * audio_encoder_bos_token_id
Yoach Lacombe's avatar
Yoach Lacombe committed
1182

1183
            if accelerator.is_main_process:
1184
                tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
Yoach Lacombe's avatar
Yoach Lacombe committed
1185
1186
1187
1188
                tmp_labels.save_to_disk(
                    os.path.join(data_args.temporary_save_to_disk, split),
                    num_proc=1 if split == "eval" else data_args.preprocessing_num_workers,
                )
1189
1190
            accelerator.wait_for_everyone()
            del all_generated_labels
Yoach Lacombe's avatar
Yoach Lacombe committed
1191

1192
            tmp_labels = datasets.load_from_disk(os.path.join(data_args.temporary_save_to_disk, split))
1193
1194
            with accelerator.main_process_first():
                vectorized_datasets[split] = concatenate_datasets([vectorized_datasets[split], tmp_labels], axis=1)
Yoach Lacombe's avatar
Yoach Lacombe committed
1195

1196
            def postprocess_dataset(labels):
1197
                # (1, codebooks, seq_len)
Yoach Lacombe's avatar
Yoach Lacombe committed
1198
                labels = torch.tensor(labels).unsqueeze(0)
1199
1200
                # add bos
                labels = torch.cat([bos_labels, labels], dim=-1)
Yoach Lacombe's avatar
Yoach Lacombe committed
1201
1202
1203
1204
1205
1206
1207
1208
1209

                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,
                )

1210
1211
1212
1213
1214
1215
                # 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]
Yoach Lacombe's avatar
Yoach Lacombe committed
1216
1217
1218
                #  - [B, B, B, B, g, h, E]
                labels = torch.where(delay_pattern_mask == -1, audio_encoder_eos_token_id, delay_pattern_mask)

1219
1220
                # 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)
1221
                output = {"labels": labels[:, 1:]}
1222
1223
                return output

1224
            # TODO(YL): done multiple times, how to deal with it.
1225
1226
1227
            with accelerator.main_process_first():
                vectorized_datasets[split] = vectorized_datasets[split].map(
                    postprocess_dataset,
Yoach Lacombe's avatar
Yoach Lacombe committed
1228
                    num_proc=data_args.preprocessing_num_workers,  # this one is resource consuming if many processor.
1229
                    input_columns=["labels"],
1230
1231
1232
1233
                    desc="Postprocessing labeling",
                )

        accelerator.free_memory()
1234
        del generate_labels, all_lens
1235

1236
        with accelerator.main_process_first():
1237
            # 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
1238
            # caching audio files is time and disk-space consuming, so we want to avoid it at all costs, especially for large (>1Kh) audio datasets.
1239
1240
            # That's also why we avoid to concat the processed datasets (vectorized_datasets) with the audio column present in raw_datasets.

1241
1242
1243
1244
1245
1246
1247
1248
1249
            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
1250

1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
            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"],
                    )
Yoach Lacombe's avatar
Yoach Lacombe committed
1268

1269
    if data_args.save_to_disk is not None and not dataset_was_precomputed:
1270
        if accelerator.is_main_process:
Yoach Lacombe's avatar
Yoach Lacombe committed
1271
1272
1273
1274
            vectorized_datasets.save_to_disk(
                data_args.save_to_disk,
                num_proc=min(data_args.preprocessing_num_workers, len(vectorized_datasets["eval"]) - 1),
            )
1275
        logger.info(f"Dataset saved at {data_args.save_to_disk}")
Yoach Lacombe's avatar
Yoach Lacombe committed
1276

1277
1278
1279
    audio_max_length = None
    if training_args.torch_compile:
        audio_max_length = max(vectorized_datasets["train"]["target_length"])
Yoach Lacombe's avatar
Yoach Lacombe committed
1280
        with accelerator.main_process_first():
1281
            max_sample = vectorized_datasets["train"].filter(
Yoach Lacombe's avatar
Yoach Lacombe committed
1282
1283
1284
1285
                lambda x: x == audio_max_length,
                num_proc=num_workers,
                input_columns=["target_length"],
            )
1286
        audio_max_length = torch.tensor(max_sample[0]["labels"]).shape[1]
1287
1288
1289
1290
1291
1292

    # 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
1293
    if data_args.preprocessing_only and data_args.save_to_disk is None:
Yoach Lacombe's avatar
Yoach Lacombe committed
1294
1295
1296
        raise ValueError(
            "`preprocessing_only=True` but `save_to_disk` is not set. The latter should indicates where to save the dataset locally."
        )
1297
1298
    elif data_args.preprocessing_only:
        logger.info(f"Data preprocessing finished. Files save at {data_args.save_to_disk}")
1299
        return
Yoach Lacombe's avatar
Yoach Lacombe committed
1300

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

Yoach Lacombe's avatar
Yoach Lacombe committed
1303
    # Let's use word CLAP similary and WER metrics as our evaluation metrics,
1304

Yoach Lacombe's avatar
Yoach Lacombe committed
1305
    # Define evaluation metrics during training, *i.e.* CLAP similarity TODO: allow using another CLAP
1306
1307
    clap = AutoModel.from_pretrained("laion/larger_clap_music_and_speech")
    clap_processor = AutoProcessor.from_pretrained("laion/larger_clap_music_and_speech")
Yoach Lacombe's avatar
Yoach Lacombe committed
1308
    metric = evaluate.load("wer")
Yoach Lacombe's avatar
Yoach Lacombe committed
1309

Yoach Lacombe's avatar
Yoach Lacombe committed
1310
1311
1312
    def clap_similarity(texts, audios, device):
        clap_inputs = clap_processor(text=texts, audios=audios, padding=True, return_tensors="pt").to(device)
        clap.to(device)
1313
        with torch.no_grad():
Yoach Lacombe's avatar
Yoach Lacombe committed
1314
1315
1316
            text_features = clap.get_text_features(
                clap_inputs["input_ids"], attention_mask=clap_inputs.get("attention_mask", None)
            )
1317
            audio_features = clap.get_audio_features(clap_inputs["input_features"])
Yoach Lacombe's avatar
Yoach Lacombe committed
1318

1319
            cosine_sim = torch.nn.functional.cosine_similarity(audio_features, text_features, dim=1, eps=1e-8)
Yoach Lacombe's avatar
Yoach Lacombe committed
1320

Yoach Lacombe's avatar
Yoach Lacombe committed
1321
1322
        clap.to("cpu")
        clap_inputs.to("cpu")
1323
        return cosine_sim.mean().to("cpu")
Yoach Lacombe's avatar
Yoach Lacombe committed
1324

Yoach Lacombe's avatar
Yoach Lacombe committed
1325
1326
    def wer(prompts, audios, device):
        asr_pipeline = pipeline(model="distil-whisper/distil-large-v2", device=device)
Yoach Lacombe's avatar
Yoach Lacombe committed
1327
1328
1329
1330
1331
1332
1333
1334
1335
        transcriptions = asr_pipeline(
            [{"raw": audio, "sampling_rate": sampling_rate} for audio in audios],
            batch_size=int(training_args.per_device_eval_batch_size),
        )

        word_error = 100 * metric.compute(
            predictions=[t["text"].lower() for t in transcriptions], references=[t.lower() for t in prompts]
        )

Yoach Lacombe's avatar
Yoach Lacombe committed
1336
        return word_error, [t["text"] for t in transcriptions]
Yoach Lacombe's avatar
Yoach Lacombe committed
1337

Yoach Lacombe's avatar
Yoach Lacombe committed
1338
    eval_methods = {"clap": clap_similarity, "wer": wer}
1339

Yoach Lacombe's avatar
Yoach Lacombe committed
1340
1341
    def compute_metrics(audios, descriptions, prompts, device="cpu"):
        input_ids = descriptions
1342
        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
Yoach Lacombe's avatar
Yoach Lacombe committed
1343
1344
        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
        audios = [a.cpu().numpy() for a in audios]
Yoach Lacombe's avatar
Yoach Lacombe committed
1345
        results = {"clap": eval_methods["clap"](texts, audios, device)}
Yoach Lacombe's avatar
Yoach Lacombe committed
1346
1347
        word_error, transcriptions = eval_methods["wer"](prompts, audios, device)
        results["wer"] = word_error
1348

Yoach Lacombe's avatar
Yoach Lacombe committed
1349
        return results, texts, prompts, audios, transcriptions
Yoach Lacombe's avatar
Yoach Lacombe committed
1350

Yoach Lacombe's avatar
Yoach Lacombe committed
1351
1352
1353
1354
1355
1356
    # 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
1357

Yoach Lacombe's avatar
Yoach Lacombe committed
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
    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
1370
        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
Yoach Lacombe's avatar
Yoach Lacombe committed
1371
1372
1373
        eval_steps = steps_per_epoch
    else:
        eval_steps = training_args.eval_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
1374

1375
    # T5 doesn't support fp16
Yoach Lacombe's avatar
Yoach Lacombe committed
1376
1377
    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))

Yoach Lacombe's avatar
Yoach Lacombe committed
1378
1379
1380
1381
1382
1383
    # 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,
1384
        weight_decay=training_args.weight_decay,
Yoach Lacombe's avatar
Yoach Lacombe committed
1385
    )
1386

Yoach Lacombe's avatar
Yoach Lacombe committed
1387
1388
1389
1390
    # 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
1391
        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
Yoach Lacombe's avatar
Yoach Lacombe committed
1392
1393
        num_training_steps=total_train_steps * accelerator.num_processes,
    )
1394
1395

    # Instantiate custom data collator
Yoach Lacombe's avatar
Yoach Lacombe committed
1396
    data_collator = DataCollatorParlerTTSWithPadding(
Yoach Lacombe's avatar
Yoach Lacombe committed
1397
1398
1399
1400
1401
1402
1403
1404
1405
        audio_feature_extractor=feature_extractor,
        feature_extractor_input_name=feature_extractor_input_name,
        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,
1406
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
1407

Yoach Lacombe's avatar
Yoach Lacombe committed
1408
1409
    # Prepare everything with accelerate
    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
Yoach Lacombe's avatar
Yoach Lacombe committed
1410

Yoach Lacombe's avatar
Yoach Lacombe committed
1411
1412
    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {total_train_steps * train_batch_size * gradient_accumulation_steps}")
1413
    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
Yoach Lacombe's avatar
Yoach Lacombe committed
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
    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
1435

Yoach Lacombe's avatar
Yoach Lacombe committed
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
    if accelerator.is_main_process:
        if training_args.push_to_hub:
            # Retrieve of infer repo_name
            repo_name = training_args.hub_model_id
            if repo_name is None:
                repo_name = Path(training_args.output_dir).absolute().name
            # Create repo and retrieve repo_id
            repo_id = create_repo(repo_name, exist_ok=True, token=training_args.hub_token).repo_id
            # Clone repo locally
            repo = Repository(training_args.output_dir, clone_from=repo_id, token=training_args.hub_token)

            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
1453

Yoach Lacombe's avatar
Yoach Lacombe committed
1454
1455
1456
1457
1458
1459
    # Now save everything to be able to create a single processor later
    # make sure all processes wait until data is saved
    with accelerator.main_process_first():
        # only the main process saves them
        if accelerator.is_main_process:
            # save feature extractor, tokenizer and config
Yoach Lacombe's avatar
Yoach Lacombe committed
1460
1461
1462
1463
1464
            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)
            ):
Yoach Lacombe's avatar
Yoach Lacombe committed
1465
1466
                prompt_tokenizer.save_pretrained(training_args.output_dir)
            else:
Yoach Lacombe's avatar
Yoach Lacombe committed
1467
1468
1469
                logger.warning(
                    "Prompt tokenizer ('{model_args.prompt_tokenizer_name}') and description tokenizer ('{model_args.description_tokenizer_name}') are not the same. Saving only the prompt tokenizer."
                )
Yoach Lacombe's avatar
Yoach Lacombe committed
1470
                prompt_tokenizer.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
1471

Yoach Lacombe's avatar
Yoach Lacombe committed
1472
1473
            feature_extractor.save_pretrained(training_args.output_dir)
            config.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490

    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):
            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
Yoach Lacombe's avatar
Yoach Lacombe committed
1491

Yoach Lacombe's avatar
Yoach Lacombe committed
1492
1493
        if training_args.max_steps < 0:
            # we know exactly the number of steps per epoch, so can skip through the required number of batches
1494
            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
1495
1496
1497
1498
1499
1500
1501
1502
        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
            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
    else:
        resume_step = None
Yoach Lacombe's avatar
Yoach Lacombe committed
1503

Yoach Lacombe's avatar
Yoach Lacombe committed
1504
1505
    gen_kwargs = {
        "do_sample": model_args.do_sample,
yoach@huggingface.co's avatar
yoach@huggingface.co committed
1506
        "temperature": model_args.temperature,
Yoach Lacombe's avatar
Yoach Lacombe committed
1507
1508
        "max_length": model_args.max_length,
    }
Yoach Lacombe's avatar
Yoach Lacombe committed
1509

Yoach Lacombe's avatar
Yoach Lacombe committed
1510
1511
1512
    # Define gradient update step fn
    def train_step(
        batch,
1513
1514
        accelerator,
        autocast_kwargs,
Yoach Lacombe's avatar
Yoach Lacombe committed
1515
1516
    ):
        model.train()
Yoach Lacombe's avatar
Yoach Lacombe committed
1517

1518
        if mixed_precision == "fp16":
1519
1520
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
1521
                if training_args.parallel_mode.value != "distributed":
Yoach Lacombe's avatar
Yoach Lacombe committed
1522
1523
1524
                    encoder_outputs = model.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
1525
                else:
Yoach Lacombe's avatar
Yoach Lacombe committed
1526
1527
1528
                    encoder_outputs = model.module.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
1529
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
1530

Yoach Lacombe's avatar
Yoach Lacombe committed
1531
1532
1533
        outputs = model(**batch)
        # CE (data) loss
        ce_loss = outputs.loss
Yoach Lacombe's avatar
Yoach Lacombe committed
1534
        # TODO: add CE per codebook
Yoach Lacombe's avatar
Yoach Lacombe committed
1535
1536
1537

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

Yoach Lacombe's avatar
Yoach Lacombe committed
1539
    # Define eval fn
Yoach Lacombe's avatar
Yoach Lacombe committed
1540
1541
1542
1543
1544
    def eval_step(
        batch,
        accelerator,
        autocast_kwargs,
    ):
Yoach Lacombe's avatar
Yoach Lacombe committed
1545
1546
1547
        eval_model = model if not training_args.torch_compile else model._orig_mod
        eval_model.eval()

1548
        if mixed_precision == "fp16":
1549
1550
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
Yoach Lacombe's avatar
Yoach Lacombe committed
1551
1552
                with torch.no_grad():
                    if training_args.parallel_mode.value != "distributed" or training_args.torch_compile:
Yoach Lacombe's avatar
Yoach Lacombe committed
1553
1554
1555
                        encoder_outputs = eval_model.text_encoder(
                            input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                        )
Yoach Lacombe's avatar
Yoach Lacombe committed
1556
                    else:
Yoach Lacombe's avatar
Yoach Lacombe committed
1557
1558
1559
                        encoder_outputs = eval_model.module.text_encoder(
                            input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                        )
1560
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
1561
1562

        with torch.no_grad():
Yoach Lacombe's avatar
Yoach Lacombe committed
1563
            outputs = eval_model(**batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
1564
1565
1566
1567
1568
1569
        # CE (data) loss
        ce_loss = outputs.loss
        metrics = {"loss": ce_loss}
        return metrics

    def generate_step(batch):
1570
        batch.pop("decoder_attention_mask", None)
Yoach Lacombe's avatar
Yoach Lacombe committed
1571
        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=mixed_precision != "fp16").eval()
Yoach Lacombe's avatar
Yoach Lacombe committed
1572
1573
1574
1575
        if training_args.torch_compile:
            eval_model = model._orig_mod

        output_audios = eval_model.generate(**batch, **gen_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
1576
1577
1578
1579
1580
        output_audios = accelerator.pad_across_processes(output_audios, dim=1, pad_index=0)
        return output_audios

    for epoch in range(epochs_trained, num_epochs):
        vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
1581
        # TODO(YL): add args
Yoach Lacombe's avatar
Yoach Lacombe committed
1582
        sampler = LengthGroupedSampler(train_batch_size, lengths=vectorized_datasets["train"]["target_length"])
Yoach Lacombe's avatar
Yoach Lacombe committed
1583
1584
1585
1586
        train_dataloader = DataLoader(
            vectorized_datasets["train"],
            collate_fn=data_collator,
            batch_size=per_device_train_batch_size,
Yoach Lacombe's avatar
Yoach Lacombe committed
1587
            sampler=sampler,
Yoach Lacombe's avatar
Yoach Lacombe committed
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
            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
            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
            resume_step = None

        for batch in train_dataloader:
            with accelerator.accumulate(model):
1602
                loss, train_metric = train_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
                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}")
1634
1635
1636
                    # safe_serialization=False to avoid shared tensors saving issue (TODO: it's a temporary fix)
                    # 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
1637
1638
1639
1640
1641
1642
                    accelerator.wait_for_everyone()
                    if accelerator.is_main_process:
                        rotate_checkpoints(training_args.save_total_limit, output_dir=training_args.output_dir)

                        if cur_step == total_train_steps:
                            # un-wrap student model for save
Yoach Lacombe's avatar
Yoach Lacombe committed
1643
1644
                            unwrapped_model = accelerator.unwrap_model(model)
                            unwrapped_model.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659

                        if training_args.push_to_hub:
                            repo.push_to_hub(
                                commit_message=f"Saving train state of step {cur_step}",
                                blocking=False,
                            )

                if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
                    train_time += time.time() - train_start
                    # ======================== Evaluating ==============================
                    eval_metrics = []
                    eval_preds = []
                    eval_descriptions = []
                    eval_prompts = []
                    eval_start = time.time()
Yoach Lacombe's avatar
Yoach Lacombe committed
1660

Yoach Lacombe's avatar
Yoach Lacombe committed
1661
1662
                    # release training input batch
                    batch = release_memory(batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
1663
1664
1665
1666
1667

                    validation_dataloader = DataLoader(
                        vectorized_datasets["eval"],
                        collate_fn=data_collator,
                        batch_size=per_device_eval_batch_size,
1668
                        drop_last=False,
Yoach Lacombe's avatar
Yoach Lacombe committed
1669
1670
1671
1672
1673
1674
1675
                        num_workers=training_args.dataloader_pin_memory,
                        pin_memory=training_args.dataloader_pin_memory,
                    )
                    validation_dataloader = accelerator.prepare(validation_dataloader)

                    for batch in tqdm(
                        validation_dataloader,
1676
                        desc=f"Evaluating - Inference ...",
Yoach Lacombe's avatar
Yoach Lacombe committed
1677
1678
1679
1680
                        position=2,
                        disable=not accelerator.is_local_main_process,
                    ):
                        # Model forward
1681
                        eval_metric = eval_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
1682
1683
1684
                        eval_metric = accelerator.gather_for_metrics(eval_metric)
                        eval_metrics.append(eval_metric)

1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
                    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
1695
                        # generation
1696
                        for batch in tqdm(
Yoach Lacombe's avatar
Yoach Lacombe committed
1697
1698
1699
1700
1701
                            validation_dataloader,
                            desc=f"Evaluating - Generation ...",
                            position=2,
                            disable=not accelerator.is_local_main_process,
                        ):
Yoach Lacombe's avatar
Yoach Lacombe committed
1702
1703
1704
1705
                            generated_audios = generate_step(batch)
                            # Gather all predictions and targets
                            # TODO: also add prompt ids
                            # TODO: better gather
Yoach Lacombe's avatar
Yoach Lacombe committed
1706
1707
1708
1709
1710
1711
                            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)
                            )
1712
1713
1714
                            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
1715
1716
1717
1718

                    eval_time = time.time() - eval_start
                    # normalize eval metrics
                    eval_metrics = {
Yoach Lacombe's avatar
Yoach Lacombe committed
1719
1720
                        key: torch.mean(torch.cat([d[key].unsqueeze(0) for d in eval_metrics]))
                        for key in eval_metrics[0]
Yoach Lacombe's avatar
Yoach Lacombe committed
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
                    }

                    # compute metrics
                    metrics_desc = ""
                    if training_args.predict_with_generate:
                        metric_values, pred_descriptions, pred_prompts, audios, transcriptions = compute_metrics(
                            eval_preds, eval_descriptions, eval_prompts, accelerator.device
                        )
                        eval_metrics.update(metric_values)
                        metrics_desc = " ".join([f"Eval {key}: {value} |" for key, value in metric_values.items()])
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
                        if "wandb" in training_args.report_to:
                            log_pred(
                                accelerator,
                                pred_descriptions,
                                pred_prompts,
                                transcriptions,
                                audios,
                                sampling_rate=sampling_rate,
                                step=cur_step,
                                prefix="eval",
                            )
Yoach Lacombe's avatar
Yoach Lacombe committed
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756

                    # Print metrics and update progress bar
                    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",
                    )
Yoach Lacombe's avatar
Yoach Lacombe committed
1757

1758
1759
1760
1761
1762
1763
1764
                    # release eval batch and relax metrics
                    eval_metrics = []
                    eval_preds = []
                    eval_descriptions = []
                    eval_prompts = []
                    batch = release_memory(batch)

Yoach Lacombe's avatar
Yoach Lacombe committed
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
                    # 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()
1777
1778
1779


if __name__ == "__main__":
1780
    set_start_method("spawn")
Yoach Lacombe's avatar
Yoach Lacombe committed
1781
    main()