{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":73047,"databundleVersionId":8149390},{"sourceType":"datasetVersion","sourceId":6707460,"datasetId":3865741,"databundleVersionId":6791840},{"sourceType":"datasetVersion","sourceId":8205178,"datasetId":4825627,"databundleVersionId":8329522},{"sourceType":"datasetVersion","sourceId":8205748,"datasetId":4812155,"databundleVersionId":8330127}],"dockerImageVersionId":30683,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **0. Import Libraries**","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"!pip install -q bitsandbytes #accelerate bitsandbytes==0.37\n!pip install -q git+https://github.com/huggingface/peft.git@main\n!pip install -q git+https://github.com/huggingface/accelerate.git@main\n# !pip install -q noisereduce","metadata":{"execution":{"iopub.status.busy":"2024-04-23T12:53:14.418672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport librosa\nimport librosa.display\nimport numpy as np\nimport IPython.display as ipd\nimport matplotlib.pyplot as plt\nimport random\nimport torch\nimport torchaudio\n\nfrom collections import Counter\nfrom sklearn.model_selection import train_test_split\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, List, Union\nfrom datasets import DatasetDict\nfrom datasets import Dataset as DS\nfrom transformers import (\n    WhisperFeatureExtractor,\n    WhisperTokenizer,\n    WhisperProcessor,\n    WhisperForConditionalGeneration,\n    BitsAndBytesConfig,\n    GenerationConfig,\n    Seq2SeqTrainingArguments,\n    Seq2SeqTrainer,\n    TrainerCallback,\n    TrainingArguments,\n    TrainerState,\n    TrainerControl,\n    EarlyStoppingCallback,\n    pipeline\n)\nfrom peft import (\n    prepare_model_for_kbit_training,\n    LoraConfig, \n    PeftModel, \n    LoraModel, \n    LoraConfig, \n    get_peft_model,\n    PeftConfig\n)\nfrom torchmetrics.text import WordErrorRate, CharErrorRate","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **1. Load Dataset**","metadata":{}},{"cell_type":"code","source":"from datasets import load_from_disk \nds_splits = load_from_disk (\"/kaggle/input/bengali-asr-dialect-preprocessed-dataset/full-preprocessed-dataset-v1\")\nprint(ds_splits)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **2. Preparing Processor & pre-trained model**","metadata":{}},{"cell_type":"code","source":"def make_inputs_require_grad(module, input, output):\n    output.requires_grad_(True)\n\n# This callback helps to save only the adapter weights and remove the base model weights.\nclass SavePeftModelCallback(TrainerCallback):\n    \n    def on_save(self,args: TrainingArguments,state: TrainerState,control: TrainerControl,**kwargs):\n        peft_model_path = model_id\n        kwargs[\"model\"].save_pretrained(peft_model_path)\n\n        pytorch_model_path = f\"{peft_model_path}/pytorch_model.bin\"\n        if os.path.exists(pytorch_model_path):os.remove(pytorch_model_path)\n        return control","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TASK = \"transcribe\"\nLANGUAGE='bn'\n# MODEL_NAME = \"openai/whisper-small\"\n# MODEL_NAME = \"openai/whisper-medium\"\n# MODEL_NAME = \"openai/whisper-large-v2\"\n# MODEL_NAME=\"/kaggle/input/bengali-ai-asr-submission/bengali-whisper-medium\"\nMODEL_NAME=\"/kaggle/input/whisper-model/bengali-whisper-medium-pretrained-augmented-dataset-lora-added\"\n\nquant_config = BitsAndBytesConfig(\n    load_in_4bit=True,\n    bnb_4bit_quant_type=\"nf4\",\n    bnb_4bit_use_double_quant=True,\n    bnb_4bit_compute_dtype=torch.bfloat16,\n)\nlora_config = LoraConfig(r=32, lora_alpha=64, target_modules=[\"q_proj\", \"v_proj\"], lora_dropout=0.05, bias=\"none\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"feature_extractor = WhisperFeatureExtractor.from_pretrained(MODEL_NAME)\ntokenizer = WhisperTokenizer.from_pretrained(MODEL_NAME, language=LANGUAGE, task=TASK)\nprocessor = WhisperProcessor.from_pretrained(MODEL_NAME, language=LANGUAGE, task=TASK)\n\nmodel = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME,quantization_config=quant_config,device_map=\"auto\")\nmodel = prepare_model_for_kbit_training(model)\nmodel.model.encoder.conv1.register_forward_hook(make_inputs_require_grad)\n# model.model.encoder.conv2.register_forward_hook(make_inputs_require_grad)\nmodel = get_peft_model(model, lora_config)\n# model.print_trainable_parameters()\n\nmodel.generation_config=GenerationConfig.from_pretrained(\"openai/whisper-medium\")\nmodel.generation_config.language = LANGUAGE\nmodel.generation_config.task = TASK\nmodel.generation_config.forced_decoder_ids = None\nmodel.config.suppress_tokens = [] # added later","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.generation_config","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **3. Training**","metadata":{}},{"cell_type":"code","source":"@dataclass\nclass DataCollatorSpeechSeq2SeqWithPadding:\n    processor: Any\n\n    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n        input_features = [{\"input_features\": feature[\"input_features\"]} for feature in features]\n        batch = self.processor.feature_extractor.pad(input_features, return_tensors=\"pt\")\n        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():\n            labels = labels[:, 1:]\n        batch[\"labels\"] = labels\n        torch.cuda.empty_cache()\n        return batch\n    \ndata_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)\n\ncer = CharErrorRate()\nwer = WordErrorRate()\n\ndef compute_metrics(pred):\n    pred_ids = pred.predictions\n    label_ids = pred.label_ids\n\n    label_ids[label_ids == -100] = tokenizer.pad_token_id\n\n    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)\n    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)\n\n    wer_res = wer(pred_str, label_str)\n    cer_res = cer(pred_str, label_str)\n    \n    print(\"WER:\",wer_res,\"| CER:\", cer_res) # to show up during running logs\n    print(\"Pred:\",pred_str[0])\n    print(\"Label:\",label_str[0])\n    \n    return {\"wer\": wer_res, \"cer\": cer_res}\n\nsteps=1690\nbatch=8\nmodel_id=\"whisper-medium-augmented+lora+full-data+2-epoch\"\ntraining_args = Seq2SeqTrainingArguments(\n    output_dir=model_id,\n    per_device_train_batch_size=batch,\n    per_device_eval_batch_size=batch,\n    gradient_accumulation_steps=1,\n    gradient_checkpointing=True,\n    learning_rate=2e-4,\n    fp16=True,\n    weight_decay=1e-2,\n    warmup_steps=50,\n    num_train_epochs=2,\n    evaluation_strategy=\"steps\", # or \"epochs\"\n    predict_with_generate=True,\n#     generation_max_length=448,\n    save_steps=steps,\n    eval_steps=steps,\n    logging_steps=steps,\n    save_total_limit=1,\n    load_best_model_at_end=True,\n    metric_for_best_model=\"wer\",\n    greater_is_better=False,\n    push_to_hub=False,\n    report_to=\"none\",\n    optim = \"paged_adamw_8bit\",\n#     remove_unused_columns=False,\n)\n\ntrainer = Seq2SeqTrainer(\n    args=training_args,\n    model=model,\n    train_dataset=ds_splits[\"train\"],\n    eval_dataset=ds_splits[\"eval\"],\n    data_collator=data_collator,\n    tokenizer=processor.feature_extractor,\n    compute_metrics=compute_metrics,\n    callbacks=[SavePeftModelCallback],\n#     callbacks=[EarlyStoppingCallback(2, 1.0)]\n#     callbacks=[SavePeftModelCallback, EarlyStoppingCallback(2, 1.0)]\n)\ntrainer.train()\ntrainer.save_model(training_args.output_dir)\nprocessor.save_pretrained(training_args.output_dir)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}