run_stable_speech_training.py 78.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" Train a text-to-speech model using 🤗 Transformers Seq2SeqTrainer"""

import functools
import json
import logging
import os
import re
import sys
Yoach Lacombe's avatar
Yoach Lacombe committed
25
import shutil
26
27
import warnings
import math
Yoach Lacombe's avatar
Yoach Lacombe committed
28
import time
29
from multiprocess import set_start_method
30
from datetime import timedelta
31

32

Yoach Lacombe's avatar
Yoach Lacombe committed
33
import evaluate
34
from tqdm import tqdm
Yoach Lacombe's avatar
Yoach Lacombe committed
35
from pathlib import Path
36
from dataclasses import dataclass, field
Yoach Lacombe's avatar
Yoach Lacombe committed
37
from typing import Dict, List, Optional, Union, Set
38
39
40
41

import datasets
import numpy as np
import torch
42
43
from torch.utils.data import DataLoader

44
45
from datasets import DatasetDict, load_dataset, Dataset, IterableDataset, interleave_datasets, concatenate_datasets

Yoach Lacombe's avatar
Yoach Lacombe committed
46
from huggingface_hub import Repository, create_repo
47
48
49
50
import transformers
from transformers import (
    AutoFeatureExtractor,
    AutoModel,
Yoach Lacombe's avatar
Yoach Lacombe committed
51
    AutoModelWithLMHead,
52
53
54
55
56
57
    AutoProcessor,
    AutoTokenizer,
    HfArgumentParser,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
)
Yoach Lacombe's avatar
Yoach Lacombe committed
58
from transformers.trainer_utils import is_main_process
Yoach Lacombe's avatar
Yoach Lacombe committed
59
from transformers.trainer_pt_utils import LengthGroupedSampler
Yoach Lacombe's avatar
Yoach Lacombe committed
60
61
from transformers import pipeline
from transformers.optimization import get_scheduler
62
63
64
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
add DAC  
Yoach Lacombe committed
65
66
from transformers import AutoConfig, AutoModel
from stable_speech import DACConfig, DACModel
67
from transformers.modeling_outputs import BaseModelOutput
Yoach Lacombe's avatar
add DAC  
Yoach Lacombe committed
68
69
70
AutoConfig.register("dac", DACConfig)
AutoModel.register(DACConfig, DACModel)

71
72

from accelerate import Accelerator
73
from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin
Yoach Lacombe's avatar
Yoach Lacombe committed
74
from accelerate.utils.memory import release_memory
75

76
from stable_speech import StableSpeechForConditionalGeneration, StableSpeechConfig, apply_delay_pattern_mask, build_delay_pattern_mask
77

Yoach Lacombe's avatar
Yoach Lacombe committed
78
79
if is_wandb_available():
    from wandb import Audio
80
81
82
83
84
85
86
87
88
89
90
91
92

# 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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
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
182
183
184
185
186
187
188
189
190
191
192
193
194
_RE_CHECKPOINT = re.compile(r"^checkpoint-(\d+)-epoch-(\d+)$")

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

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

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)

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)

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,
        )
        
        # wandb can only loads 100 audios per step
        wandb_tracker.log({
                "Speech samples": [
                    Audio(
                        audio,
                        caption=f"{pred_prompts[i]} --- DESCRIPTION: {pred_descriptions[i]}",
                        sample_rate=sampling_rate,
                    )
                    for (i, audio) in enumerate(audios[:min(len(audios), 100)])
                ]},
                step=step)
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240

@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """
    # 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(
        default=None, metadata={"help": "Pretrained prompt tokenizer name or path if not the same as description_tokenizer_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    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
241
242
243
244
    do_sample: bool = field(
        default=False,
        metadata={"help": "Whether to do sampling or greedy decoding."},
    )
yoach@huggingface.co's avatar
yoach@huggingface.co committed
245
246
247
248
    temperature: float = field(
        default=0.4,
        metadata={"help": "Temperature if sampling."},
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
249
    max_length: int = field(
250
251
        default=2580,
        metadata={"help": "Generation max length."},
Yoach Lacombe's avatar
Yoach Lacombe committed
252
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
253
    bandwidth: float = field(
Yoach Lacombe's avatar
Yoach Lacombe committed
254
        default=6, # TODO
Yoach Lacombe's avatar
Yoach Lacombe committed
255
256
        metadata={"help": "Audio encoder bandwidth."},
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
257
    
258
259
260


@dataclass
Yoach Lacombe's avatar
Yoach Lacombe committed
261
class DataTrainingArguments:
262
263
264
265
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
336
337
338
339
340
341
342
343
344
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
    """
    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'`."
        },
    )
    target_audio_column_name: str = field( # TODO
        default="audio",
        metadata={"help": "The name of the dataset column containing the target audio data. Defaults to 'audio'"},
    )
    description_column_name: str = field( #TODO
        default=None,
        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'None'."},
    )
    prompt_column_name: str = field( #TODO
        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": (
373
374
                "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`."
375
376
377
378
379
380
            )
        },
    )
    min_duration_in_seconds: float = field(
        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
    )
381
    max_text_length: int = field(
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
        default=500, metadata={"help": "If set, max description lengths in number of characters."}
    )
    max_prompt_token_length: int = field(
        default=None, metadata={
            "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`."
            )
        }
    )
    max_description_token_length: int = field(
        default=None, metadata={
            "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`."
            )
        }    
    )
    pad_to_max_length: bool = field(
        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`."
                )
            }
407
    )
408
409
410
411
412
413
414
    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"
415
416
                " 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. "
417
418
419
420
421
422
423
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
454
455
456
            )
        },
    )
    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,
        metadata={
            "help": "If set and if `wandb` in args.report_to, will add generated audio samples to wandb logs."
        }
    )
    id_column_name: str = field(
        default=None,
        metadata={
            "help": "id column name."
        }
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
457
458
459
460
    wandb_project: str = field(
        default="stable-speech",
        metadata={"help": "The name of the wandb project."},
    )
461
462
463
464
465
466
    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."
        }
    )
467
468
469
470
471
472
    temporary_save_to_disk: str = field(
        default=None,
        metadata={
            "help": "Temporarily save audio labels here."
        }
    )
473
474
475
476
477
478
479
480
    pad_to_multiple_of: Optional[int] = field(
        default=2,
        metadata={
            "help": (
                "Pad to multiple of for tokenizers."
            )
        },
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
481
482
483
484
485
486
487
488
489
490
491
492
    
@dataclass
class StableSpeechTrainingArguments(Seq2SeqTrainingArguments):
    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
493
494
495
496
497
498
499
500
    audio_encode_per_device_eval_batch_size: int = field(
        default=8,
        metadata={
            "help": (
                "TODO"
            )
        },
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
501

502
503
504
@dataclass
class DataCollatorEncodecWithPadding:
    """
505
506
    Data collator that will dynamically pad the inputs received to the longest sequence in the batch or 
    to `max_length` if `max_length` is set and `padding=max_length`.
507
508
509
    """

    feature_extractor: AutoFeatureExtractor
510
    audio_column_name: str
511
    feature_extractor_input_name: Optional[str] = "input_values"
512
    max_length: Optional[int] = None
513
    padding: Optional[str] = "longest"
514

515

516
517
518
    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
519
        audios = [feature[self.audio_column_name]["array"] for feature in features]
520
        len_audio = [len(audio) for audio in audios]
521
522

        batch = self.feature_extractor(audios, return_tensors="pt", padding=self.padding, max_length=self.max_length)
523
524
        batch["len_audio"] = torch.tensor(len_audio).unsqueeze(1)
        return batch
525
526
527

    
@dataclass
Yoach Lacombe's avatar
Yoach Lacombe committed
528
class DataCollatorStableSpeechWithPadding:
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
    """
    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
559
560
561
    prompt_max_length: Optional[int] = None
    description_max_length: Optional[int] = None
    audio_max_length: Optional[int] = None
562
563
564
565
566
567
568
569
570

    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
        
        
        labels = [torch.tensor(feature["labels"]).transpose(0,1) for feature in features]
        # (bsz, seq_len, num_codebooks)
        labels = torch.nn.utils.rnn.pad_sequence(labels,batch_first=True,padding_value=-100)
571
572
        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)))
573
574
        
        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
575
576
577
578
579
580
581
582
583

        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":
            # 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
584

585
        prompt_input_ids = [{"input_ids": feature["prompt_input_ids"]} for feature in features]
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
590
591
592
        
        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"]
        
        if self.feature_extractor_input_name in features[0]:
593
            # TODO (YL): verify that it works - IMPORTANT -> probably not working
594
595
596
597
598
599
            input_values = [{self.feature_extractor_input_name: feature[self.feature_extractor_input_name]} for feature in features]
            input_values = self.feature_extractor.pad(input_values, return_tensors="pt")
            
            batch[self.feature_extractor_input_name: input_values]
        
        return batch
600

601
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
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(
660
    accelerator: Accelerator,
661
662
663
664
665
666
667
668
669
670
    dataset_names: Union[List, str],
    dataset_config_names: Union[List, str],
    metadata_dataset_names: Optional[str]=None,
    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
671
    columns_to_keep: Optional[Set[str]] = None,
672
    prompt_column_name: Optional[str] = None,
673
674
    sampling_rate: Optional[int] = None,
    audio_column_name: Optional[str] = None,
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
    **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..."):
690
691
692
        with accelerator.main_process_first():
            dataset = load_dataset(
                dataset_dict["name"],
693
694
695
696
697
                dataset_dict["config"],
                split=dataset_dict["split"],
                streaming=streaming,
                **kwargs,
            )
698
699
            dataset_features = dataset.features.keys()
            
700
701
702
703
704
705
            if sampling_rate is not None and audio_column_name is not None:
                # resample target audio
                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:
708
                logger.info(f'Merging {dataset_dict["name"]} - {dataset_dict["split"]} with {metadata_dataset_name} - {dataset_dict["split"]}')
709
710
711
712
713
714
715
                metadata_dataset = load_dataset(
                    metadata_dataset_name,
                    dataset_dict["config"],
                    split=dataset_dict["split"],
                    streaming=streaming,
                    **kwargs,
                )
716
717
                
                # TODO(YL): I forgot to create unique ids for MLS english.
718
                # To iterate faster, I bypass the original id check and do another one. - Done once because assuming it won't change next time
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
                # if dataset_dict["name"] == "stable-speech/mls_eng_10k":
                #     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}")

                if dataset_dict["name"] != "stable-speech/mls_eng_10k":
                    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))}."
                            )
                    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))}."
                            )
                    elif id_column_name is not None:
                        metadata_dataset = metadata_dataset.rename_column(id_column_name, f"metadata_{id_column_name}")
739
                    
740
                metadata_columns_to_remove = set(metadata_dataset.column_names).intersection(set(dataset.column_names))
741
                
742
743
744
745
                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:
746
                        logger.info(f"REMOVE {prompt_column_name} from dataset {dataset_dict['name']} - dataset_dict['split']")
747
748
                        dataset.remove_columns(prompt_column_name)

749
750
                metadata_columns_to_remove = set(metadata_dataset.column_names).intersection(set(dataset.column_names))
                metadata_dataset = metadata_dataset.remove_columns(metadata_columns_to_remove)
751

752
                dataset = concatenate_datasets([dataset, metadata_dataset], axis=1)
753
                
754
                if id_column_name is not None and dataset_dict["name"] != "stable-speech/mls_eng_10k":
755
756
                    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']}")
757

758
                dataset_features = dataset.features.keys()
Yoach Lacombe's avatar
Yoach Lacombe committed
759

760
761
            if columns_to_keep is not None:
                dataset = dataset.remove_columns(set(dataset_features - columns_to_keep))
762
763
764
765
766
767
768
769
770
771
772
773
774
775
        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:
776
777
        with accelerator.main_process_first():
            interleaved_dataset = concatenate_datasets(all_datasets)
778
779
780
781
782
783
784
785
786
787

    return interleaved_dataset

    
    
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
788
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, StableSpeechTrainingArguments))
789
790
791
792
793
794
795
796
797
798
    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.
    send_example_telemetry("run_stable_speech", model_args, data_args)
799
    
Yoach Lacombe's avatar
Yoach Lacombe committed
800
801
802
803
804
805
    if training_args.dtype == "float16":
        mixed_precision = "fp16"
    elif training_args.dtype == "bfloat16":
        mixed_precision = "bf16"
    else:
        mixed_precision = "no"
806
807
808
809
810
        
    if data_args.pad_to_max_length and (data_args.max_duration_in_seconds is None or data_args.max_prompt_token_length is None or data_args.max_description_token_length is None):
        raise ValueError("`pad_to_max_length` is `True` but one of the following parameters has not been set: `max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`")

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

812
    ####### A. Preparation
813
814
815
    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
816
        kwargs_handlers.append(TorchDynamoPlugin(backend="inductor", mode="default")) #reduce-overhead
817
        
Yoach Lacombe's avatar
Yoach Lacombe committed
818
819
820
821
822
    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,
823
        kwargs_handlers=kwargs_handlers,
Yoach Lacombe's avatar
Yoach Lacombe committed
824
825
    )
    
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
    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,
yoach@huggingface.co's avatar
yoach@huggingface.co committed
841
        "temperature": model_args.temperature,
842
        })
Yoach Lacombe's avatar
Yoach Lacombe committed
843
844
845
    
    
    # Detecting last checkpoint and eventually continue from last checkpoint
846
847
848
849
850
851
852
853
    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
854
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
855
856
857
858
            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."
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
859
    
860
861
862
863
864
865
866

    # 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)],
    )
867
    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
868

Yoach Lacombe's avatar
Yoach Lacombe committed
869
    # Log a small summary on each proces
870
871
872
873
    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
874
875
876
877

    # 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()
878
        transformers.utils.logging.set_verbosity_info()
Yoach Lacombe's avatar
Yoach Lacombe committed
879
880
881
882
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

883

884
885
886
887
    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed before initializing model.
    set_seed(training_args.seed)
888
    num_workers = data_args.preprocessing_num_workers
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
    
    # 1. First, lett's instantiate the feature extractor, tokenizers and model
    # Note for distributed training, the .from_pretrained methods guarantee that only
    # one local process can concurrently download model & vocab.
    
    # load feature extractor
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        model_args.feature_extractor_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
    sampling_rate = feature_extractor.sampling_rate
    
    # load prompt tokenizer
    prompt_tokenizer = AutoTokenizer.from_pretrained(
        model_args.prompt_tokenizer_name or model_args.description_tokenizer_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
        use_fast=model_args.use_fast_tokenizer,
        padding_side="left", # prompt has to be padded on the left bc it's preprend to codebooks hidden states
    )
        
    # load description tokenizer
    description_tokenizer = AutoTokenizer.from_pretrained(
        model_args.description_tokenizer_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
        use_fast=model_args.use_fast_tokenizer,
    )
    
    if model_args.use_fast_tokenizer:
        logger.warning("Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235")
        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
926

927
    # 2. Now, let's load the dataset
Yoach Lacombe's avatar
Yoach Lacombe committed
928
    
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
    if data_args.save_to_disk is not None:
        os.makedirs(data_args.save_to_disk, exist_ok=True)
    
    # assume that the dataset has been saved to `save_to_disk` if the latter is not empty
    dataset_was_precomputed = len(os.listdir(data_args.save_to_disk)) > 0
    if dataset_was_precomputed:
        vectorized_datasets = datasets.load_from_disk(data_args.save_to_disk)
    else:    
        raw_datasets = DatasetDict()

        columns_to_keep = {
            "target_audio_column_name": data_args.target_audio_column_name,
            "prompt_column_name": data_args.prompt_column_name
        }
        if data_args.description_column_name is not None:
944
            columns_to_keep["description_column_name"] = data_args.description_column_name
945
946
947
948
949
950
951
952
953
954
955
956
957
958
            
        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(),
959
                prompt_column_name=data_args.prompt_column_name,
960
961
                audio_column_name=data_args.target_audio_column_name,
                sampling_rate=sampling_rate,
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
                # streaming=data_args.streaming, TODO(SG): optionally enable streaming mode
            )
            
            for key in columns_to_keep:
                if columns_to_keep[key] not in raw_datasets["train"].column_names:
                    raise ValueError(
                        f"--{key} '{columns_to_keep[key]}' not found in dataset '{data_args.train_dataset_name}'."
                        f" Make sure to set `--{key}` to the correct audio column - one of"
                        f" {', '.join(raw_datasets['train'].column_names)}."
                    )        

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

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

993
            if data_args.max_eval_samples is not None:
994
                raw_datasets["eval"] = raw_datasets["eval"].shuffle(seed=training_args.seed).select(range(data_args.max_eval_samples))
995
996


997
998
    # 3. Next, let's load the config.
    # TODO(YL): add the option to create the config from scratch
999
1000
1001
1002
1003
1004
1005
1006
    config = StableSpeechConfig.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
    
    # update pad token id and decoder_start_token_id
1007
    # TODO(YL): verify if this makes sense, maybe should do it for model.decoder
1008
1009
1010
1011
    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,
    })
1012
    
1013
    # create model + TODO(YL): not from_pretrained probably
1014
1015
1016
1017
1018
1019
1020
1021
    model = StableSpeechForConditionalGeneration.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        config=config,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
    
1022
1023
1024
    # enable gradient checkpointing if necessary
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
1025
    
1026
    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
1027
1028
1029
1030
1031
    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,
    # so that we just need to set the correct target sampling rate and normalize the input
    # via the `feature_extractor`
    
    # derive max & min input length for sample rate & max duration
1032
1033
1034
    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
1035
1036
1037
1038
    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
1039
1040
    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
1041
1042
1043
    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
1044
    bandwidth = model_args.bandwidth
1045
    
1046
1047
    # Freeze Encoders
    model.freeze_encoders(model_args.freeze_text_encoder)
1048
1049
1050
1051
1052
1053
1054
1055
    
    # 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()
    
1056
1057
    if not dataset_was_precomputed:        
        # Filter on text length
1058
        if description_column_name is not None and data_args.max_text_length is not None:
1059
1060
1061
1062
1063
1064
1065
            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],
                )
1066

1067
1068
1069
1070
1071
1072
1073
1074
        # Preprocessing the dataset.
        # We need to tokenize the texts.
        def pass_through_processors(description, prompt):
            batch = {}
            
            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"]
1075
1076

            return batch
1077
        
1078
        with accelerator.main_process_first():
1079
            # this is a trick to avoid to rewrite the entire audio column which takes ages
1080
            vectorized_datasets = raw_datasets.map(
1081
1082
                pass_through_processors,
                remove_columns=next(iter(raw_datasets.values())).column_names,
1083
                input_columns=[description_column_name, prompt_column_name],
1084
1085
1086
                num_proc=num_workers,
                desc="preprocess datasets",
            )
1087

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

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

1095
        logger.info("*** Encode target audio with encodec ***")
1096
        
1097
1098
        # 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
1099
1100
1101
1102
        if training_args.torch_compile:
            audio_decoder = accelerator.prepare_model(model.audio_encoder, evaluation_mode=True)
        else:
            audio_decoder = model.audio_encoder
1103

1104
        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)
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115

        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)
            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
1116
            return output
1117

1118
1119
        for split in vectorized_datasets:
            data_loader = DataLoader(
1120
                raw_datasets[split],
1121
1122
1123
1124
                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,
1125
            )
1126
            data_loader = accelerator.prepare(data_loader)        
1127
            
1128
1129
1130
1131
1132
1133
1134
            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)
                
1135
                if accelerator.is_main_process:
1136
1137
1138
1139
1140
1141
1142
                    lab = generate_labels["labels"].cpu().transpose(1,2).to(torch.int16)
                    rat = generate_labels["ratio"].cpu().squeeze()
                    lens = generate_labels["len_audio"].cpu().squeeze()
                    lab = [l[:, :int(ratio*length)] for (l, ratio, length) in zip(lab, rat, lens)]
                    
                    all_generated_labels.extend(lab)
                    all_lens.extend(lens)
1143
1144
1145
                
            # (1, codebooks, seq_len) where seq_len=1
            bos_labels = torch.ones((1, num_codebooks, 1)) * audio_encoder_bos_token_id
1146
1147
            
            if accelerator.is_main_process:
1148
1149
                tmp_labels = Dataset.from_dict({"labels": all_generated_labels, "target_length": all_lens})
                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)
1150
1151
1152
            accelerator.wait_for_everyone()
            del all_generated_labels
            
1153
            tmp_labels = datasets.load_from_disk(os.path.join(data_args.temporary_save_to_disk, split))
1154
1155
1156
1157
            with accelerator.main_process_first():
                vectorized_datasets[split] = concatenate_datasets([vectorized_datasets[split], tmp_labels], axis=1)
            
            
1158
            def postprocess_dataset(labels):
1159
                # (1, codebooks, seq_len)
1160
                labels = torch.tensor(labels).unsqueeze(0)         
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
                # add bos
                labels = torch.cat([bos_labels, labels], dim=-1)
                
                labels, delay_pattern_mask = build_delay_pattern_mask(labels, 
                                                        bos_token_id=audio_encoder_bos_token_id,
                                                        pad_token_id=audio_encoder_eos_token_id,
                                                        max_length=labels.shape[-1] + num_codebooks,
                                                        num_codebooks=num_codebooks)
                
                # the first ids of the delay pattern mask are precisely labels, we use the rest of the labels mask
                # to take care of EOS
                # we want labels to look like this:
                #  - [B, a, b, E, E, E, E]
                #  - [B, B, c, d, E, E, E]
                #  - [B, B, B, e, f, E, E]
                #  - [B, B, B, B, g, h, E] 
                labels = torch.where(delay_pattern_mask==-1, audio_encoder_eos_token_id, delay_pattern_mask)
                            
                # the first timestamp is associated to a row full of BOS, let's get rid of it
                # we also remove the last timestampts (full of PAD)
1181
                output = {"labels": labels[:, 1:]}
1182
                return output
1183
            
1184

1185
            # TODO(YL): done multiple times, how to deal with it.
1186
1187
1188
            with accelerator.main_process_first():
                vectorized_datasets[split] = vectorized_datasets[split].map(
                    postprocess_dataset,
1189
                    num_proc=data_args.preprocessing_num_workers, # this one is resource consuming if many processor.
1190
                    input_columns=["labels"],
1191
1192
1193
1194
1195
                    desc="Postprocessing labeling",
                )

                
        accelerator.free_memory()
1196
        del generate_labels, all_lens
1197
        
1198

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

1204
1205
1206
1207
1208
1209
1210
1211
1212
            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"],
            )
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
            
            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"],
                    )
1231
        
1232
    if data_args.save_to_disk is not None and not dataset_was_precomputed:
1233
        if accelerator.is_main_process:
1234
            vectorized_datasets.save_to_disk(data_args.save_to_disk, num_proc=min(data_args.preprocessing_num_workers, len(vectorized_datasets["eval"])-1))
1235
        logger.info(f"Dataset saved at {data_args.save_to_disk}")
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
        
    
    audio_max_length = None
    if training_args.torch_compile:
        audio_max_length = max(vectorized_datasets["train"]["target_length"])
        with accelerator.main_process_first(): 
            max_sample = vectorized_datasets["train"].filter(
                        lambda x: x == audio_max_length,
                        num_proc=num_workers,
                        input_columns=["target_length"],
                    )
        audio_max_length = torch.tensor(max_sample[0]["labels"]).shape[1]
1248
1249
1250
1251
1252
1253

    # 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
1254
1255
1256
1257
    if data_args.preprocessing_only and data_args.save_to_disk is None:
        raise ValueError("`preprocessing_only=True` but `save_to_disk` is not set. The latter should indicates where to save the dataset locally.")
    elif data_args.preprocessing_only:
        logger.info(f"Data preprocessing finished. Files save at {data_args.save_to_disk}")
1258
1259
        return
    
Yoach Lacombe's avatar
Yoach Lacombe committed
1260
    
1261
    # 6. Next, we can prepare the training.
Yoach Lacombe's avatar
Yoach Lacombe committed
1262
1263
    
    # Let's use word CLAP similary and WER metrics as our evaluation metrics,
1264

Yoach Lacombe's avatar
Yoach Lacombe committed
1265
    # Define evaluation metrics during training, *i.e.* CLAP similarity TODO: allow using another CLAP
1266
1267
    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
1268
    metric = evaluate.load("wer")
1269
1270
    
    
Yoach Lacombe's avatar
Yoach Lacombe committed
1271
1272
1273
    def clap_similarity(texts, audios, device):
        clap_inputs = clap_processor(text=texts, audios=audios, padding=True, return_tensors="pt").to(device)
        clap.to(device)
1274
1275
1276
1277
1278
        with torch.no_grad():
            text_features = clap.get_text_features(clap_inputs["input_ids"], attention_mask=clap_inputs.get("attention_mask", None))
            audio_features = clap.get_audio_features(clap_inputs["input_features"])
            
            cosine_sim = torch.nn.functional.cosine_similarity(audio_features, text_features, dim=1, eps=1e-8)
1279
        
Yoach Lacombe's avatar
Yoach Lacombe committed
1280
1281
        clap.to("cpu")
        clap_inputs.to("cpu")
1282
        return cosine_sim.mean().to("cpu")
1283
    
Yoach Lacombe's avatar
Yoach Lacombe committed
1284
1285
    def wer(prompts, audios, device):
        asr_pipeline = pipeline(model="distil-whisper/distil-large-v2", device=device)
Yoach Lacombe's avatar
Yoach Lacombe committed
1286
        transcriptions = asr_pipeline([{'raw': audio, 'sampling_rate': sampling_rate} for audio in audios], batch_size=int(training_args.per_device_eval_batch_size))
Yoach Lacombe's avatar
Yoach Lacombe committed
1287
1288
1289
1290
1291
1292
1293
        
        word_error = 100 * metric.compute(predictions=[t["text"].lower() for t in transcriptions], references=[t.lower() for t in prompts])
        
        return word_error, [t["text"] for t in transcriptions]
        
    
    eval_methods = {"clap": clap_similarity, "wer": wer}
1294

Yoach Lacombe's avatar
Yoach Lacombe committed
1295
1296
    def compute_metrics(audios, descriptions, prompts, device="cpu"):
        input_ids = descriptions
1297
        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
Yoach Lacombe's avatar
Yoach Lacombe committed
1298
1299
1300
1301
1302
1303
1304
        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
        audios = [a.cpu().numpy() for a in audios]
        results = {
            "clap": eval_methods["clap"](texts, audios, device)
        }
        word_error, transcriptions = eval_methods["wer"](prompts, audios, device)
        results["wer"] = word_error
1305

Yoach Lacombe's avatar
Yoach Lacombe committed
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
        return results, texts, prompts, audios, transcriptions
    
    # 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)
    
    
    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:
        logger.info(
            f"eval_steps is not set, evaluating at the end of each epoch"
        )
        eval_steps = steps_per_epoch
    else:
        eval_steps = training_args.eval_steps
        
1335
1336
1337
    # T5 doesn't support fp16
    autocast_kwargs = AutocastKwargs(enabled= (mixed_precision != "fp16"))
        
Yoach Lacombe's avatar
Yoach Lacombe committed
1338
1339
1340
1341
1342
1343
    # 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,
1344
        weight_decay=training_args.weight_decay,
Yoach Lacombe's avatar
Yoach Lacombe committed
1345
    )
1346

Yoach Lacombe's avatar
Yoach Lacombe committed
1347
1348
1349
1350
    # 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
1351
        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
Yoach Lacombe's avatar
Yoach Lacombe committed
1352
1353
        num_training_steps=total_train_steps * accelerator.num_processes,
    )
1354
1355

    # Instantiate custom data collator
Yoach Lacombe's avatar
Yoach Lacombe committed
1356
    data_collator = DataCollatorStableSpeechWithPadding(
1357
1358
        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
1359
1360
    )
    
Yoach Lacombe's avatar
Yoach Lacombe committed
1361
1362
1363
    
    # Prepare everything with accelerate
    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
Yoach Lacombe's avatar
Yoach Lacombe committed
1364
    
Yoach Lacombe's avatar
Yoach Lacombe committed
1365
1366
    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {total_train_steps * train_batch_size * gradient_accumulation_steps}")
1367
    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
Yoach Lacombe's avatar
Yoach Lacombe committed
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
    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
    
            
    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
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
    # 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
            if model_args.prompt_tokenizer_name is None and model_args.description_tokenizer_name or (model_args.prompt_tokenizer_name==model_args.description_tokenizer_name):
                prompt_tokenizer.save_pretrained(training_args.output_dir)
            else:
                logger.warning("Prompt tokenizer ('{model_args.prompt_tokenizer_name}') and description tokenizer ('{model_args.description_tokenizer_name}') are not the same. Saving only the prompt tokenizer.")
                prompt_tokenizer.save_pretrained(training_args.output_dir)
            
            feature_extractor.save_pretrained(training_args.output_dir)
            config.save_pretrained(training_args.output_dir)
            
    
Yoach Lacombe's avatar
Yoach Lacombe committed
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441

    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)
1442
            
Yoach Lacombe's avatar
Yoach Lacombe committed
1443
1444
        if training_args.max_steps < 0:
            # we know exactly the number of steps per epoch, so can skip through the required number of batches
1445
            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
        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
        
    gen_kwargs = {
        "do_sample": model_args.do_sample,
yoach@huggingface.co's avatar
yoach@huggingface.co committed
1457
        "temperature": model_args.temperature,
Yoach Lacombe's avatar
Yoach Lacombe committed
1458
1459
1460
1461
1462
1463
        "max_length": model_args.max_length,
    }
    
    # Define gradient update step fn
    def train_step(
        batch,
1464
1465
        accelerator,
        autocast_kwargs,
Yoach Lacombe's avatar
Yoach Lacombe committed
1466
1467
    ):
        model.train()
1468
        
1469
        if mixed_precision == "fp16":
1470
1471
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
1472
1473
1474
1475
                if training_args.parallel_mode.value != "distributed":
                    encoder_outputs = model.text_encoder(input_ids= batch.get("input_ids"), attention_mask=batch.get("attention_mask", None))
                else:
                    encoder_outputs = model.module.text_encoder(input_ids= batch.get("input_ids"), attention_mask=batch.get("attention_mask", None))
1476
1477
                batch["encoder_outputs"] = encoder_outputs
       
Yoach Lacombe's avatar
Yoach Lacombe committed
1478
1479
1480
1481
1482
1483
1484
1485
1486
        outputs = model(**batch)
        # CE (data) loss
        ce_loss = outputs.loss
        # TODO: add CE per codebook 

        metrics = {"loss": ce_loss}
        return ce_loss, metrics
    
    # Define eval fn
1487
    def eval_step(batch, accelerator, autocast_kwargs,):
Yoach Lacombe's avatar
Yoach Lacombe committed
1488
1489
1490
        eval_model = model if not training_args.torch_compile else model._orig_mod
        eval_model.eval()

1491
        if mixed_precision == "fp16":
1492
1493
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
Yoach Lacombe's avatar
Yoach Lacombe committed
1494
1495
1496
1497
1498
                with torch.no_grad():
                    if training_args.parallel_mode.value != "distributed" or training_args.torch_compile:
                        encoder_outputs = eval_model.text_encoder(input_ids= batch.get("input_ids"), attention_mask=batch.get("attention_mask", None))
                    else:
                        encoder_outputs = eval_model.module.text_encoder(input_ids= batch.get("input_ids"), attention_mask=batch.get("attention_mask", None))
1499
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
1500
1501

        with torch.no_grad():
Yoach Lacombe's avatar
Yoach Lacombe committed
1502
            outputs = eval_model(**batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
1503
1504
1505
1506
1507
1508
        # CE (data) loss
        ce_loss = outputs.loss
        metrics = {"loss": ce_loss}
        return metrics

    def generate_step(batch):
1509
        batch.pop("decoder_attention_mask", None)
Yoach Lacombe's avatar
Yoach Lacombe committed
1510
1511
1512
1513
1514
        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper = mixed_precision != "fp16").eval() 
        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
1515
1516
1517
1518
1519
        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)
1520
        # TODO(YL): add args
1521
        sampler = LengthGroupedSampler(train_batch_size, lengths = vectorized_datasets["train"]["target_length"])
Yoach Lacombe's avatar
Yoach Lacombe committed
1522
1523
1524
1525
        train_dataloader = DataLoader(
            vectorized_datasets["train"],
            collate_fn=data_collator,
            batch_size=per_device_train_batch_size,
Yoach Lacombe's avatar
Yoach Lacombe committed
1526
            sampler=sampler,
Yoach Lacombe's avatar
Yoach Lacombe committed
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
            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):
1541
                loss, train_metric = train_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
                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}")
1574
1575
1576
                    # 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
1577
1578
1579
1580
1581
1582
                    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
1583
1584
                            unwrapped_model = accelerator.unwrap_model(model)
                            unwrapped_model.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599

                        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
1600
1601
1602
                    
                    # release training input batch
                    batch = release_memory(batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
1603
1604
1605
1606
1607

                    validation_dataloader = DataLoader(
                        vectorized_datasets["eval"],
                        collate_fn=data_collator,
                        batch_size=per_device_eval_batch_size,
1608
                        drop_last=False,
Yoach Lacombe's avatar
Yoach Lacombe committed
1609
1610
1611
1612
1613
1614
1615
                        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,
1616
                        desc=f"Evaluating - Inference ...",
Yoach Lacombe's avatar
Yoach Lacombe committed
1617
1618
1619
1620
                        position=2,
                        disable=not accelerator.is_local_main_process,
                    ):
                        # Model forward
1621
                        eval_metric = eval_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
1622
1623
1624
                        eval_metric = accelerator.gather_for_metrics(eval_metric)
                        eval_metrics.append(eval_metric)

1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
                    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
1635
                        # generation
1636
1637
1638
1639
1640
1641
                        for batch in tqdm(
                                validation_dataloader,
                                desc=f"Evaluating - Generation ...",
                                position=2,
                                disable=not accelerator.is_local_main_process,
                            ):
Yoach Lacombe's avatar
Yoach Lacombe committed
1642
1643
1644
1645
                            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
1646
                            generated_audios, input_ids, prompts = accelerator.pad_across_processes((generated_audios, batch["input_ids"], batch["prompt_input_ids"]), dim=1, pad_index=0) 
1647
                            generated_audios, input_ids, prompts = accelerator.gather_for_metrics((generated_audios, input_ids, prompts))
1648
1649
1650
                            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
1651
1652
1653
1654

                    eval_time = time.time() - eval_start
                    # normalize eval metrics
                    eval_metrics = {
1655
                        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
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
                    }

                    # 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()])
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
                        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
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691

                    # 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",
                    )
1692
1693
1694
1695
1696
1697
1698
1699
                    
                    # 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
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713

                    # 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()
    
1714
1715
1716


if __name__ == "__main__":
1717
    set_start_method("spawn")
1718
    main()