Upload student-teacher-distillation-streaming.py

student-teacher-distillation-streaming.py

@@ -0,0 +1,426 @@
+# Run teacher student distillation for Whisper model
+from transformers import AutoTokenizer, AutoProcessor, AutoConfig, AutoFeatureExtractor, AutoTokenizer, AutoProcessor, \
+    AutoModelForSpeechSeq2Seq, set_seed, get_linear_schedule_with_warmup
+from datasets import load_dataset, DatasetDict, interleave_datasets, IterableDatasetDict
+from transformers.models.whisper.english_normalizer import BasicTextNormalizer
+import transformers
+import argparse
+import datasets
+import evaluate
+import string
+from accelerate import Accelerator
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+import torch
+import os
+from tqdm.auto import tqdm
+import numpy as np
+
+
+def load_streaming_dataset(dataset_name, dataset_config_name, split, **kwargs):
+    if "+" in split:
+        # load multiple splits separated by the `+` symbol *with* streaming mode
+        dataset_splits = [load_dataset(dataset_name, dataset_config_name, split=split_name, streaming=True, **kwargs) for split_name in split.split("+")]
+        # interleave multiple splits to form one dataset
+        interleaved_dataset = interleave_datasets(dataset_splits)
+        return interleaved_dataset
+    else:
+        # load a single split *with* streaming mode
+        dataset = load_dataset(dataset_name, dataset_config_name, split=split, streaming=True, **kwargs)
+        return dataset
+    
+
+def train(args, accelerator: Accelerator):
+    # load dataset in streaming mode
+    raw_datasets = IterableDatasetDict()
+    raw_datasets["train"] = load_streaming_dataset(args.train_dataset_name, args.train_dataset_config_name,
+                                                   split=args.train_split_name,
+                                                   cache_dir=args.data_cache_dir)
+    raw_datasets["validation"] = load_streaming_dataset(args.validation_dataset_name,
+                                                        args.validation_dataset_config_name,
+                                                        split=args.validation_split_name,
+                                                        cache_dir=args.data_cache_dir) 
+    
+    # raw_datasets = raw_datasets.remove_columns(["file", "speaker_id", "chapter_id", "id"])
+    # raw_datasets = raw_datasets.rename_columns({'audio': 'audio', 'text': 'text'})
+    
+    assert args.audio_column in raw_datasets["train"].column_names
+    assert args.text_column in raw_datasets["train"].column_names
+    
+    with accelerator.main_process_first():
+        if args.max_train_samples is not None:
+            raw_datasets["train"] = raw_datasets["train"].take(args.max_train_samples)
+        if args.max_val_samples is not None:
+            raw_datasets["validation"] = raw_datasets["validation"].take(args.max_val_samples)
+    
+    student_config = AutoConfig.from_pretrained(args.student_model_name_or_path, cache_dir=args.student_cache_dir)
+    teacher_config = AutoConfig.from_pretrained(args.teacher_model_name_or_path, cache_dir=args.teacher_cache_dir)
+    
+    # assuming student and teacher uses same feature extractor, tokenizer and processor
+    feature_extractor = AutoFeatureExtractor.from_pretrained(args.teacher_model_name_or_path, cache_dir=args.teacher_cache_dir)
+    tokenizer = AutoTokenizer.from_pretrained(args.teacher_model_name_or_path, cache_dir=args.teacher_cache_dir)
+    processor = AutoProcessor.from_pretrained(args.teacher_model_name_or_path, cache_dir=args.teacher_cache_dir)
+    
+    # make sure decoder_start_token_id is defined for both
+    assert teacher_config.decoder_start_token_id is not None
+    assert student_config.decoder_start_token_id is not None
+    
+    # We need to set the language and task ids for previously multilingual checkpoints, default is English and transcribe
+    # Set to None if the model is not multilingual
+    student_config.forced_decoder_ids = None
+    # tokenizer.get_decoder_prompt_ids(language=args.language, task=args.task, no_timestamps=True)
+    teacher_config.forced_decoder_ids = None
+    # tokenizer.get_decoder_prompt_ids(language=args.language, task=args.task, no_timestamps=True)
+    student_config.suppress_tokens = []
+    teacher_config.suppress_tokens = []
+    
+    student_model = AutoModelForSpeechSeq2Seq.from_pretrained(args.student_model_name_or_path, config=student_config)
+    teacher_model = AutoModelForSpeechSeq2Seq.from_pretrained(args.teacher_model_name_or_path, config=teacher_config,
+                                                              cache_dir=args.teacher_cache_dir)
+    
+    accelerator.print(
+        f"Loaded the model on device: student: {student_model.device}, teacher:{teacher_model.device}, accelerator:{accelerator.device}")
+    
+    # freeze teacher model
+    for p in teacher_model.parameters():
+        p.requires_grad = False
+
+    if args.freeze_encoder:
+        accelerator.print("Freezing encoder")
+        student_model.freeze_encoder()
+        student_model.model.encoder.gradient_checkpointing = False
+    
+    # Resample speech dataset: so we just need to set the correct target sampling rate
+    with accelerator.main_process_first():
+        # raw_datasets = raw_datasets.cast_column(args.audio_column,
+        #                                         datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate))
+        raw_datasets = raw_datasets.cast_column(args.audio_column, datasets.Audio(sampling_rate=16000))
+
+    # Preprocessing the raw_datasets, need to read the audio files as arrays and tokenize the targets.
+    # might need to change the normalizer depending on language and task
+    normalizer = BasicTextNormalizer()
+    
+    def prepare_dataset(batch):
+        # process audio
+        sample = batch[args.audio_column]
+        # compute log-Mel input features from input audio array
+        batch["input_features"] = \
+        processor.feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"]).input_features[0]
+        # process audio length
+        batch["input_length"] = len(sample["array"]) / sample["sampling_rate"]
+        # process 
+        transcription = batch[args.text_column]
+        if not args.keep_case:
+            transcription = transcription.lower()
+        if not args.keep_punctuation:
+            transcription = normalizer(transcription).strip()
+        batch["labels"] = processor.tokenizer(transcription).input_ids
+        return batch
+    
+    with accelerator.main_process_first():
+        vectorized_datasets = raw_datasets.map(prepare_dataset,
+                                               remove_columns=raw_datasets["train"].column_names)
+    
+    # filter training data with inputs longer than max_input_length
+    def is_audio_in_length_range(length):
+        # return min_input_length <= length <= max_input_length
+        return args.min_duration_in_seconds <= length <= args.max_duration_in_seconds
+    
+    with accelerator.main_process_first():
+        vectorized_datasets = vectorized_datasets.filter(is_audio_in_length_range,
+                                                         input_columns=["input_length"])
+    
+    @dataclass
+    class DataCollatorForSeq2SeqWithPadding:
+        processor: Any
+        
+        def __call__(self, features: List[Union[Dict[str, torch.Tensor], Dict[str, Any]]]) -> Dict[str, torch.Tensor]:
+            # split inputs and labels since they have to be of different lengths and need different padding methods
+            # first treat the audio inputs by simply returning torch tensors
+            input_features = [{"input_features": feature["input_features"]} for feature in features]
+            batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+            
+            # then pad the labels
+            label_features = [{"input_ids": feature["labels"]} for feature in features]
+            labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+            # replace -100 in labels by tokenizer.pad_token_id to ignore padding in loss
+            labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+            
+            # if bos token is appended in previous step, remove it here as it's appended again in the forward pass
+            if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+                labels = labels[:, 1:]
+            
+            batch["labels"] = labels
+            return batch
+    
+    data_collator = DataCollatorForSeq2SeqWithPadding(processor=processor)
+    
+    # now define data loaders
+    train_dataloader = torch.utils.data.DataLoader(vectorized_datasets["train"], shuffle=False, collate_fn=data_collator,
+                                                   batch_size=args.per_device_train_batch_size)
+    eval_dataloader = torch.utils.data.DataLoader(vectorized_datasets["validation"], shuffle=False,
+                                                  collate_fn=data_collator, batch_size=args.per_device_eval_batch_size)
+   
+    # define optimizer
+    optimizer = torch.optim.AdamW(list(student_model.parameters()), lr=args.learning_rate)
+    
+    # scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps,
+                                                   num_training_steps=args.train_steps)
+    
+    # accelerator setup for distributed training, this handles all the device mapping, gradient accumulation, fp16 training etc.
+    # add eval_dataloader to accelerator.prepare for distributed evaluation
+    student_model, teacher_model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        student_model, teacher_model, optimizer, train_dataloader, lr_scheduler)
+    accelerator.print(
+        f"Distributed the model on device: student: {student_model.device}, teacher:{teacher_model.device}, accelerator:{accelerator.device}")
+    
+    global_step = 0  # global step for logging
+    total_loss = 0  # total loss before each eval
+    total_kl_loss = 0  # total kl loss before each eval
+    total_ce_loss = 0  # total ce loss before each eval
+    
+    if args.resume_from_checkpoint is not None:
+        accelerator.print(f"Loading checkpoint: {args.resume_from_checkpoint}")
+        accelerator.load_state(args.resume_from_checkpoint)
+        steps_completed = int(args.resume_from_checkpoint.split("-")[-1])
+        global_step += steps_completed
+        train_dataloader = accelerator.skip_first_batches(train_dataloader, steps_completed)
+    
+    # load metric
+    wer_metric = evaluate.load("wer")
+    cer_metric = evaluate.load("cer")
+    all_punctuations = list(string.punctuation.replace("'", ""))
+    
+    def compute_metrics(preds, labels):
+        # replace padded labels by padding token
+        for idx in range(len(labels)):
+            labels[idx][labels[idx] == -100] = tokenizer.pad_token_id
+        
+        pred_str = tokenizer.batch_decode(preds, skip_special_tokens=True)
+        label_str = tokenizer.batch_decode(labels, skip_special_tokens=True)
+        
+        pred_str = [_.strip() for _ in pred_str]
+        label_str = [_.strip() for _ in label_str]
+        
+        # space punctuation for orthographic WER
+        spaced_pred_str = [pred_str[i].replace(punctuation, "") for punctuation in all_punctuations for i in
+                           range(len(pred_str))]
+        spaced_label_str = [label_str[i].replace(punctuation, "") for punctuation in all_punctuations for i in
+                            range(len(label_str))]
+        
+        # compute WER
+        wer_ortho = 100 * wer_metric.compute(predictions=spaced_pred_str, references=spaced_label_str)
+        cer_ortho = 100 * cer_metric.compute(predictions=spaced_pred_str, references=spaced_label_str)
+        accelerator.print(
+            f"\nspaced_pred_str: {[_ for i, _ in enumerate(spaced_pred_str) if i < 3]}, \nspaced_label_str: {[_ for i, _ in enumerate(spaced_label_str) if i < 3]}")
+        # normalize everything and re-compute the WER
+        norm_pred_str = [normalizer(pred) for pred in pred_str]
+        norm_label_str = [normalizer(label) for label in label_str]
+        # filtering step to only evaluate the samples that correspond to non-zero normlized references:
+        norm_pred_str = [norm_pred_str[i] for i in range(len(norm_label_str)) if len(norm_label_str[i]) > 0]
+        norm_label_str = [norm_label_str[i] for i in range(len(norm_label_str)) if len(norm_label_str[i]) > 0]
+        accelerator.print(
+            f"\nnorm_pred_str: {[_ for i, _ in enumerate(norm_pred_str) if i < 3]}, \nnorm_label_str: {[_ for i, _ in enumerate(norm_label_str) if i < 3]}")
+        
+        wer = 100 * wer_metric.compute(predictions=norm_pred_str, references=norm_label_str)
+        cer = 100 * cer_metric.compute(predictions=norm_pred_str, references=norm_label_str)
+        
+        return {"wer": wer, "wer_ortho": wer_ortho, "cer": cer, "cer_ortho": cer_ortho}, pred_str, label_str
+    
+    # save feature extractor, tokenizer, config and generation config
+    with accelerator.main_process_first():
+        output_dir = args.output_dir
+        feature_extractor.save_pretrained(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        student_config.save_pretrained(output_dir)
+        teacher_config.save_pretrained(output_dir)
+    
+    progress_bar = tqdm(range(global_step, args.train_steps), disable=not accelerator.is_main_process)
+    # define training step
+    while global_step < args.train_steps:
+        student_model.train()
+        for batch in train_dataloader:
+            with accelerator.accumulate(student_model):
+                # forward pass
+                outputs = student_model(**batch)
+                ce_loss = outputs.loss
+                logits = outputs.logits
+                with torch.no_grad():
+                    teacher_logits = teacher_model(**batch).logits
+                # compute kl loss
+                kl_loss = torch.nn.functional.kl_div(torch.nn.functional.log_softmax(logits / args.temperature, dim=-1),
+                                                     torch.nn.functional.softmax(teacher_logits / args.temperature,
+                                                                                 dim=-1),
+                                                     reduction="batchmean") * (args.temperature ** 2)
+                # compute total loss
+                loss = args.alpha_ce * ce_loss + args.alpha_distil * kl_loss
+                
+                total_kl_loss += kl_loss.detach().item()
+                total_ce_loss += ce_loss.detach().item()
+                total_loss += loss.detach().item()
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+                
+                global_step += 1
+                progress_bar.update(1)
+                
+                # log metrics
+                eval_metrics = {}
+                eval_preds = []
+                eval_labels = []
+                
+                if global_step % args.eval_steps == 0:
+                    student_model.eval()
+                    valid_loss = 0
+                    # validation_progress_bar = tqdm(range(0, len(eval_dataloader)), disable=not accelerator.is_main_process)
+                    for batch in eval_dataloader:
+                        with torch.no_grad():
+                            batch.to(accelerator.device)
+                            references = batch.labels
+                            if not args.predict_without_generate:
+                                # accelerator.print("\nPredicting with generate")
+                                # for modules wrapped in DataParallel or DistributedDataParallel, we need to use .module to access the underlying module
+                                if accelerator.num_processes > 1:
+                                    # accelerator.print("Distributed eval")
+                                    predictions = student_model.module.generate(batch.input_features)
+                                else:
+                                    predictions = student_model.generate(batch.input_features)
+                            else:
+                                # accelerator.print("\nPredicting without generate")
+                                outputs = student_model(**batch)
+                                valid_loss += outputs.loss.detach().item()
+                                pred_logits = outputs.logits
+                                predictions = pred_logits.argmax(-1)
+
+                        # accelerator.print("Before gather")
+                        # accelerator.print(f"len of predictions: {len(predictions)}, len of references: {len(references)}")
+                        # accelerator.print(f"All types for gather has to be tensor: \ntype of predictions: {type(predictions)}, type of references: {type(references)}")
+                        predictions, references = accelerator.gather_for_metrics((predictions, references))
+                        # accelerator.print("After gather")
+                        # accelerator.print(f"len of predictions: {len(predictions)}, len of references: {len(references)}")
+
+                        ###########################
+                        # convert any token after after first tokenizer.eos_token_id to eos_token_id
+                        for idx, pred in enumerate(predictions):
+                            first_eos_token_idx = (pred == tokenizer.eos_token_id).nonzero(as_tuple=True)[0]
+                            if len(first_eos_token_idx) > 0:
+                                predictions[idx, first_eos_token_idx[0] + 1:] = tokenizer.eos_token_id
+                        ###########################
+                        
+                        eval_preds.extend(predictions)
+                        eval_labels.extend(references)
+                        # validation_progress_bar.update(1)
+                    
+                    accelerator.print(f"\npredictions: {eval_preds[:3]}, \nreferences: {eval_preds[:3]}")
+                    accelerator.print(f"\nlen(eval_preds): {len(eval_preds)}, \nlen(eval_labels): {len(eval_labels)}")
+                    eval_metrics, eval_preds, eval_labels = compute_metrics(eval_preds, eval_labels)
+                    train_loss = total_loss / (
+                                args.eval_steps * args.per_device_train_batch_size * accelerator.num_processes)
+                    train_kl_loss = total_kl_loss / (
+                                args.eval_steps * args.per_device_train_batch_size * accelerator.num_processes)
+                    train_ce_loss = total_ce_loss / (
+                                args.eval_steps * args.per_device_train_batch_size * accelerator.num_processes)
+                    
+                    accelerator.print(
+                        f"Step: {global_step}, Train Loss: {train_loss}, Train KL Loss: {train_kl_loss}, Train CE Loss: {train_ce_loss}, \
+                                       Eval WER: {eval_metrics['wer']}, Eval WER Ortho: {eval_metrics['wer_ortho']}, Eval CER: {eval_metrics['cer']}, \
+                                      Eval CER Ortho: {eval_metrics['cer_ortho']}")
+                    accelerator.log(
+                        {"cer": eval_metrics["cer"], "cer_ortho": eval_metrics["cer_ortho"], "wer": eval_metrics["wer"],
+                         "wer_ortho": eval_metrics["wer_ortho"],
+                         "train_loss": train_loss,
+                         "train_kl_loss": train_kl_loss,
+                         "train_ce_loss": train_ce_loss,
+                        #  "eval_loss": valid_loss / (len(eval_dataloader))
+                        })
+                    
+                    output_dir = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+                    # accelerator.save(student_model.state_dict(), output_dir)
+                    accelerator.save_state(output_dir)
+                    accelerator.wait_for_everyone()
+                    unwrapped_model = accelerator.unwrap_model(student_model)
+                    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save,
+                                                    is_main_process=accelerator.is_main_process)
+                    total_loss = 0
+                    total_kl_loss = 0
+                    total_ce_loss = 0
+                    student_model.train()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--teacher_model_name_or_path", type=str, default="openai/whisper-large-v2")
+    parser.add_argument("--student_model_name_or_path", type=str, default="distil-whisper/large-v2-8")
+    parser.add_argument("--output_dir", type=str, default="output")
+    parser.add_argument("--per_device_train_batch_size", type=int, default=16)
+    parser.add_argument("--per_device_eval_batch_size", type=int, default=16)
+    parser.add_argument("--learning_rate", type=float, default=2e-5)
+    parser.add_argument("--freeze_encoder", action="store_true")
+    parser.add_argument("--temperature", type=float, default=2.0)
+    parser.add_argument("--alpha_ce", type=float, default=0.5)
+    parser.add_argument("--alpha_distil", type=float, default=0.5)
+    parser.add_argument("--language", type=str, default="en")
+    parser.add_argument("--task", type=str, default="transcribe")
+    parser.add_argument("--train_steps", type=int, default=100000)
+    parser.add_argument("--eval_steps", type=int, default=100)
+    parser.add_argument("--warmup_steps", type=int, default=2000)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=4) # increase by 2x for every 2x decrease in batch size 64
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--data_cache_dir", type=str, default="data/cache")
+    parser.add_argument("--teacher_cache_dir", type=str, default="model/cache")
+    parser.add_argument("--student_cache_dir", type=str, default="model/cache")
+    parser.add_argument("--mixed_precision", type=str, default="fp16")
+    parser.add_argument("--max_train_samples", type=int, default=None)
+    parser.add_argument("--max_val_samples", type=int, default=None)
+    parser.add_argument("--max_test_samples", type=int, default=None)
+    parser.add_argument("--audio_column", type=str, default="audio")
+    parser.add_argument("--text_column", type=str, default="text")
+    parser.add_argument("--max_duration_in_seconds", type=float, default=30)
+    parser.add_argument("--min_duration_in_seconds", type=float, default=1)
+    parser.add_argument("--keep_case", action="store_true")
+    parser.add_argument("--keep_punctuation", action="store_true")
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None)
+    parser.add_argument("--num_workers", type=int, default=16)
+    parser.add_argument("--predict_without_generate", action="store_true")
+    parser.add_argument("--train_dataset_name", type=str, default="librispeech_asr")
+    parser.add_argument("--train_dataset_config_name", type=str, default="all")
+    parser.add_argument("--train_split_name", type=str, default="train.clean.100+train.clean.360+train.other.500")
+    parser.add_argument("--validation_dataset_name", type=str, default="librispeech_asr")
+    parser.add_argument("--validation_dataset_config_name", type=str, default="all")
+    parser.add_argument("--validation_split_name", type=str, default="validation.clean")
+    
+    args = parser.parse_args()
+    
+    # set seed
+    set_seed(args.seed)
+    
+    if args.teacher_model_name_or_path is None or args.student_model_name_or_path is None:
+        raise ValueError("teacher_model_name_or_path and student_model_name_or_path cannot be None")
+    
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, gradient_accumulation_steps=1,
+                              log_with="tensorboard", logging_dir=args.output_dir)
+    
+    # have only one message per logs of transformers or datasets, so logging verbosity INFO only for the main process
+    if accelerator.is_main_process:
+        datasets.utils.logging.set_verbosity_info()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+    
+    # establish trackers for logging
+    track_config = {"lr": args.learning_rate,
+                    "train_batch_size": args.per_device_train_batch_size,
+                    "eval_batch_size": args.per_device_eval_batch_size,
+                    "seed": args.seed,
+                    "train_steps": args.train_steps}
+    accelerator.init_trackers('runs', track_config)
+    train(args, accelerator)
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main()