sanchit-gandhi HF staff commited on
Commit
8415611
1 Parent(s): 73da21d

Add scripts and weights

Browse files
.gitattributes CHANGED
@@ -30,3 +30,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
30
  *.zip filter=lfs diff=lfs merge=lfs -text
31
  *.zst filter=lfs diff=lfs merge=lfs -text
32
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
30
  *.zip filter=lfs diff=lfs merge=lfs -text
31
  *.zst filter=lfs diff=lfs merge=lfs -text
32
  *tfevents* filter=lfs diff=lfs merge=lfs -text
33
+ *.whisper filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language:
3
+ - en
4
+ tags:
5
+ - esc
6
+ datasets:
7
+ - chime4
8
+ ---
9
+ To reproduce this run, execute:
10
+ ```python
11
+ #!/usr/bin/env bash
12
+ CUDA_VISIBLE_DEVICES=0 python run_speech_recognition_whisper.py \
13
+ --model_name_or_path="medium.en" \
14
+ --dataset_name="esc/esc-datasets" \
15
+ --dataset_config_name="chime4" \
16
+ --max_steps="2500" \
17
+ --output_dir="./" \
18
+ --run_name="whisper-chime4" \
19
+ --dropout_rate="0.1" \
20
+ --wandb_project="whisper" \
21
+ --per_device_train_batch_size="64" \
22
+ --per_device_eval_batch_size="16" \
23
+ --logging_steps="25" \
24
+ --learning_rate="1e-4" \
25
+ --warmup_steps="500" \
26
+ --report_to="wandb" \
27
+ --preprocessing_num_workers="16" \
28
+ --evaluation_strategy="steps" \
29
+ --eval_steps="500" \
30
+ --save_strategy="steps" \
31
+ --save_steps="500" \
32
+ --generation_max_length="224" \
33
+ --length_column_name="input_lengths" \
34
+ --gradient_checkpointing \
35
+ --group_by_length \
36
+ --freeze_encoder \
37
+ --fp16 \
38
+ --overwrite_output_dir \
39
+ --do_train \
40
+ --do_eval \
41
+ --do_predict \
42
+ --predict_with_generate \
43
+ --use_auth_token
44
+
45
+ ```
medium.en.whisper ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c1f25d9a74cf9fe130c73b9613cf9516fb774ce8c51bb8b88cf2fa2447d86bf1
3
+ size 3055771163
run_chime4.sh ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ CUDA_VISIBLE_DEVICES=0 python run_speech_recognition_whisper.py \
3
+ --model_name_or_path="medium.en" \
4
+ --dataset_name="esc/esc-datasets" \
5
+ --dataset_config_name="chime4" \
6
+ --max_steps="2500" \
7
+ --output_dir="./" \
8
+ --run_name="whisper-chime4" \
9
+ --dropout_rate="0.1" \
10
+ --wandb_project="whisper" \
11
+ --per_device_train_batch_size="64" \
12
+ --per_device_eval_batch_size="16" \
13
+ --logging_steps="25" \
14
+ --learning_rate="1e-4" \
15
+ --warmup_steps="500" \
16
+ --report_to="wandb" \
17
+ --preprocessing_num_workers="16" \
18
+ --evaluation_strategy="steps" \
19
+ --eval_steps="500" \
20
+ --save_strategy="steps" \
21
+ --save_steps="500" \
22
+ --generation_max_length="224" \
23
+ --length_column_name="input_lengths" \
24
+ --gradient_checkpointing \
25
+ --group_by_length \
26
+ --freeze_encoder \
27
+ --fp16 \
28
+ --overwrite_output_dir \
29
+ --do_train \
30
+ --do_eval \
31
+ --do_predict \
32
+ --predict_with_generate \
33
+ --use_auth_token
run_speech_recognition_whisper.py ADDED
@@ -0,0 +1,747 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ # coding=utf-8
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """
16
+ Fine-tuning OpenAI Whisper models for speech recognition.
17
+ """
18
+ # You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
19
+ # flake8: noqa: E501
20
+ import logging
21
+ import os
22
+
23
+ import whisper
24
+ import sys
25
+ from dataclasses import dataclass, field
26
+ import tempfile
27
+
28
+ from typing import Optional, Dict, Union, List
29
+
30
+ import numpy as np
31
+ import torch
32
+
33
+ import datasets
34
+ from datasets import DatasetDict, load_dataset
35
+ import transformers
36
+ from torch import nn
37
+ from transformers import (
38
+ HfArgumentParser,
39
+ Seq2SeqTrainingArguments,
40
+ set_seed,
41
+ Seq2SeqTrainer,
42
+ )
43
+ from transformers.trainer_utils import get_last_checkpoint, is_main_process
44
+ from transformers.utils import check_min_version
45
+ from transformers.utils.versions import require_version
46
+
47
+ import wandb
48
+
49
+ # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
50
+ check_min_version("4.17.0.dev0")
51
+
52
+ require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt")
53
+
54
+ logger = logging.getLogger(__name__)
55
+
56
+
57
+ @dataclass
58
+ class ModelArguments:
59
+ """
60
+ Arguments pertaining to which model/tokenizer we are going to fine-tune from.
61
+ """
62
+ model_name_or_path: Optional[str] = field(
63
+ default=None,
64
+ metadata={"help": "Path to pretrained model or model identifier from OpenAI Whisper NGC."}
65
+ )
66
+ cache_dir: Optional[str] = field(
67
+ default=None,
68
+ metadata={"help": "Where to store the pretrained models downloaded from huggingface.co or OpenAI Whisper NGC."},
69
+ )
70
+ use_auth_token: bool = field(
71
+ default=False,
72
+ metadata={
73
+ "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
74
+ "with private models)."
75
+ },
76
+ )
77
+ manifest_path: str = field(
78
+ default="data",
79
+ metadata={
80
+ "help": "Manifest path."
81
+ },
82
+ )
83
+ tokenizer_path: str = field(
84
+ default="tokenizers",
85
+ metadata={
86
+ "help": "Tokenizer path."
87
+ },
88
+ )
89
+ freeze_encoder: bool = field(
90
+ default=False,
91
+ metadata={"help": "Freeze the acoustic encoder of the model. Recommend when fine-tuning on small datasets."}
92
+ )
93
+ num_beams: int = field(
94
+ default=1,
95
+ metadata={"help": "Number of beams for evaluation."},
96
+ )
97
+ length_penalty: float = field(
98
+ default=1.0,
99
+ metadata={"help": "Length penalty for evaluation."},
100
+ )
101
+ use_adam8bit: bool = field(
102
+ default=False,
103
+ metadata={"help": "Whether to use bitsandbytes 8bit AdamW optimiser."}
104
+ )
105
+ dropout_rate: float = field(
106
+ default=0.0,
107
+ metadata={"help": "The dropout ratio for all dropout layers (default=0)."}
108
+ )
109
+
110
+
111
+ @dataclass
112
+ class DataTrainingArguments:
113
+ """
114
+ Arguments pertaining to what data we are going to input our model for training and eval.
115
+ """
116
+
117
+ dataset_name: str = field(
118
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
119
+ )
120
+ dataset_config_name: Optional[str] = field(
121
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
122
+ )
123
+ text_column: Optional[str] = field(
124
+ default=None,
125
+ metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
126
+ )
127
+ dataset_cache_dir: Optional[str] = field(
128
+ default=None, metadata={"help": "Path to cache directory for saving and loading datasets"}
129
+ )
130
+ overwrite_cache: bool = field(
131
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
132
+ )
133
+ preprocessing_num_workers: Optional[int] = field(
134
+ default=None,
135
+ metadata={"help": "The number of processes to use for the preprocessing."},
136
+ )
137
+ max_train_samples: Optional[int] = field(
138
+ default=None,
139
+ metadata={
140
+ "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
141
+ "value if set."
142
+ },
143
+ )
144
+ max_eval_samples: Optional[int] = field(
145
+ default=None,
146
+ metadata={
147
+ "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
148
+ "value if set."
149
+ },
150
+ )
151
+ max_predict_samples: Optional[int] = field(
152
+ default=None,
153
+ metadata={
154
+ "help": "For debugging purposes or quicker training, truncate the number of test examples to this "
155
+ "value if set."
156
+ },
157
+ )
158
+ audio_column_name: str = field(
159
+ default="audio",
160
+ metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},
161
+ )
162
+ text_column_name: str = field(
163
+ default="text",
164
+ metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},
165
+ )
166
+ max_duration_in_seconds: float = field(
167
+ default=20.0,
168
+ metadata={
169
+ "help": "Truncate training audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_seconds`"
170
+ },
171
+ )
172
+ min_duration_in_seconds: float = field(
173
+ default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}
174
+ )
175
+ max_eval_duration_in_seconds: float = field(
176
+ default=None,
177
+ metadata={
178
+ "help": "Truncate eval/test audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_seconds`"
179
+ },
180
+ )
181
+ max_target_length: Optional[int] = field(
182
+ default=128,
183
+ metadata={
184
+ "help": "The maximum total sequence length for target text after tokenization. Sequences longer "
185
+ "than this will be truncated, sequences shorter will be padded."
186
+ },
187
+ )
188
+ min_target_length: Optional[int] = field(
189
+ default=0,
190
+ metadata={
191
+ "help": "The minimum total sequence length for target text after tokenization. Sequences shorter "
192
+ "than this will be filtered."
193
+ },
194
+ )
195
+ preprocessing_only: bool = field(
196
+ default=False,
197
+ metadata={
198
+ "help": "Whether to only do data preprocessing and skip training. "
199
+ "This is especially useful when data preprocessing errors out in distributed training due to timeout. "
200
+ "In this case, one should run the preprocessing in a non-distributed setup with `preprocessing_only=True` "
201
+ "so that the cached datasets can consequently be loaded in distributed training"
202
+ },
203
+ )
204
+ train_split_name: str = field(
205
+ default="train",
206
+ metadata={
207
+ "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
208
+ },
209
+ )
210
+ eval_split_name: str = field(
211
+ default="validation",
212
+ metadata={
213
+ "help": "The name of the evaluation data set split to use (via the datasets library). Defaults to 'validation'"
214
+ },
215
+ )
216
+ test_split_name: str = field(
217
+ default="test",
218
+ metadata={"help": "The name of the test data set split to use (via the datasets library). Defaults to 'test'"},
219
+ )
220
+ wandb_project: str = field(
221
+ default="speech-recognition-whisper",
222
+ metadata={"help": "The name of the wandb project."},
223
+ )
224
+
225
+
226
+ def write_wandb_pred(pred_str, label_str, prefix="eval"):
227
+ # convert str data to a wandb compatible format
228
+ str_data = [[label_str[i], pred_str[i]] for i in range(len(pred_str))]
229
+ # we'll log all predictions for the last epoch
230
+ wandb.log(
231
+ {
232
+ f"{prefix}/predictions": wandb.Table(
233
+ columns=["label_str", "pred_str"], data=str_data
234
+ )
235
+ },
236
+ )
237
+
238
+
239
+ def to_pad_to_mel(array):
240
+ """Static function which:
241
+ 1. Pads/trims a list of audio arrays to a max length of 30s
242
+ 2. Computes log-mel filter coefficients from padded/trimmed audio sequences
243
+ Inputs:
244
+ array: list of audio arrays
245
+ Returns:
246
+ input_ids: torch.tensor of log-mel filter bank coefficients
247
+ """
248
+ padded_input = whisper.pad_or_trim(np.asarray(array, dtype=np.float32))
249
+ input_ids = whisper.log_mel_spectrogram(padded_input)
250
+ return input_ids
251
+
252
+
253
+ def to_mel_to_pad(array):
254
+ """Static function which:
255
+ 1. Computes log-mel filter coefficients from padded/trimmed audio sequences
256
+ 2. Pads/trims a list of audio arrays to a max length of 30s
257
+ Inputs:
258
+ array: list of audio arrays
259
+ Returns:
260
+ input_ids: torch.tensor of log-mel filter bank coefficients
261
+ """
262
+ mels = whisper.log_mel_spectrogram(np.asarray(array, dtype=np.float32))
263
+ input_ids = whisper.pad_or_trim(mels, 3000)
264
+ return input_ids
265
+
266
+
267
+ @dataclass
268
+ class WhisperDataCollatorWithPadding:
269
+ """
270
+ Data collator that dynamically pads the audio inputs received. An EOS token is appended to the labels sequences.
271
+ They are then dynamically padded to max length.
272
+ Args:
273
+ eos_token_id (`int`)
274
+ The end-of-sentence token for the Whisper tokenizer. Ensure to set for sequences to terminate before
275
+ generation max length.
276
+ """
277
+
278
+ eos_token_id: int
279
+
280
+ def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
281
+ """
282
+ Since Whisper models don't have a HF processor defined (feature extractor + tokenizer), we'll pad by hand...
283
+ """
284
+ # split inputs and labels since they have to be of different lengths
285
+ # and need different padding methods
286
+ input_ids = [feature["input_ids"] for feature in features]
287
+ labels = [feature["labels"] for feature in features]
288
+
289
+ # first, pad the audio inputs to max_len
290
+ input_ids = torch.concat([to_pad_to_mel(input_val)[None, :] for input_val in input_ids])
291
+
292
+ # next, append the eos token to our sequence of labels
293
+ labels = [lab + [self.eos_token_id] for lab in labels]
294
+ # finally, pad the target labels to max_len
295
+ label_lengths = [len(lab) for lab in labels]
296
+ max_label_len = max(label_lengths)
297
+ labels = [np.pad(lab, (0, max_label_len - lab_len), 'constant', constant_values=-100) for lab, lab_len in zip(labels, label_lengths)]
298
+
299
+ batch = {"labels": labels}
300
+ batch = {k: torch.tensor(np.array(v), requires_grad=False) for k, v in batch.items()}
301
+
302
+ batch["input_ids"] = input_ids
303
+
304
+ return batch
305
+
306
+
307
+ def main():
308
+ # See all possible arguments in src/transformers/training_args.py
309
+ # or by passing the --help flag to this script.
310
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
311
+
312
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
313
+ if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
314
+ # If we pass only one argument to the script and it's the path to a json file,
315
+ # let's parse it to get our arguments.
316
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
317
+ else:
318
+ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
319
+
320
+ # Set wandb project ID before instantiating the Trainer
321
+ os.environ["WANDB_PROJECT"] = data_args.wandb_project
322
+ report_to_wandb = "wandb" in training_args.report_to
323
+
324
+ sample_rate = 16_000
325
+
326
+ # Detecting last checkpoint.
327
+ last_checkpoint = None
328
+ if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
329
+ last_checkpoint = get_last_checkpoint(training_args.output_dir)
330
+ if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
331
+ raise ValueError(
332
+ f"Output directory ({training_args.output_dir}) already exists and is not empty. "
333
+ "Use --overwrite_output_dir to overcome."
334
+ )
335
+ elif last_checkpoint is not None:
336
+ logger.info(
337
+ f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
338
+ "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
339
+ )
340
+
341
+ # Setup logging
342
+ logging.basicConfig(
343
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
344
+ datefmt="%m/%d/%Y %H:%M:%S",
345
+ handlers=[logging.StreamHandler(sys.stdout)],
346
+ )
347
+ logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)
348
+
349
+ # Log on each process the small summary:
350
+ logger.warning(
351
+ f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
352
+ f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
353
+ )
354
+ # Set the verbosity to info of the Transformers logger (on main process only):
355
+ if is_main_process(training_args.local_rank):
356
+ transformers.utils.logging.set_verbosity_info()
357
+ logger.info("Training/evaluation parameters %s", training_args)
358
+
359
+ # Set seed before initializing model.
360
+ set_seed(training_args.seed)
361
+
362
+ # load the model
363
+ if os.path.isfile(model_args.model_name_or_path):
364
+ checkpoint = torch.load(model_args.model_name_or_path)
365
+ need_to_rewrite_checkpoint = any(k.startswith("decoder.blocks") and ".mlp.3" in k for k in checkpoint.keys())
366
+ if need_to_rewrite_checkpoint:
367
+ new_checkpoint = {}
368
+ for k, v in checkpoint.items():
369
+ if k.startswith("decoder.blocks") and "mlp" in k.split("."):
370
+ if int(k.split(".mlp.")[-1].split(".")[0]) in [2, 4]:
371
+ continue
372
+ elif int(k.split(".mlp.")[-1].split(".")[0]) == 3:
373
+ k = k.replace(".mlp.3", ".mlp.2")
374
+
375
+ new_checkpoint[k] = v
376
+
377
+ with tempfile.TemporaryDirectory() as tmp:
378
+ file = os.path.join(tmp, "model.pt")
379
+ torch.save(new_checkpoint, file)
380
+ model = whisper.Whisper.load_trained(file)
381
+ else:
382
+ model = whisper.Whisper.load_trained(model_args.model_name_or_path)
383
+ del checkpoint
384
+ else:
385
+ model = whisper.load_model(model_args.model_name_or_path, dropout_rate=model_args.dropout_rate)
386
+
387
+ if training_args.do_train:
388
+ # set the dropout for the MLP layers -> we do this here as the MLP layers are written as a 'sequential'
389
+ # so changing the modelling code gives mis-matches in the state-dict
390
+ if not model_args.freeze_encoder:
391
+ # only apply dropout when training the encoder
392
+ for block_idx in range(len(model.encoder.blocks)):
393
+ mlp_layer = model.encoder.blocks[block_idx].mlp
394
+ # going very verbose to explain what we're doing here!
395
+ fc1 = mlp_layer[0]
396
+ act_fn = mlp_layer[1]
397
+ dropout = nn.Dropout(p=model_args.dropout_rate)
398
+ fc2 = mlp_layer[2]
399
+ model.encoder.blocks[block_idx].mlp = nn.Sequential(fc1, act_fn, dropout, fc2, dropout)
400
+
401
+ for block_idx in range(len(model.decoder.blocks)):
402
+ mlp_layer = model.decoder.blocks[block_idx].mlp
403
+ fc1 = mlp_layer[0]
404
+ act_fn = mlp_layer[1]
405
+ dropout_1 = nn.Dropout(p=model_args.dropout_rate)
406
+ fc2 = mlp_layer[2]
407
+ dropout_2 = nn.Dropout(p=model_args.dropout_rate)
408
+ model.decoder.blocks[block_idx].mlp = nn.Sequential(fc1, act_fn, dropout_1, fc2, dropout_2)
409
+ for block_idx in range(len(model.decoder.blocks)):
410
+ mlp_layer = model.decoder.blocks[block_idx].mlp
411
+ fc1 = mlp_layer[0]
412
+ act_fn = mlp_layer[1]
413
+ dropout1 = nn.Dropout(p=model_args.dropout_rate)
414
+ fc2 = mlp_layer[2]
415
+ dropout2 = nn.Dropout(p=model_args.dropout_rate)
416
+ model.decoder.blocks[block_idx].mlp = nn.Sequential(fc1, act_fn, dropout1, fc2, dropout2)
417
+
418
+ # load the tokenizer
419
+ whisper_tok = whisper.tokenizer.get_tokenizer(False, task="transcribe", language="en")
420
+ tokenizer = whisper_tok.tokenizer
421
+ tokenizer.pad_token = tokenizer.eos_token
422
+
423
+ # 4. Load dataset
424
+ raw_datasets = DatasetDict()
425
+
426
+ if training_args.do_train:
427
+ raw_datasets["train"] = load_dataset(
428
+ data_args.dataset_name,
429
+ data_args.dataset_config_name,
430
+ split=data_args.train_split_name,
431
+ cache_dir=data_args.dataset_cache_dir,
432
+ use_auth_token=True if model_args.use_auth_token else None,
433
+ )
434
+
435
+ if training_args.do_eval:
436
+ raw_datasets["eval"] = load_dataset(
437
+ data_args.dataset_name,
438
+ data_args.dataset_config_name,
439
+ split=data_args.eval_split_name,
440
+ cache_dir=data_args.dataset_cache_dir,
441
+ use_auth_token=True if model_args.use_auth_token else None,
442
+ )
443
+
444
+ if training_args.do_predict:
445
+ test_split = data_args.test_split_name.split("+")
446
+ for split in test_split:
447
+ raw_datasets[split] = load_dataset(
448
+ data_args.dataset_name,
449
+ data_args.dataset_config_name,
450
+ split=split,
451
+ cache_dir=data_args.dataset_cache_dir,
452
+ use_auth_token=True if model_args.use_auth_token else None,
453
+ )
454
+
455
+ if not training_args.do_train and not training_args.do_eval and not training_args.do_predict:
456
+ raise ValueError(
457
+ "Cannot not train, not do evaluation and not do prediction. At least one of "
458
+ "training, evaluation or prediction has to be done."
459
+ )
460
+
461
+ # if not training, there is no need to run multiple epochs
462
+ if not training_args.do_train:
463
+ training_args.num_train_epochs = 1
464
+
465
+ if data_args.audio_column_name not in next(iter(raw_datasets.values())).column_names:
466
+ raise ValueError(
467
+ f"--audio_column_name '{data_args.audio_column_name}' not found in dataset '{data_args.dataset_name}'. "
468
+ "Make sure to set `--audio_column_name` to the correct audio column - one of "
469
+ f"{', '.join(next(iter(raw_datasets.values())).column_names)}."
470
+ )
471
+
472
+ if data_args.text_column_name not in next(iter(raw_datasets.values())).column_names:
473
+ raise ValueError(
474
+ f"--text_column_name {data_args.text_column_name} not found in dataset '{data_args.dataset_name}'. "
475
+ "Make sure to set `--text_column_name` to the correct text column - one of "
476
+ f"{', '.join(next(iter(raw_datasets.values())).column_names)}."
477
+ )
478
+
479
+ # 6. Resample speech dataset ALWAYS
480
+ raw_datasets = raw_datasets.cast_column(
481
+ data_args.audio_column_name, datasets.features.Audio(sampling_rate=sample_rate)
482
+ )
483
+
484
+ # 7. Preprocessing the datasets.
485
+ # We need to read the audio files as arrays and tokenize the targets.
486
+ max_input_length = int(data_args.max_duration_in_seconds * sample_rate)
487
+ min_input_length = min(int(data_args.min_duration_in_seconds * sample_rate), 1)
488
+ max_eval_input_length = int(data_args.max_eval_duration_in_seconds * sample_rate) if data_args.max_eval_duration_in_seconds else None
489
+ max_target_length = data_args.max_target_length
490
+ min_target_length = data_args.min_target_length
491
+ audio_column_name = data_args.audio_column_name
492
+ num_workers = data_args.preprocessing_num_workers
493
+ text_column_name = data_args.text_column_name
494
+
495
+ if training_args.do_train and data_args.max_train_samples is not None:
496
+ raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))
497
+
498
+ if training_args.do_eval and data_args.max_eval_samples is not None:
499
+ raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))
500
+
501
+ if training_args.do_predict and data_args.max_predict_samples is not None:
502
+ for split in test_split:
503
+ raw_datasets[split] = raw_datasets[split].select(range(data_args.max_predict_samples))
504
+
505
+
506
+ def prepare_dataset(batch):
507
+ # pre-process audio
508
+ sample = batch[audio_column_name]
509
+
510
+ # For training Whisper we perform the audio preprocessing in the WhisperDataCollator
511
+ # => we only need to supply it with the raw audio values
512
+ batch["input_ids"] = sample["array"]
513
+ batch["input_lengths"] = len(batch["input_ids"])
514
+
515
+ input_str = batch[text_column_name]
516
+ batch["labels"] = tokenizer(input_str).input_ids
517
+ return batch
518
+
519
+ vectorized_datasets = raw_datasets.map(
520
+ prepare_dataset,
521
+ remove_columns=next(iter(raw_datasets.values())).column_names,
522
+ num_proc=num_workers,
523
+ desc="preprocess train dataset",
524
+ )
525
+
526
+ # filter training data with inputs longer than max_input_length
527
+ def is_audio_in_length_range(input_length):
528
+ return min_input_length < input_length < max_input_length
529
+
530
+ if training_args.do_train:
531
+ vectorized_datasets["train"] = vectorized_datasets["train"].filter(
532
+ is_audio_in_length_range,
533
+ num_proc=num_workers,
534
+ input_columns=["input_lengths"],
535
+ )
536
+
537
+ if max_eval_input_length is not None:
538
+ # filter training data with inputs longer than max_input_length
539
+ def is_eval_audio_in_length_range(input_length):
540
+ return min_input_length < input_length < max_eval_input_length
541
+
542
+ if training_args.do_eval:
543
+ vectorized_datasets["eval"] = vectorized_datasets["eval"].filter(
544
+ is_eval_audio_in_length_range,
545
+ num_proc=num_workers,
546
+ input_columns=["input_lengths"],
547
+ )
548
+
549
+ if training_args.do_predict:
550
+ for split in test_split:
551
+ vectorized_datasets[split] = vectorized_datasets[split].filter(
552
+ is_eval_audio_in_length_range,
553
+ num_proc=num_workers,
554
+ input_columns=["input_lengths"],
555
+ )
556
+
557
+ # filter training data with targets shorter than min_target_length or longer than max_target_length
558
+ def is_labels_in_length_range(labels):
559
+ return min_target_length < len(labels) < max_target_length
560
+
561
+ if training_args.do_train:
562
+ vectorized_datasets["train"] = vectorized_datasets["train"].filter(
563
+ is_labels_in_length_range,
564
+ num_proc=num_workers,
565
+ input_columns=["labels"],
566
+ )
567
+
568
+ # for large datasets it is advised to run the preprocessing on a
569
+ # single machine first with `args.preprocessing_only` since there will mostly likely
570
+ # be a timeout when running the script in distributed mode.
571
+ # In a second step `args.preprocessing_only` can then be set to `False` to load the
572
+ # cached dataset
573
+ if data_args.preprocessing_only:
574
+ cache = {k: v.cache_files for k, v in vectorized_datasets.items()}
575
+ logger.info(f"Data preprocessing finished. Files cached at {cache}.")
576
+ return
577
+
578
+ if model_args.freeze_encoder:
579
+ model.freeze_encoder()
580
+ logging.info("Model encoder has been frozen")
581
+
582
+ # 8. Load Metric
583
+ metric_wer = datasets.load_metric("wer")
584
+ metric_cer = datasets.load_metric("cer")
585
+
586
+ def compute_metrics(pred):
587
+ pred_ids = pred.predictions
588
+ pred.label_ids[pred.label_ids == -100] = tokenizer.eos_token_id
589
+
590
+ pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
591
+ pred_str = [x.lstrip().strip() for x in pred_str]
592
+
593
+ # we do not want to group tokens when computing the metrics
594
+ label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True)
595
+
596
+ wer = metric_wer.compute(predictions=pred_str, references=label_str)
597
+ cer = metric_cer.compute(predictions=pred_str, references=label_str)
598
+
599
+ return {"wer": wer, "cer": cer}
600
+
601
+ def compute_metrics_and_predictions(pred):
602
+ pred_ids = pred.predictions
603
+ pred.label_ids[pred.label_ids == -100] = tokenizer.eos_token_id
604
+
605
+ pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
606
+ pred_str = [x.lstrip().strip() for x in pred_str]
607
+
608
+ # we do not want to group tokens when computing the metrics
609
+ label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True)
610
+
611
+ wer = metric_wer.compute(predictions=pred_str, references=label_str)
612
+ cer = metric_cer.compute(predictions=pred_str, references=label_str)
613
+
614
+ return {"wer": wer, "cer": cer, "pred_str": pred_str, "label_str": label_str}
615
+
616
+ class WhisperTrainer(Seq2SeqTrainer):
617
+ def _save(self, output_dir: Optional[str] = None, state_dict=None):
618
+ # If we are executing this function, we are the process zero, so we don't check for that.
619
+ output_dir = output_dir if output_dir is not None else self.args.output_dir
620
+ os.makedirs(output_dir, exist_ok=True)
621
+ logger.info(f"Saving model checkpoint to {output_dir}")
622
+ # Save a trained model and configuration using `save_pretrained()`.
623
+ # They can then be reloaded using `from_pretrained()`
624
+ self.model.save_to(save_path=os.path.join(output_dir, model_args.model_name_or_path + ".whisper"))
625
+ # Good practice: save your training arguments together with the trained model
626
+ torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
627
+
628
+ # Define data collator
629
+ whisper_data_collator = WhisperDataCollatorWithPadding(eos_token_id=tokenizer.eos_token_id)
630
+
631
+ # make sure model uses 50257 as BOS
632
+ bos = tokenizer("<|startoftranscript|>").input_ids[0]
633
+ model.config.decoder_start_token_id = bos
634
+
635
+ # Initialize Trainer
636
+ trainer = WhisperTrainer(
637
+ model=model,
638
+ args=training_args,
639
+ compute_metrics=compute_metrics,
640
+ train_dataset=vectorized_datasets['train'] if training_args.do_train else None,
641
+ eval_dataset=vectorized_datasets['eval'] if training_args.do_eval else None,
642
+ data_collator=whisper_data_collator,
643
+ )
644
+
645
+ # 8. Finally, we can start training
646
+
647
+ # Training
648
+ if training_args.do_train:
649
+
650
+ # use last checkpoint if exist
651
+ if last_checkpoint is not None:
652
+ checkpoint = last_checkpoint
653
+ elif model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path):
654
+ checkpoint = model_args.model_name_or_path
655
+ else:
656
+ checkpoint = None
657
+
658
+ train_result = trainer.train(resume_from_checkpoint=checkpoint)
659
+ trainer.save_model()
660
+
661
+ metrics = train_result.metrics
662
+ max_train_samples = (
663
+ data_args.max_train_samples
664
+ if data_args.max_train_samples is not None
665
+ else len(vectorized_datasets["train"])
666
+ )
667
+ metrics["train_samples"] = min(max_train_samples, len(vectorized_datasets["train"]))
668
+
669
+ trainer.log_metrics("train", metrics)
670
+ trainer.save_metrics("train", metrics)
671
+ trainer.save_state()
672
+
673
+ # Change decoding strategy for final eval/predict
674
+ # if training_args.do_eval or training_args.do_predict:
675
+ # trainer.model.num_beams = 2
676
+
677
+ trainer.compute_metrics = compute_metrics_and_predictions
678
+
679
+ results = {}
680
+ if training_args.do_eval:
681
+ if not training_args.do_train and report_to_wandb:
682
+ # manually init wandb
683
+ wandb.init(project=data_args.wandb_project, name=training_args.run_name)
684
+ # Have to run this as a predict step, otherwise trainer will try to log the pred/label strings to wandb
685
+ eval_results = trainer.predict(vectorized_datasets["eval"], metric_key_prefix="eval", num_beams=model_args.num_beams, length_penalty=model_args.length_penalty)
686
+ metrics = eval_results.metrics
687
+ max_eval_samples = (
688
+ data_args.max_eval_samples if data_args.max_eval_samples is not None else len(vectorized_datasets["eval"])
689
+ )
690
+ metrics["eval_samples"] = min(max_eval_samples, len(vectorized_datasets["eval"]))
691
+ pred_str = metrics.pop("eval_pred_str", None)
692
+ label_str = metrics.pop("eval_label_str", None)
693
+
694
+ trainer.log_metrics("eval", metrics)
695
+ trainer.save_metrics("eval", metrics)
696
+
697
+ if report_to_wandb:
698
+ metrics = {os.path.join("eval", k[len("eval") + 1:]): v for k, v in metrics.items()}
699
+ wandb.log(metrics)
700
+ write_wandb_pred(pred_str, label_str, prefix="eval")
701
+
702
+ if training_args.do_predict:
703
+ if not training_args.do_train and not training_args.do_eval and report_to_wandb:
704
+ # manually init wandb
705
+ wandb.init(project=data_args.wandb_project, name=training_args.run_name)
706
+ for split in test_split:
707
+ predict_results = trainer.predict(
708
+ vectorized_datasets[split], metric_key_prefix=split, num_beams=model_args.num_beams, length_penalty=model_args.length_penalty)
709
+ metrics = predict_results.metrics
710
+ max_predict_samples = (
711
+ data_args.max_predict_samples if data_args.max_predict_samples is not None else len(vectorized_datasets[split])
712
+ )
713
+ metrics[f"{split}_samples"] = min(max_predict_samples, len(vectorized_datasets[split]))
714
+ pred_str = metrics.pop(f"{split}_pred_str", None)
715
+ label_str = metrics.pop(f"{split}_label_str", None)
716
+
717
+ trainer.log_metrics(split, metrics)
718
+ trainer.save_metrics(split, metrics)
719
+
720
+ if report_to_wandb:
721
+ metrics = {os.path.join(split, k[len(split)+1:]): v for k, v in metrics.items()}
722
+ wandb.log(metrics)
723
+ write_wandb_pred(pred_str, label_str, prefix=split)
724
+
725
+ # Write model card and (optionally) push to hub
726
+ config_name = data_args.dataset_config_name if data_args.dataset_config_name is not None else "na"
727
+ kwargs = {
728
+ "finetuned_from": model_args.model_name_or_path,
729
+ "tasks": "speech-recognition",
730
+ "tags": ["automatic-speech-recognition", data_args.dataset_name],
731
+ "dataset_args": (
732
+ f"Config: {config_name}, Training split: {data_args.train_split_name}, Eval split:"
733
+ f" {data_args.eval_split_name}"
734
+ ),
735
+ "dataset": f"{data_args.dataset_name.upper()} - {config_name.upper()}",
736
+ }
737
+ if "common_voice" in data_args.dataset_name:
738
+ kwargs["language"] = config_name
739
+
740
+ if training_args.push_to_hub:
741
+ trainer.push_to_hub(**kwargs)
742
+
743
+ return results
744
+
745
+
746
+ if __name__ == "__main__":
747
+ main()