4_0_prepare_finetuned_whisper.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5753cc2a-3e5f-453c-b2bf-f142ca8a0de7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token will not been saved to git credential helper. Pass `add_to_git_credential=True` if you want to set the git credential as well.\n",
      "Token is valid (permission: write).\n",
      "Your token has been saved to /home/zhang.jinda1/.cache/huggingface/token\n",
      "Login successful\n"
     ]
    }
   ],
   "source": [
    "# This notebook uses https://github.com/openai/whisper with edits to the whisper_openAI/decoding.py to generate multiple hypothesis\n",
    "import sys\n",
    "import os\n",
    "import argparse\n",
    "import re\n",
    "import json\n",
    "import torch\n",
    "import logging\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from datasets import load_dataset, DatasetDict, Audio\n",
    "from dataclasses import dataclass\n",
    "from typing import Dict, List, Union\n",
    "from evaluate import load\n",
    "from tqdm import tqdm\n",
    "from datetime import datetime\n",
    "\n",
    "!huggingface-cli login --token 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0338206e-5029-44f2-af94-01be319a4d03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Speaker ID: F01\n",
      "Learning rate: 0.0001\n",
      "Training batch size: 4\n",
      "Evaluation batch size: 4\n",
      "Random seed: 42\n",
      "Gradient accumulation steps: 2\n",
      "Optimizer type: adamw_torch\n",
      "Learning rate scheduler type: linear\n",
      "Number of epochs: 20\n",
      "Keep all data: False\n",
      "Debug mode: False\n",
      "Repository suffix: \n"
     ]
    }
   ],
   "source": [
    "speaker_id = \"F01\"  # Example value; replace with the actual Speaker ID as needed\n",
    "learning_rate = 0.0001\n",
    "train_batch_size = 4\n",
    "eval_batch_size = 4\n",
    "seed = 42\n",
    "gradient_accumulation_steps = 2\n",
    "optimizer = \"adamw_torch\"\n",
    "lr_scheduler_type = \"linear\"\n",
    "num_epochs = 20\n",
    "keep_all_data = False\n",
    "debug = False\n",
    "repo_suffix = \"\"\n",
    "\n",
    "print(f\"Speaker ID: {speaker_id}\")\n",
    "print(f\"Learning rate: {learning_rate}\")\n",
    "print(f\"Training batch size: {train_batch_size}\")\n",
    "print(f\"Evaluation batch size: {eval_batch_size}\")\n",
    "print(f\"Random seed: {seed}\")\n",
    "print(f\"Gradient accumulation steps: {gradient_accumulation_steps}\")\n",
    "print(f\"Optimizer type: {optimizer}\")\n",
    "print(f\"Learning rate scheduler type: {lr_scheduler_type}\")\n",
    "print(f\"Number of epochs: {num_epochs}\")\n",
    "print(f\"Keep all data: {keep_all_data}\")\n",
    "print(f\"Debug mode: {debug}\")\n",
    "print(f\"Repository suffix: {repo_suffix}\")\n",
    "\n",
    "if not re.match(r'^[MF]C?[0-9]{2}$', speaker_id):\n",
    "    print(\"Please provide a valid speaker ID.\")\n",
    "    sys.exit(1)\n",
    "test_speaker = speaker_id\n",
    "\n",
    "if repo_suffix and not re.match(r'^[_-]', args.repo_suffix):\n",
    "    repo_suffix = '_' + repo_suffix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fe6642c7-ee1f-407d-8969-839e6d245caa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The CSV file does not exist.\n",
      "torgo_dataset_path: /work/van-speech-nlp/data/torgo\n",
      "torgo_dataset_dir_path: /work/van-speech-nlp/data/torgo/\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import logging\n",
    "\n",
    "# Define the path to the CSV file\n",
    "torgo_csv_path = \"data_preparation/torgo.csv\"\n",
    "\n",
    "# Check if the path exists and is a file\n",
    "if os.path.exists(torgo_csv_path) and os.path.isfile(torgo_csv_path):\n",
    "    print(\"The CSV file exists.\")\n",
    "else:\n",
    "    print(\"The CSV file does not exist.\")\n",
    "\n",
    "torgo_dataset_path = '/work/van-speech-nlp/data/torgo'\n",
    "torgo_dataset_dir_path = torgo_dataset_path + \\\n",
    "        '/' if torgo_dataset_path[-1] != '/' else torgo_dataset_path\n",
    "output_path = 'output'\n",
    "print(f'torgo_dataset_path: {torgo_dataset_path}')\n",
    "print(f'torgo_dataset_dir_path: {torgo_dataset_dir_path}')\n",
    "\n",
    "repo_name = f'torgo_tiny_finetune_{test_speaker}{repo_suffix}'\n",
    "repo_path = f'jindaxz/{repo_name}'\n",
    "\n",
    "# Path to save model / checkpoints{repo_name}'\n",
    "model_local_path = output_path + '/model/' + repo_name\n",
    "\n",
    "pretrained_model_name = \"openai/whisper-tiny\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6c3172be-fc35-4a17-811e-46f775a0d895",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA\n"
     ]
    }
   ],
   "source": [
    "# moving to the whisper folder ; make sure you have the whisper environment on\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8bf16a5d-64b2-48ba-b789-ff3b6b4b2088",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3_prepare_for_torgo_baseline.py       \u001b[0m\u001b[01;34mimages\u001b[0m/\n",
      "3_whisper_finetuned_generate_json.py  \u001b[01;34mInference\u001b[0m/\n",
      "4_output_feature_pt.py                LICENSE\n",
      "convert_hf_checkpoint.py              \u001b[01;34mlit_llama\u001b[0m/\n",
      "\u001b[01;34mdata_preparation\u001b[0m/                     README.md\n",
      "demo.ipynb                            requirements.txt\n",
      "EMNLP__4032_Whispering_LLaMA_.pdf     \u001b[01;34mruns\u001b[0m/\n",
      "environment.yml                       \u001b[01;34mtraining\u001b[0m/\n",
      "\u001b[01;34mfinetuned_whisper_output\u001b[0m/             Untitled.ipynb\n",
      "\u001b[01;34mgenerate\u001b[0m/                             \u001b[01;34mweights\u001b[0m/\n",
      "\u001b[01;32mgitignore.sh\u001b[0m*                         \u001b[01;34mwhisper_openAI\u001b[0m/\n"
     ]
    }
   ],
   "source": [
    "ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a1f6c9fb-59c0-4b3f-8730-9e958c9e9781",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The file exists.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "MODEL_PATH = f\"finetuned_whisper_output/model/torgo_tiny_finetune_{speaker_id}/pytorch_model.bin\"\n",
    "\n",
    "if os.path.exists(MODEL_PATH):\n",
    "    print(\"The file exists.\")\n",
    "else:\n",
    "    print(\"The file does not exist.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "30cd2546-7ec2-4d44-a9e4-3454473d3acf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import re\n",
    "import whisper_openAI.whisper as whisper\n",
    "\n",
    "# https://github.com/openai/whisper/discussions/830\n",
    "def hf_to_whisper_states(text):\n",
    "    text = re.sub('.layers.', '.blocks.', text)\n",
    "    text = re.sub('.self_attn.', '.attn.', text)\n",
    "    text = re.sub('.q_proj.', '.query.', text)\n",
    "    text = re.sub('.k_proj.', '.key.', text)\n",
    "    text = re.sub('.v_proj.', '.value.', text)\n",
    "    text = re.sub('.out_proj.', '.out.', text)\n",
    "    text = re.sub('.fc1.', '.mlp.0.', text)\n",
    "    text = re.sub('.fc2.', '.mlp.2.', text)\n",
    "    text = re.sub('.fc3.', '.mlp.3.', text)\n",
    "    text = re.sub('.fc3.', '.mlp.3.', text)\n",
    "    text = re.sub('.encoder_attn.', '.cross_attn.', text)\n",
    "    text = re.sub('.cross_attn.ln.', '.cross_attn_ln.', text)\n",
    "    text = re.sub('.embed_positions.weight', '.positional_embedding', text)\n",
    "    text = re.sub('.embed_tokens.', '.token_embedding.', text)\n",
    "    text = re.sub('model.', '', text)\n",
    "    text = re.sub('attn.layer_norm.', 'attn_ln.', text)\n",
    "    text = re.sub('.final_layer_norm.', '.mlp_ln.', text)\n",
    "    text = re.sub('encoder.layer_norm.', 'encoder.ln_post.', text)\n",
    "    text = re.sub('decoder.layer_norm.', 'decoder.ln.', text)\n",
    "    text = re.sub('proj_out.weight', 'decoder.token_embedding.weight', text)\n",
    "    return text\n",
    "\n",
    "# Load HF Model\n",
    "hf_state_dict = torch.load(MODEL_PATH)    # pytorch_model.bin file\n",
    "# print(hf_state_dict)\n",
    "\n",
    "# Rename layers\n",
    "for key in list(hf_state_dict.keys())[:]:\n",
    "    new_key = hf_to_whisper_states(key)\n",
    "    hf_state_dict[new_key] = hf_state_dict.pop(key)\n",
    "\n",
    "# Init Whisper Model and replace model weights\n",
    "model,_ = whisper.load_model('tiny')\n",
    "model.load_state_dict(hf_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f1b919fa-911e-4950-99a0-62adfe374e1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_df = pd.read_csv(torgo_csv_path)\n",
    "dataset_csv = load_dataset('csv', data_files=torgo_csv_path)\n",
    "\n",
    "# Check if the following columns exist in the dataset ['session', 'audio', 'text', 'speaker_id']\n",
    "expected_columns = ['session', 'audio', 'text', 'speaker_id']\n",
    "not_found_columns = []\n",
    "for column in expected_columns:\n",
    "    if column not in dataset_csv['train'].column_names:\n",
    "        not_found_columns.append(column)\n",
    "\n",
    "if len(not_found_columns) > 0:\n",
    "    logging.error(\n",
    "        \"The following columns are not found in the dataset:\" + \" [\" + \", \".join(not_found_columns) + \"]\")\n",
    "    sys.exit(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "62137e32-3fe1-49d8-91d2-e5c95824b673",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90bc44ca55df4233b897a5c5d7045ed7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/16394 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3c2b2b67b39d44cb8e136a577e6914cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/16394 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dcc44672bc7940cf8188ffb5e3f2a525",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/16394 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "logging.info(\n",
    "    \"Splitting the dataset into training / validation / test sets...\")\n",
    "\n",
    "# Extract the unique speakers in the dataset\n",
    "speakers = data_df['speaker_id'].unique()\n",
    "\n",
    "logging.info(\"Unique speakers found in the dataset:\")\n",
    "logging.info(str(speakers) + '\\n')\n",
    "\n",
    "if test_speaker not in speakers:\n",
    "    logging.error(\"Test Speaker not found in the dataset.\")\n",
    "    sys.exit(1)\n",
    "\n",
    "valid_speaker = 'F03' if test_speaker != 'F03' else 'F04'\n",
    "train_speaker = [s for s in speakers if s not in [\n",
    "    test_speaker, valid_speaker]]\n",
    "\n",
    "torgo_dataset = DatasetDict()\n",
    "torgo_dataset['train'] = dataset_csv['train'].filter(\n",
    "    lambda x: x in train_speaker, input_columns=['speaker_id'])\n",
    "torgo_dataset['validation'] = dataset_csv['train'].filter(\n",
    "    lambda x: x == valid_speaker, input_columns=['speaker_id'])\n",
    "torgo_dataset['test'] = dataset_csv['train'].filter(\n",
    "    lambda x: x == test_speaker, input_columns=['speaker_id'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9661b983-a215-4188-b015-e49df718639d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb9bffb20f3b4e1094be76104c8c1a0c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/15091 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e97392ea3a3149ad9dc31fc0a7aa3ecf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/1075 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea6c6dcbbade4b079d40ee1e442533b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/228 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "original_data_count = {'train': len(torgo_dataset['train']), 'validation': len(\n",
    "    torgo_dataset['validation']), 'test': len(torgo_dataset['test'])}\n",
    "\n",
    "if not keep_all_data:\n",
    "    # Update the three dataset splits (if ['test_data'] == 1, keep in test, if ['test_data'] == 0, keep in train and validation)\n",
    "    torgo_dataset['train'] = torgo_dataset['train'].filter(\n",
    "        lambda x: x['test_data'] == 0)\n",
    "    torgo_dataset['validation'] = torgo_dataset['validation'].filter(\n",
    "        lambda x: x['test_data'] == 0)\n",
    "    torgo_dataset['test'] = torgo_dataset['test'].filter(\n",
    "        lambda x: x['test_data'] == 1)\n",
    "\n",
    "    # Drop the 'test_data' column\n",
    "    torgo_dataset['train'] = torgo_dataset['train'].remove_columns([\n",
    "                                                                   'test_data'])\n",
    "    torgo_dataset['validation'] = torgo_dataset['validation'].remove_columns([\n",
    "                                                                             'test_data'])\n",
    "    torgo_dataset['test'] = torgo_dataset['test'].remove_columns([\n",
    "                                                                 'test_data'])\n",
    "    logging.info(\n",
    "        f\"After removal of repeated prompts, the number of data in each dataset is:\")\n",
    "    logging.info(\n",
    "        f'Train:       {len(torgo_dataset[\"train\"])}/{original_data_count[\"train\"]} ({len(torgo_dataset[\"train\"]) * 100 // original_data_count[\"train\"]}%)')\n",
    "    logging.info(\n",
    "        f'Validation:  {len(torgo_dataset[\"validation\"])}/{original_data_count[\"validation\"]} ({len(torgo_dataset[\"validation\"]) * 100 // original_data_count[\"validation\"]}%)')\n",
    "    logging.info(\n",
    "        f'Test:        {len(torgo_dataset[\"test\"])}/{original_data_count[\"test\"]} ({len(torgo_dataset[\"test\"]) * 100 // original_data_count[\"test\"]}%)\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "27ceed6c-6a19-42f4-bb21-354d705b47d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83ed0ffb2b09493e9e402dfb021c8fd6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/9749 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5103bf6e12044e8a2119cac3a374f1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/483 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea895fd33cf84ae6b264e91e9b42bfd9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/126 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Remove special characters from the text\n",
    "chars_to_ignore_regex = '[\\,\\?\\.\\!\\-\\;\\:\\\"\\“\\%\\‘\\”\\`\\�0-9]'\n",
    "\n",
    "\n",
    "def remove_special_characters(batch):\n",
    "    batch['text'] = re.sub(chars_to_ignore_regex,\n",
    "                           ' ', batch['text']).lower()\n",
    "    return batch\n",
    "\n",
    "torgo_dataset = torgo_dataset.map(remove_special_characters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b7828ba7-163f-4079-bd1e-b7eac2f36f16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "air\n"
     ]
    }
   ],
   "source": [
    "print(torgo_dataset['train'][2]['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "7509ddf9-1877-49ef-a886-56160bc02b92",
   "metadata": {},
   "outputs": [],
   "source": [
    "## convert the sample rate of every audio files using cast_column function\n",
    "torgo_dataset = torgo_dataset.cast_column(\"audio\", Audio(sampling_rate=16000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "00ea9751-9bd6-47ef-a8b8-21d270dc7b87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e008d10c9b8948548201590c13fca544",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/9749 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d40e03bd7bbb42ccb6e787c98e2bf9b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/483 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3435b7887e304834b424feb814bc6c29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/126 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59f43868284f486fa4993b8d5c5fc5a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/9669 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b2d4c83c5454cf8999cb249d6070cf5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/482 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa4a189174254042addce6f53d2bd92b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/122 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Define the minimum and maximum input length in seconds\n",
    "min_input_length_in_sec = 1.0\n",
    "max_input_length_in_sec = 10.0\n",
    "sampling_rate=16000\n",
    "\n",
    "# Define the filtering functions based on input length\n",
    "def filter_min_length(example):\n",
    "    return example[\"audio\"][\"array\"].shape[0] > min_input_length_in_sec * sampling_rate\n",
    "\n",
    "def filter_max_length(example):\n",
    "    return example[\"audio\"][\"array\"].shape[0] < max_input_length_in_sec * sampling_rate\n",
    "\n",
    "# Apply the filters\n",
    "torgo_dataset = torgo_dataset.filter(filter_max_length)\n",
    "torgo_dataset = torgo_dataset.filter(filter_min_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b085e4e5-06e1-415d-8f1d-81d7b5947186",
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.info(\n",
    "    f\"After filter, the number of data in each dataset is:\")\n",
    "logging.info(\n",
    "    f'Train:       {len(torgo_dataset[\"train\"])}/{original_data_count[\"train\"]} ({len(torgo_dataset[\"train\"]) * 100 // original_data_count[\"train\"]}%)')\n",
    "logging.info(\n",
    "    f'Validation:  {len(torgo_dataset[\"validation\"])}/{original_data_count[\"validation\"]} ({len(torgo_dataset[\"validation\"]) * 100 // original_data_count[\"validation\"]}%)')\n",
    "logging.info(\n",
    "    f'Test:        {len(torgo_dataset[\"test\"])}/{original_data_count[\"test\"]} ({len(torgo_dataset[\"test\"]) * 100 // original_data_count[\"test\"]}%)\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "cd000aab-930a-4912-b8cb-5fdfea67e6b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = torgo_dataset[\"train\"]\n",
    "validation_dataset = torgo_dataset[\"validation\"]\n",
    "test_dataset = torgo_dataset[\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "0882a400-7a8c-49ec-b6ab-586dab78f258",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "460"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(validation_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "187c8f89-646a-48bd-9a85-3225d199ee19",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  9%|▉         | 874/9638 [04:12<42:14,  3.46it/s]  \n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[32], line 52\u001b[0m\n\u001b[1;32m     48\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mopen\u001b[39m(save_path, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mw\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m file:\n\u001b[1;32m     49\u001b[0m         json\u001b[38;5;241m.\u001b[39mdump(to_json, file, indent\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m4\u001b[39m)\n\u001b[0;32m---> 52\u001b[0m \u001b[43mgenerate_inference_json\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrain_dataset\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mtorgo_train_\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mspeaker_id\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[32], line 30\u001b[0m, in \u001b[0;36mgenerate_inference_json\u001b[0;34m(dataset, dataset_name)\u001b[0m\n\u001b[1;32m     28\u001b[0m         random_temperature \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.1\u001b[39m\n\u001b[1;32m     29\u001b[0m     options \u001b[38;5;241m=\u001b[39m whisper\u001b[38;5;241m.\u001b[39mDecodingOptions(fp16\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, without_timestamps\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, temperature\u001b[38;5;241m=\u001b[39mrandom_temperature, best_of\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m200\u001b[39m,language\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124menglish\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m---> 30\u001b[0m     result, _ \u001b[38;5;241m=\u001b[39m \u001b[43mwhisper\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdecode\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptions\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(result)\n\u001b[1;32m     33\u001b[0m to_json\u001b[38;5;241m.\u001b[39mappend({\n\u001b[1;32m     34\u001b[0m     item[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124msession\u001b[39m\u001b[38;5;124m'\u001b[39m]: {\n\u001b[1;32m     35\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtemp\u001b[39m\u001b[38;5;124m'\u001b[39m: random_temperature,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     43\u001b[0m     }\n\u001b[1;32m     44\u001b[0m })\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/mmenv/lib/python3.9/site-packages/torch/utils/_contextlib.py:115\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    114\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 115\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA/whisper_openAI/whisper/decoding.py:843\u001b[0m, in \u001b[0;36mdecode\u001b[0;34m(model, mel, options, **kwargs)\u001b[0m\n\u001b[1;32m    840\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m kwargs:\n\u001b[1;32m    841\u001b[0m     options \u001b[38;5;241m=\u001b[39m replace(options, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 843\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mDecodingTask\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptions\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmel\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    845\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/mmenv/lib/python3.9/site-packages/torch/utils/_contextlib.py:115\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    114\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 115\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA/whisper_openAI/whisper/decoding.py:740\u001b[0m, in \u001b[0;36mDecodingTask.run\u001b[0;34m(self, mel)\u001b[0m\n\u001b[1;32m    737\u001b[0m tokens \u001b[38;5;241m=\u001b[39m tokens\u001b[38;5;241m.\u001b[39mrepeat_interleave(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn_group, dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\u001b[38;5;241m.\u001b[39mto(audio_features\u001b[38;5;241m.\u001b[39mdevice)  \u001b[38;5;66;03m# ([200, 4])\u001b[39;00m\n\u001b[1;32m    739\u001b[0m \u001b[38;5;66;03m# call the main sampling loop\u001b[39;00m\n\u001b[0;32m--> 740\u001b[0m tokens, sum_logprobs, no_speech_probs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_main_loop\u001b[49m\u001b[43m(\u001b[49m\u001b[43maudio_features\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtokens\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    742\u001b[0m \u001b[38;5;66;03m# reshape the tensors to have (n_audio, n_group) as the first two dimensions (lost the no of samples)\u001b[39;00m\n\u001b[1;32m    743\u001b[0m audio_features \u001b[38;5;241m=\u001b[39m audio_features[:: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn_group]\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA/whisper_openAI/whisper/decoding.py:689\u001b[0m, in \u001b[0;36mDecodingTask._main_loop\u001b[0;34m(self, audio_features, tokens)\u001b[0m\n\u001b[1;32m    687\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    688\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msample_len):\n\u001b[0;32m--> 689\u001b[0m         logits \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minference\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlogits\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtokens\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maudio_features\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    691\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m    692\u001b[0m             i \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtokenizer\u001b[38;5;241m.\u001b[39mno_speech \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    693\u001b[0m         ):  \u001b[38;5;66;03m# save no_speech_probs\u001b[39;00m\n\u001b[1;32m    694\u001b[0m             probs_at_sot \u001b[38;5;241m=\u001b[39m logits[:, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msot_index]\u001b[38;5;241m.\u001b[39mfloat()\u001b[38;5;241m.\u001b[39msoftmax(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA/whisper_openAI/whisper/decoding.py:162\u001b[0m, in \u001b[0;36mPyTorchInference.logits\u001b[0;34m(self, tokens, audio_features)\u001b[0m\n\u001b[1;32m    158\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m tokens\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m] \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minitial_token_length:\n\u001b[1;32m    159\u001b[0m     \u001b[38;5;66;03m# only need to use the last token except in the first forward pass\u001b[39;00m\n\u001b[1;32m    160\u001b[0m     tokens \u001b[38;5;241m=\u001b[39m tokens[:, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m:]\n\u001b[0;32m--> 162\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdecoder\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtokens\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maudio_features\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkv_cache\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkv_cache\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/mmenv/lib/python3.9/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/mmenv/lib/python3.9/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA/whisper_openAI/whisper/model.py:211\u001b[0m, in \u001b[0;36mTextDecoder.forward\u001b[0;34m(self, x, xa, kv_cache)\u001b[0m\n\u001b[1;32m    208\u001b[0m x \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mto(xa\u001b[38;5;241m.\u001b[39mdtype)\n\u001b[1;32m    210\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m block \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mblocks:\n\u001b[0;32m--> 211\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[43mblock\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mxa\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmask\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkv_cache\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkv_cache\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    213\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mln(x)\n\u001b[1;32m    214\u001b[0m logits \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    215\u001b[0m     x \u001b[38;5;241m@\u001b[39m torch\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtoken_embedding\u001b[38;5;241m.\u001b[39mweight\u001b[38;5;241m.\u001b[39mto(x\u001b[38;5;241m.\u001b[39mdtype), \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m    216\u001b[0m )\u001b[38;5;241m.\u001b[39mfloat()\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/mmenv/lib/python3.9/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/mmenv/lib/python3.9/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA/whisper_openAI/whisper/model.py:139\u001b[0m, in \u001b[0;36mResidualAttentionBlock.forward\u001b[0;34m(self, x, xa, mask, kv_cache)\u001b[0m\n\u001b[1;32m    137\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcross_attn:\n\u001b[1;32m    138\u001b[0m     x \u001b[38;5;241m=\u001b[39m x \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcross_attn(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcross_attn_ln(x), xa, kv_cache\u001b[38;5;241m=\u001b[39mkv_cache)[\u001b[38;5;241m0\u001b[39m]\n\u001b[0;32m--> 139\u001b[0m x \u001b[38;5;241m=\u001b[39m x \u001b[38;5;241m+\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmlp\u001b[49m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmlp_ln(x))\n\u001b[1;32m    140\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m x\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/mmenv/lib/python3.9/site-packages/torch/nn/modules/module.py:1692\u001b[0m, in \u001b[0;36mModule.__getattr__\u001b[0;34m(self, name)\u001b[0m\n\u001b[1;32m   1690\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m _buffers[name]\n\u001b[1;32m   1691\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m_modules\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__dict__\u001b[39m:\n\u001b[0;32m-> 1692\u001b[0m     modules \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__dict__\u001b[39;49m[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m_modules\u001b[39m\u001b[38;5;124m'\u001b[39m]\n\u001b[1;32m   1693\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m name \u001b[38;5;129;01min\u001b[39;00m modules:\n\u001b[1;32m   1694\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m modules[name]\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import json\n",
    "import os\n",
    "import tqdm\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def generate_inference_json(dataset, dataset_name):\n",
    "    to_json = []\n",
    "    for i, item in enumerate(tqdm.tqdm(dataset)):\n",
    "        # print(item)\n",
    "        audio = item['audio']['array'].astype(np.single)\n",
    "        audio = whisper.pad_or_trim(audio)\n",
    "        mel = whisper.log_mel_spectrogram(audio).to(model.device)\n",
    "        ground_truth = item['text'].replace(' <COMMA>', ',').replace(' <PERIOD>', '.').replace(' <QUESTIONMARK>', '?').replace(' <EXCLAMATIONPOINT>', '!').lower()\n",
    "        source = 'NP-Torgo'\n",
    "        cat = 'NP-Torgo'\n",
    "        time = len(audio)/16000\n",
    "        path_to_file = item['audio']['path']\n",
    "        random_temperature = np.random.randint(70, 81) / 100\n",
    "        options = whisper.DecodingOptions(fp16=True, without_timestamps=True, temperature=random_temperature, best_of=200)\n",
    "        result, _ = whisper.decode(model, mel, options)\n",
    "        result = list(result)\n",
    "\n",
    "        if len(result) <= 10:\n",
    "            if random_temperature < 0.75:\n",
    "                random_temperature += 0.2\n",
    "            else:\n",
    "                random_temperature += 0.1\n",
    "            options = whisper.DecodingOptions(fp16=True, without_timestamps=True, temperature=random_temperature, best_of=200,language='english')\n",
    "            result, _ = whisper.decode(model, mel, options)\n",
    "            result = list(result)\n",
    "\n",
    "        to_json.append({\n",
    "            item['session']: {\n",
    "                'temp': random_temperature,\n",
    "                'path': path_to_file,\n",
    "                'ground_truth': ground_truth,\n",
    "                'inference': result,\n",
    "                'source': source,\n",
    "                'category': cat,\n",
    "                'time': time,\n",
    "                'path': path_to_file\n",
    "            }\n",
    "        })\n",
    "\n",
    "    os.makedirs(f\"Inference/gs_inferences\", exist_ok=True)\n",
    "    save_path = f'Inference/gs_inferences/{str(dataset_name)}.json'\n",
    "    with open(save_path, \"w\") as file:\n",
    "        json.dump(to_json, file, indent=4)\n",
    "\n",
    "\n",
    "generate_inference_json(train_dataset, f'finetuned_torgo_train_{speaker_id}')\n",
    "generate_inference_json(validation_dataset, f'finetuned_torgo_val_{speaker_id}')\n",
    "generate_inference_json(test_dataset, f'finetuned_torgo_test_{speaker_id}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "17ec8e30-7f74-4b42-a4fb-a006bf12fa9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dataclasses import dataclass, field, replace\n",
    "from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Sequence, Tuple, Union\n",
    "import whisper_openAI.whisper as whisper\n",
    "import torch\n",
    "from whisper_openAI.whisper.tokenizer import Tokenizer, get_tokenizer\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor\n",
    "\n",
    "# We get the acoustic embeddings from Whisper Large V2\n",
    "model,processor = whisper.load_model(\"large-v2\")\n",
    "# model,processor = whisper.load_model(\"medium\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0f883085-5df8-470f-9480-f78577aa85af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Whisper(\n",
       "  (encoder): AudioEncoder(\n",
       "    (conv1): Conv1d(80, 1280, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "    (conv2): Conv1d(1280, 1280, kernel_size=(3,), stride=(2,), padding=(1,))\n",
       "    (blocks): ModuleList(\n",
       "      (0-31): 32 x ResidualAttentionBlock(\n",
       "        (attn): MultiHeadAttention(\n",
       "          (query): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (key): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "          (value): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "        (attn_ln): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (0): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "          (1): GELU(approximate='none')\n",
       "          (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "        )\n",
       "        (mlp_ln): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "    )\n",
       "    (ln_post): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (decoder): TextDecoder(\n",
       "    (token_embedding): Embedding(51865, 1280)\n",
       "    (blocks): ModuleList(\n",
       "      (0-31): 32 x ResidualAttentionBlock(\n",
       "        (attn): MultiHeadAttention(\n",
       "          (query): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (key): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "          (value): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "        (attn_ln): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        (cross_attn): MultiHeadAttention(\n",
       "          (query): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (key): Linear(in_features=1280, out_features=1280, bias=False)\n",
       "          (value): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "          (out): Linear(in_features=1280, out_features=1280, bias=True)\n",
       "        )\n",
       "        (cross_attn_ln): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (0): Linear(in_features=1280, out_features=5120, bias=True)\n",
       "          (1): GELU(approximate='none')\n",
       "          (2): Linear(in_features=5120, out_features=1280, bias=True)\n",
       "        )\n",
       "        (mlp_ln): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "    )\n",
       "    (ln): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6b3b2d8e-d4da-4847-a692-dcf696a14b0a",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'speaker_id' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 5\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mjson\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# The below is the json file you can generate using the \"To generatn-best hyporhesis.ipynb\" notebook; Need to further tokenize the hypothesis\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mopen\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mInference/gs_inferences/torgo_train_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mspeaker_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.json\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m file:  \u001b[38;5;66;03m# Change the file path and name here\u001b[39;00m\n\u001b[1;32m      6\u001b[0m     train_data \u001b[38;5;241m=\u001b[39m json\u001b[38;5;241m.\u001b[39mload(file)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'speaker_id' is not defined"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "# The below is the json file you can generate using the \"To generatn-best hyporhesis.ipynb\" notebook; Need to further tokenize the hypothesis\n",
    "\n",
    "with open(f'Inference/gs_inferences/torgo_train_{speaker_id}.json', \"r\") as file:  # Change the file path and name here\n",
    "    train_data = json.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60a961a9-d60f-4531-bbc6-620a43c55505",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'Inference/gs_inferences/torgo_val_{speaker_id}.json', \"r\") as valid_file:\n",
    "    val_data = json.load(valid_file)\n",
    "\n",
    "# Load the test set\n",
    "with open(f'Inference/gs_inferences/torgo_test_{speaker_id}.json', \"r\") as test_file:\n",
    "    test_data = json.load(test_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5678e1c3-ec82-45b3-a735-a4529987dec2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Implementation derived from https://github.com/tloen/alpaca-lora\"\"\"\n",
    "import sys\n",
    "from pathlib import Path\n",
    "import torch\n",
    "import requests\n",
    "import json\n",
    "import os \n",
    "\n",
    "from lit_llama.tokenizer import Tokenizer\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2875f9ca-7ce5-4514-9a13-98db217cad68",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_path: Path = Path(\"weights/tokenizer.model\")\n",
    "tokenizer = Tokenizer(tokenizer_path)\n",
    "print(f\"train has {len(train_data):,} samples\")\n",
    "print(\"Processing train split ...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a11b42-f260-4d15-b80c-fa5e04f833ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(tokenizer: Tokenizer, string: str, max_length: int, eos=True) -> torch.Tensor:\n",
    "    return tokenizer.encode(string, bos=True, eos=eos, max_length=max_length)\n",
    "    \n",
    "def process_train_data(train_data):\n",
    "    instruction = 'You are an ASR transcript selector. You have a few transcripts generated by an automatic speech recognition model. Your task is to generate the most likely transcript from them. If the generated transcripts have grammatical or logical errors, you will modify them accordingly to produce the most accurate and coherent transcript.'\n",
    "    result = []\n",
    "\n",
    "    for i in tqdm(range(len(train_data))):        \n",
    "        for name in train_data[i].keys():\n",
    "            ip = train_data[i][name]\n",
    "        inference = ip['inference']\n",
    "        gt = ip['ground_truth']\n",
    "            \n",
    "        # Removing the ground_truth, if present among the inferences for the prompt\n",
    "        if gt in inference:\n",
    "            inference.remove(gt)\n",
    "                \n",
    "        # Joining the inputs with '\\n'\n",
    "        for_input = '\\n'.join(inference[:15])\n",
    "        # The prompt follows the Alpaca template\n",
    "        full_prompt = f\"\"\"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\\n\\n### Instruction:\\n{instruction}\\n\\n### Input:\\n{for_input}\\n\\n### Response:\"\"\"\n",
    "        full_prompt_and_response = full_prompt + gt\n",
    "\n",
    "        encoded_full_prompt = tokenize(tokenizer, full_prompt, max_length=2048, eos=False)\n",
    "        encoded_full_prompt_and_response = tokenize(tokenizer, full_prompt_and_response, eos=True, max_length=2048)\n",
    "        labels = encoded_full_prompt_and_response.clone()\n",
    "        labels_with_masked_input = encoded_full_prompt_and_response.clone()\n",
    "        labels_with_masked_input[:len(encoded_full_prompt)] = -1\n",
    "        \n",
    "        path = ip['path']\n",
    "        audio = whisper.load_audio(path)  \n",
    "        audio = whisper.pad_or_trim(audio)            \n",
    "        mel = whisper.log_mel_spectrogram(audio).to(model.device)  # Adjust as needed for your model\n",
    "        mel = mel.unsqueeze(0)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            audio_features = model.encoder(mel)\n",
    "        \n",
    "        result.append({**ip, 'index': name, \"input_ids\": encoded_full_prompt_and_response, \"input_ids_no_response\": encoded_full_prompt, \"labels\": labels, 'labels_with_masked_input': labels_with_masked_input, 'audio_features': audio_features.bfloat16()})\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c61f584-a74c-4e4d-9101-548d68a51f0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "split = \"train\"\n",
    "result = process_train_data(train_data)\n",
    "torch.save(result,f'Inference/gs_inferences/finetuned_torgo_{speaker_id}_{split}.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88ea6dea-63aa-4f22-8d58-b7015ae11aca",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "result = process_train_data(validation_dataset)\n",
    "torch.save(result,f'Inference/gs_inferences/finetuned_torgo_{speaker_id}_{split}_en.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c5f0f66-9cf6-43c0-a9dc-e555dffe026b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "result = process_train_data(test_data)\n",
    "torch.save(result,f'Inference/gs_inferences/finetuned_torgo_{speaker_id}_{split}_en.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b914ae70-99ca-49dc-9077-bae62d27f30a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}