{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/tesla/miniconda3/envs/DUY/lib/python3.9/site-packages/whisper/__init__.py\n"
     ]
    }
   ],
   "source": [
    "import whisper\n",
    "print(whisper.__file__)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "2\n",
      "0\n",
      "<torch.cuda.device object at 0x7f69e1e31eb0>\n",
      "Tesla T4\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "print(torch.cuda.is_available())\n",
    "\n",
    "\n",
    "print(torch.cuda.device_count())\n",
    "\n",
    "\n",
    "print(torch.cuda.current_device())\n",
    "print(torch.cuda.device(0))\n",
    "\n",
    "print(torch.cuda.get_device_name(0))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f290d4efc37a4112a662c062e621e482",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name_or_path = \"openai/whisper-medium\"\n",
    "task = \"transcribe\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = \"Vin100h_MITI_private\"\n",
    "language = \"Vietnamese\"\n",
    "language_abbr = \"vi\" # Short hand code for the language we want to fine-tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['input_features', 'input_length', 'attention_mask', 'labels'],\n",
      "        num_rows: 1679\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['input_features', 'input_length', 'attention_mask', 'labels'],\n",
      "        num_rows: 420\n",
      "    })\n",
      "})\n",
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['input_features', 'input_length', 'attention_mask', 'labels'],\n",
      "        num_rows: 6735\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['input_features', 'input_length', 'attention_mask', 'labels'],\n",
      "        num_rows: 1688\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "    # Load Dataset\n",
    "from datasets import load_dataset, DatasetDict, load_from_disk\n",
    "processed_dataset = DatasetDict()\n",
    "processed_dataset = load_from_disk(\"./MITI_clean\")\n",
    "processed_dataset2 = load_from_disk(\"./vin_10h/\")\n",
    "\n",
    "print(processed_dataset)\n",
    "print(processed_dataset2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "# Assuming you have already loaded your dataset\n",
    "# processed_dataset2 = ...\n",
    "\n",
    "# Randomly select 5000 indices from the train dataset\n",
    "import random\n",
    "num_samples_train = 5000\n",
    "num_samples_test = 600\n",
    "random_indices_train = random.sample(range(len(processed_dataset2['train'])), num_samples_train)\n",
    "random_indices_test = random.sample(range(len(processed_dataset2['test'])), num_samples_test)\n",
    "\n",
    "# Initialize lists for train dataset\n",
    "input_features_train = []\n",
    "input_length_train = []\n",
    "attention_mask_train = []\n",
    "labels_train = []\n",
    "\n",
    "# Initialize lists for test dataset\n",
    "input_features_test = []\n",
    "input_length_test = []\n",
    "attention_mask_test = []\n",
    "labels_test = []\n",
    "\n",
    "# Populate lists for train dataset\n",
    "for i in random_indices_train:\n",
    "    input_features_train.append(processed_dataset2['train'][i]['input_features'])\n",
    "    input_length_train.append(processed_dataset2['train'][i]['input_length'])\n",
    "    attention_mask_train.append(processed_dataset2['train'][i]['attention_mask'])\n",
    "    labels_train.append(processed_dataset2['train'][i]['labels'])\n",
    "\n",
    "# Populate lists for test dataset\n",
    "for i in random_indices_test:\n",
    "    input_features_test.append(processed_dataset2['test'][i]['input_features'])\n",
    "    input_length_test.append(processed_dataset2['test'][i]['input_length'])\n",
    "    attention_mask_test.append(processed_dataset2['test'][i]['attention_mask'])\n",
    "    labels_test.append(processed_dataset2['test'][i]['labels'])\n",
    "\n",
    "# Create a new dataset with the randomly selected rows\n",
    "random_subset = Dataset.from_dict({\n",
    "    'train': {\n",
    "        'input_features': input_features_train,\n",
    "        'input_length': input_length_train,\n",
    "        'attention_mask': attention_mask_train,\n",
    "        'labels': labels_train,\n",
    "    },\n",
    "    'test': {\n",
    "        'input_features': input_features_test,\n",
    "        'input_length': input_length_test,\n",
    "        'attention_mask': attention_mask_test,\n",
    "        'labels': labels_test,\n",
    "    }\n",
    "})\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "concat = DatasetDict()\n",
    "concat[\"train\"] = datasets.concatenate_datasets([processed_dataset[\"train\"], processed_dataset2[\"train\"]])\n",
    "concat['test']= datasets.concatenate_datasets([processed_dataset[\"test\"], processed_dataset2[\"test\"]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_features', 'input_length', 'attention_mask', 'labels'],\n",
       "        num_rows: 8414\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_features', 'input_length', 'attention_mask', 'labels'],\n",
       "        num_rows: 2108\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperFeatureExtractor\n",
    "\n",
    "feature_extractor = WhisperFeatureExtractor.from_pretrained(model_name_or_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from transformers import WhisperTokenizer\n",
    "\n",
    "tokenizer = WhisperTokenizer.from_pretrained(model_name_or_path, language=language, task=task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('./Viet_ASR/tokenizer_config.json',\n",
       " './Viet_ASR/special_tokens_map.json',\n",
       " './Viet_ASR/vocab.json',\n",
       " './Viet_ASR/merges.txt',\n",
       " './Viet_ASR/normalizer.json',\n",
       " './Viet_ASR/added_tokens.json')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.save_pretrained('./Viet_ASR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import Any, Dict, List, Union\n",
    "from transformers import WhisperProcessor\n",
    "\n",
    "processor = WhisperProcessor.from_pretrained(model_name_or_path, language=language, task=task)\n",
    "@dataclass\n",
    "class DataCollatorSpeechSeq2SeqWithPadding:\n",
    "    processor: Any\n",
    "\n",
    "    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n",
    "        # split inputs and labels since they have to be of different lengths and need different padding methods\n",
    "        # first treat the audio inputs by simply returning torch tensors\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",
    "\n",
    "        # get the tokenized label sequences\n",
    "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
    "\n",
    "\n",
    "        # pad the labels to max length\n",
    "        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n",
    "\n",
    "         # replace padding with -100 to ignore loss correctly\n",
    "        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n",
    "\n",
    "        # if bos token is appended in previous tokenization step,\n",
    "        # cut bos token here as it's append later anyways\n",
    "        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():\n",
    "            labels = labels[:, 1:]\n",
    "\n",
    "        batch[\"labels\"] = labels\n",
    "\n",
    "        return batch\n",
    "data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "\n",
    "metric = evaluate.load(\"wer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperForConditionalGeneration\n",
    "\n",
    "model = WhisperForConditionalGeneration.from_pretrained('openai/whisper-medium', load_in_8bit=True, device_map=\"auto\" )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.config.forced_decoder_ids = None\n",
    "model.config.suppress_tokens = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch.utils.hooks.RemovableHandle at 0x7f1a9445da60>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from peft import prepare_model_for_kbit_training\n",
    "\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "def make_inputs_require_grad(module, input, output):\n",
    "    output.requires_grad_(True)\n",
    "\n",
    "model.model.encoder.conv1.register_forward_hook(make_inputs_require_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 9,437,184 || all params: 773,295,104 || trainable%: 1.2203858463844612\n"
     ]
    }
   ],
   "source": [
    "from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model\n",
    "#target_modules = [\"k_proj\", \"q_proj\", \"v_proj\", \"out_proj\", \"fc1\", \"fc2\"] #will it better ?\n",
    "target_modules=[\"q_proj\", \"v_proj\"]\n",
    "config = LoraConfig(r=32, lora_alpha=64, target_modules=target_modules, lora_dropout=0.05, bias=\"none\")\n",
    "\n",
    "model = get_peft_model(model, config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Seq2SeqTrainingArguments\n",
    "\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=\"./Vietnamese_ASR\", \n",
    "    per_device_train_batch_size=10,\n",
    "    #auto_find_batch_size  = True,\n",
    "    gradient_accumulation_steps=2,  # increase by 2x for every 2x decrease in batch size\n",
    "    learning_rate=5e-5,\n",
    "    warmup_steps=50,\n",
    "    num_train_epochs=3,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    gradient_checkpointing=True,\n",
    "    optim=\"adamw_torch\",\n",
    "    fp16=True,\n",
    "    per_device_eval_batch_size=8,\n",
    "    generation_max_length=225,\n",
    "    logging_steps=100,\n",
    "    report_to=[\"tensorboard\"],\n",
    "    predict_with_generate=True,\n",
    "    # load_best_model_at_end=True,\n",
    "    greater_is_better=False,\n",
    "    save_strategy = \"epoch\",\n",
    "    # required as the PeftModel forward doesn't have the signature of the wrapped model's forward\n",
    "    remove_unused_columns=False,\n",
    "    label_names=[\"labels\"],  # same reason as above\n",
    "    push_to_hub=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/tesla/New Volume/DEMO/DUY/Vietnamese_ASR/./Vietnamese_ASR is already a clone of https://huggingface.co/DuyTa/Vietnamese_ASR. Make sure you pull the latest changes with `repo.git_pull()`.\n"
     ]
    }
   ],
   "source": [
    "from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl\n",
    "from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR\n",
    "\n",
    "\n",
    "class SavePeftModelCallback(TrainerCallback):\n",
    "    def on_save(\n",
    "        self,\n",
    "        args: TrainingArguments,\n",
    "        state: TrainerState,\n",
    "        control: TrainerControl,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        checkpoint_folder = os.path.join(args.output_dir, f\"{PREFIX_CHECKPOINT_DIR}-{state.global_step}\")\n",
    "\n",
    "        peft_model_path = os.path.join(checkpoint_folder, \"adapter_model\")\n",
    "        kwargs[\"model\"].save_pretrained(peft_model_path)\n",
    "\n",
    "        pytorch_model_path = os.path.join(checkpoint_folder, \"pytorch_model.bin\")\n",
    "        if os.path.exists(pytorch_model_path):\n",
    "            os.remove(pytorch_model_path)\n",
    "        return control\n",
    "\n",
    "\n",
    "trainer = Seq2SeqTrainer(\n",
    "    args=training_args,\n",
    "    model=model,\n",
    "    train_dataset=concat[\"train\"],\n",
    "    eval_dataset=concat[\"test\"],\n",
    "    data_collator=data_collator,\n",
    "    # compute_metrics=compute_metrics,\n",
    "    tokenizer=processor.feature_extractor,\n",
    "    callbacks=[SavePeftModelCallback],\n",
    ")\n",
    "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0aa0180f6e64eaa8951a4c940aa518f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1263 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tesla/miniconda3/envs/DUY/lib/python3.9/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 1.9814, 'learning_rate': 4.814509480626546e-05, 'epoch': 0.24}\n",
      "{'loss': 0.6861, 'learning_rate': 4.402308326463314e-05, 'epoch': 0.48}\n",
      "{'loss': 0.3736, 'learning_rate': 3.9901071723000826e-05, 'epoch': 0.71}\n",
      "{'loss': 0.332, 'learning_rate': 3.577906018136851e-05, 'epoch': 0.95}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e9a8f06d39e448a9523d9a29699cadc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/264 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.3133259117603302, 'eval_runtime': 887.0949, 'eval_samples_per_second': 2.376, 'eval_steps_per_second': 0.298, 'epoch': 1.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tesla/miniconda3/envs/DUY/lib/python3.9/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.3005, 'learning_rate': 3.165704863973619e-05, 'epoch': 1.19}\n",
      "{'loss': 0.307, 'learning_rate': 2.753503709810388e-05, 'epoch': 1.43}\n",
      "{'loss': 0.2838, 'learning_rate': 2.341302555647156e-05, 'epoch': 1.66}\n",
      "{'loss': 0.2746, 'learning_rate': 1.9291014014839242e-05, 'epoch': 1.9}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e65ecdbc96246b8b4721505b4252a8a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/264 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.28433552384376526, 'eval_runtime': 880.1965, 'eval_samples_per_second': 2.395, 'eval_steps_per_second': 0.3, 'epoch': 2.0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tesla/miniconda3/envs/DUY/lib/python3.9/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.2857, 'learning_rate': 1.5169002473206925e-05, 'epoch': 2.14}\n",
      "{'loss': 0.2643, 'learning_rate': 1.104699093157461e-05, 'epoch': 2.38}\n",
      "{'loss': 0.2604, 'learning_rate': 6.924979389942292e-06, 'epoch': 2.61}\n",
      "{'loss': 0.2505, 'learning_rate': 2.8029678483099755e-06, 'epoch': 2.85}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a4f479ef36f4f00b3c591503b411e5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/264 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.27759623527526855, 'eval_runtime': 879.7333, 'eval_samples_per_second': 2.396, 'eval_steps_per_second': 0.3, 'epoch': 3.0}\n",
      "{'train_runtime': 35575.7347, 'train_samples_per_second': 0.71, 'train_steps_per_second': 0.036, 'train_loss': 0.4555940831925127, 'epoch': 3.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=1263, training_loss=0.4555940831925127, metrics={'train_runtime': 35575.7347, 'train_samples_per_second': 0.71, 'train_steps_per_second': 0.036, 'train_loss': 0.4555940831925127, 'epoch': 3.0})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DuyTa/Vietnamese_ASR\n"
     ]
    }
   ],
   "source": [
    "peft_model_id = \"DuyTa/Vietnamese_ASR\"\n",
    "model.push_to_hub(peft_model_id)\n",
    "print(peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import WhisperForConditionalGeneration\n",
    "peft_model_id = \"./Vietnamese_ASR\"\n",
    "peft_config = PeftConfig.from_pretrained(peft_model_id)\n",
    "model = WhisperForConditionalGeneration.from_pretrained(\n",
    "    peft_config.base_model_name_or_path, load_in_8bit=True, device_map=\"auto\"\n",
    ")\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/88 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tesla/miniconda3/envs/DUY/lib/python3.9/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n",
      "100%|██████████| 88/88 [59:07<00:00, 40.31s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wer=15.57082617523036\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import gc\n",
    "\n",
    "eval_dataloader = DataLoader(concat[\"test\"], batch_size=24, collate_fn=data_collator)\n",
    "\n",
    "model.eval()\n",
    "for step, batch in enumerate(tqdm(eval_dataloader)):\n",
    "    with torch.cuda.amp.autocast():\n",
    "        with torch.no_grad():\n",
    "            generated_tokens = (\n",
    "                model.generate(\n",
    "                    input_features=batch[\"input_features\"].to(\"cuda\"),\n",
    "                    decoder_input_ids=batch[\"labels\"][:, :4].to(\"cuda\"),\n",
    "                    max_new_tokens=255,\n",
    "                )\n",
    "                .cpu()\n",
    "                .numpy()\n",
    "            )\n",
    "            labels = batch[\"labels\"].cpu().numpy()\n",
    "            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "            decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)\n",
    "            decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "            metric.add_batch(\n",
    "                predictions=decoded_preds,\n",
    "                references=decoded_labels,\n",
    "            )\n",
    "    del generated_tokens, labels, batch\n",
    "    gc.collect()\n",
    "wer = 100 * metric.compute()\n",
    "print(f\"{wer=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text Norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# using Vietnamese text normalization after take whisper out token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 88/88 [58:18<00:00, 39.75s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "normalized_wer=14.601364195460137\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import gc\n",
    "from transformers.models.whisper.english_normalizer import BasicTextNormalizer\n",
    "normalizer = BasicTextNormalizer()\n",
    "forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)\n",
    "\n",
    "model.eval()\n",
    "for step, batch in enumerate(tqdm(eval_dataloader)):\n",
    "    with torch.cuda.amp.autocast():\n",
    "        with torch.no_grad():\n",
    "            generated_tokens= model.generate(input_features=batch[\"input_features\"].to(\"cuda\"),\n",
    "                                             forced_decoder_ids=forced_decoder_ids,\n",
    "                                             max_new_tokens=255).cpu().numpy()\n",
    "            labels = batch[\"labels\"].cpu().numpy()\n",
    "            labels = np.where(labels != -100, labels, processor.tokenizer.pad_token_id)\n",
    "            decoded_preds = processor.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)\n",
    "            decoded_labels = processor.tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "            metric.add_batch(\n",
    "                predictions=[normalizer(pred).strip() for pred in decoded_preds],\n",
    "                references=[normalizer(label).strip() for label in decoded_labels],\n",
    "            )\n",
    "    # if step==0:\n",
    "    #     break\n",
    "    del generated_tokens, labels, batch\n",
    "    gc.collect()\n",
    "normalized_wer = 100 * metric.compute()\n",
    "print(f\"{normalized_wer=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pred='toàn bộ phi hành đoàn đã bị giết chết khiến con tàu quân sự đậm thẳng vào ishimura', label='toàn bộ phi hành đoàn đã bị giết chết khiến con tàu quân sự đâm thẳng vào ishimura'\n",
      "pred='đủ kinh nghiệm để mình quản lý nhân viên và mình làm sao để mình đưa ra một cái dịch vụ tốt nhất', label='đủ kinh nghiệm để mình quản lý nhân viên và mình làm sao để mình đưa ra một cái dịch vụ tốt nhất'\n",
      "pred='nói một trong một cái chương trình trong tương lai về ngành thu y thì ở mỹ tất cả các đại học nào lớn đều có ngành thu y hết', label='<unk> nói một trong một cái chương trình trong tương lai về ngành thú y thì ở mỹ tất cả các đại học nào lớn đều có ngành thú y hết'\n",
      "pred='phấn đấu đến năm hai ngàn không trăm hai mười có từ tám trăm đến một ngàn kinh nghiệp tham gia sàn sâu dịch thu mại điện tử của bộ công thương và quốc tế năm mươi phần trăm số này vô bắn trên sàn', label='phấn đấu đến năm hai ngàn không trăm hai mươi có từ tám trăm đến một ngàn doanh nghiệp tham gia sàn giao dịch thương mại điện tử của bộ công thương và quốc tế năm mươi phần trăm số này luôn bán trên sàn'\n",
      "pred='còn trách nhiệm kiểm tra thanh tra là của ủy ban nhân dân các cấp', label='còn trách nhiệm kiểm tra thanh tra là của ủy ban nhân dân các cấp'\n",
      "pred='vậy mà cậu im lặng khóa trái tìm mình chắc địa dành cho cái gì đó vĩ đại hơn chăng', label='vậy mà cậu im lặng khóa trái tim mình chắc để giành cho cái gì đó vĩ đại hơn chăng'\n",
      "pred='khi nộp phiếu trả lời trắc nghiệm thí sinh phải ghi tên và danh sách thí sinh nộp bài', label='khi nộp phiếu trả lời trắc nghiệm thí sinh phải ký tên vào danh sách thí sinh nộp bài'\n",
      "pred='khi nghĩ rằng mình đã khỏi ai ngờ ung thư lại tái phát và tôi đã lắng nghe câu chuyện của tất cả mọi người', label='khi nghĩ rằng mình đã khỏi ai ngờ ung thư lại tái phát và tôi đã lắng nghe câu chuyện của tất cả mọi người'\n",
      "pred='người cùng ấp là trương thật từng muốn kết giao với giám vì ông từ chối', label='người cùng ấp là trương thực từng muốn kết giao với giám bị ông từ chối'\n",
      "pred='bài thơ với những dòng thơ rất xúc động như sau', label='bài thơ với những dòng thơ rất xúc động như sau'\n",
      "pred='công bố chỉ số niềm tin kinh doanh của doanh nghiệp', label='công bố chỉ số niềm tin kinh doanh của doanh nghiệp'\n",
      "pred='khi quanh hồ tổng tới thăng lông đúng lúc tô trung tự đang đánh nhau to với đồ quảng', label='khi quân hộ tống tới thăng long đúng lúc tô trung từ đang đánh nhau to với đỗ quảng'\n",
      "pred='chứ không lẽ bây giờ kêu men trai', label='chứ hổng lẽ bây giờ kêu mê trai'\n",
      "pred='trong thời gian đó anh ấy hãy tâm sự với tôi', label='trong thời gian đó anh ấy hay tâm sự với tôi'\n",
      "pred='mi mo sa lại cho màu sắc lá đẹp không cần dùng đến màu nhuộng hoàng sực dỡ từ duy nhẹ bé đó đã giúp vườn mi mo sa của bà nhanh chóng đem lại lợi nhuộn', label='mi mo sa lại cho màu sắc lá đẹp không cần dùng đến màu nhuộm hoa rực rỡ tư duy nhạy bén đó đã giúp vườn mi mo sa của bà nhanh chóng đem lại lợi nhuận'\n",
      "pred='chơi tìm kiếm tài năng thiên đỉnh god thai lần thế các táo đâu hết cả rồi', label='chơi tìm kiếm tài năng thiên đình gót thai lừn thế các táo đâu hết cả rồi'\n",
      "pred='dù đức và pháp bất đồng sâu sắc nhưng chính kiến của họ thì đều sai', label='dù đức và pháp bất đồng sâu sắc nhưng chính kiến của họ thì đều sai'\n",
      "pred='đại ca bảo không hình anh ra mà ngồi anh đánh đi', label='đại ca bảo buông anh ra mà thôi anh'\n",
      "pred='khi mà mang thai bác thị cũng cảnh báo rồi', label='khi mà mang thai thì bác sĩ cũng cảnh báo rồi'\n",
      "pred='là tăng giảm thất thường và đột xuất kéo dài', label='mà tăng giảm thất thường và đột xuất kéo dài'\n"
     ]
    }
   ],
   "source": [
    "for pred,label in zip(decoded_preds,decoded_labels):\n",
    "  print(f\"{pred=}, {label=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import (\n",
    "    AutomaticSpeechRecognitionPipeline,\n",
    "    WhisperForConditionalGeneration,\n",
    "    WhisperTokenizer,\n",
    "    WhisperProcessor,\n",
    ")\n",
    "from peft import PeftModel, PeftConfig\n",
    "\n",
    "\n",
    "peft_model_id = \"./Vietnamese_ASR\"\n",
    "language = \"Vietnamese\"\n",
    "task = \"transcribe\"\n",
    "\n",
    "peft_config = PeftConfig.from_pretrained(peft_model_id)\n",
    "model = WhisperForConditionalGeneration.from_pretrained(\n",
    "    peft_config.base_model_name_or_path\n",
    ")\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)\n",
    "merged_model = model.merge_and_unload()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_model.save_pretrained(\"./Vietnamese_ASR/merged\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperTokenizer\n",
    "\n",
    "tokenizer = WhisperTokenizer.from_pretrained('openai/whisper-medium', language=language, task=task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('./Vietnamese_ASR/merged/tokenizer_config.json',\n",
       " './Vietnamese_ASR/merged/special_tokens_map.json',\n",
       " './Vietnamese_ASR/merged/vocab.json',\n",
       " './Vietnamese_ASR/merged/merges.txt',\n",
       " './Vietnamese_ASR/merged/normalizer.json',\n",
       " './Vietnamese_ASR/merged/added_tokens.json')"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.save_pretrained('./Vietnamese_ASR/merged')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: /home/tesla/miniconda3/lib/libtinfo.so.6: no version information available (required by /bin/bash)\n"
     ]
    }
   ],
   "source": [
    "!ct2-transformers-converter --model ./Vietnamese_ASR/merged --output_dir ./Vietnamese_ASR/ct2ranslate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: /home/tesla/miniconda3/lib/libtinfo.so.6: no version information available (required by /bin/bash)\n"
     ]
    }
   ],
   "source": [
    "!ct2-transformers-converter --model ./Vietnamese_ASR/merged --output_dir ./Vietnamese_ASR/ct2ranslate/quantized --quantization float16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/tesla/New Volume/DEMO/DUY/Vietnamese_ASR/Vietnamese_ASR/src/Vietnamese_ASR is already a clone of https://huggingface.co/DuyTa/Vietnamese_ASR. Make sure you pull the latest changes with `repo.git_pull()`.\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import Repository\n",
    "repo = Repository(local_dir=\"\", clone_from='DuyTa/Vietnamese_ASR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "061e5ea903e04d2e95bc3ff8a8de434b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Clean file runs/Aug17_22-42-43_tesla-T4/events.out.tfevents.1692289257.tesla-T4.201346.0:  14%|#4        | 1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "repo.git_pull(rebase=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "repo.git_add(\".\")\n",
    "repo.git_commit(commit_message=\"3 epochs finetuning and quantized model )\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Several commits (3) will be pushed upstream.\n",
      "The progress bars may be unreliable.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "To https://huggingface.co/DuyTa/Vietnamese_ASR\n",
      "   63bacc4..82e8e84  main -> main\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'https://huggingface.co/DuyTa/Vietnamese_ASR/commit/82e8e84fe4f1ffee17eff82c39a163f4b81335d5'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repo.git_push()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_model.push_to_hub(\"DuyTa/MITI_Whisper\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import gradio as gr\n",
    "from transformers import (\n",
    "    AutomaticSpeechRecognitionPipeline,\n",
    "    WhisperForConditionalGeneration,\n",
    "    WhisperTokenizer,\n",
    "    WhisperProcessor,\n",
    ")\n",
    "from peft import PeftModel, PeftConfig\n",
    "\n",
    "\n",
    "peft_model_id = \"DuyTa/MITI_Whisper\"\n",
    "language = \"Vietnamese\"\n",
    "task = \"transcribe\"\n",
    "peft_config = PeftConfig.from_pretrained(peft_model_id)\n",
    "model = WhisperForConditionalGeneration.from_pretrained(\n",
    "    peft_config.base_model_name_or_path, load_in_8bit=True, device_map=\"auto\"\n",
    ")\n",
    "\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)\n",
    "tokenizer = WhisperTokenizer.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)\n",
    "processor = WhisperProcessor.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)\n",
    "feature_extractor = processor.feature_extractor\n",
    "forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)\n",
    "pipe = AutomaticSpeechRecognitionPipeline(model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)\n",
    "\n",
    "\n",
    "def transcribe(audio):\n",
    "    with torch.cuda.amp.autocast():\n",
    "        text = pipe(audio, generate_kwargs={\"forced_decoder_ids\": forced_decoder_ids}, max_new_tokens=255)[\"text\"]\n",
    "    return text\n",
    "\n",
    "\n",
    "iface = gr.Interface(\n",
    "    fn=transcribe,\n",
    "    inputs=gr.Audio(source=\"upload\", type=\"filepath\"),\n",
    "    outputs=\"text\",\n",
    "    title=\"PEFT LoRA\",\n",
    "    description=\"Realtime demo for Vietnamese speech recognition using `PEFT-LoRA+INT8` fine-tuned Whisper Medium .\",\n",
    ")\n",
    "\n",
    "iface.launch(share=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DUY",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}