3_prepare_torgo_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 17,
   "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": 18,
   "id": "0338206e-5029-44f2-af94-01be319a4d03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Speaker ID: M01\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 = \"M01\"  # 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": 4,
   "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": 5,
   "id": "3a40f8d0-7522-404c-b968-15edc5038513",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "# Renamed the Whisepr repo (https://github.com/openai/whisper) with the changed decoding.py file as whisper_openAI\n",
    "import whisper_openAI.whisper as whisper\n",
    "import torch\n",
    "import tqdm\n",
    "model, _ = whisper.load_model(\"tiny\") # you can change the whisper model here to largev2 or large to swap the  model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "id": "62137e32-3fe1-49d8-91d2-e5c95824b673",
   "metadata": {},
   "outputs": [],
   "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": 8,
   "id": "9661b983-a215-4188-b015-e49df718639d",
   "metadata": {},
   "outputs": [],
   "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": 9,
   "id": "27ceed6c-6a19-42f4-bb21-354d705b47d1",
   "metadata": {},
   "outputs": [],
   "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": 10,
   "id": "b7828ba7-163f-4079-bd1e-b7eac2f36f16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "meat\n"
     ]
    }
   ],
   "source": [
    "print(torgo_dataset['train'][2]['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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": 12,
   "id": "00ea9751-9bd6-47ef-a8b8-21d270dc7b87",
   "metadata": {},
   "outputs": [],
   "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": 13,
   "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": 14,
   "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": 15,
   "id": "0882a400-7a8c-49ec-b6ab-586dab78f258",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "460"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(validation_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "187c8f89-646a-48bd-9a85-3225d199ee19",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 208/9502 [02:29<1:51:34,  1.39it/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[16], 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[16], line 21\u001b[0m, in \u001b[0;36mgenerate_inference_json\u001b[0;34m(dataset, dataset_name)\u001b[0m\n\u001b[1;32m     19\u001b[0m random_temperature \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mrandint(\u001b[38;5;241m70\u001b[39m, \u001b[38;5;241m81\u001b[39m) \u001b[38;5;241m/\u001b[39m \u001b[38;5;241m100\u001b[39m\n\u001b[1;32m     20\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)\n\u001b[0;32m---> 21\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     22\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(result)\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(result) \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m10\u001b[39m:\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:705\u001b[0m, in \u001b[0;36mDecodingTask._main_loop\u001b[0;34m(self, audio_features, tokens)\u001b[0m\n\u001b[1;32m    702\u001b[0m     logit_filter\u001b[38;5;241m.\u001b[39mapply(logits, tokens)\n\u001b[1;32m    704\u001b[0m \u001b[38;5;66;03m# expand the tokens tensor with the selected next tokens\u001b[39;00m\n\u001b[0;32m--> 705\u001b[0m tokens, completed \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdecoder\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mupdate\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtokens\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogits\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msum_logprobs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    707\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m completed \u001b[38;5;129;01mor\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[39mn_ctx:\n\u001b[1;32m    708\u001b[0m     \u001b[38;5;28;01mbreak\u001b[39;00m\n",
      "File \u001b[0;32m/work/van-speech-nlp/jindaznb/jslpnb/mllm_expriments/Whispering-LLaMA/whisper_openAI/whisper/decoding.py:290\u001b[0m, in \u001b[0;36mGreedyDecoder.update\u001b[0;34m(self, tokens, logits, sum_logprobs)\u001b[0m\n\u001b[1;32m    287\u001b[0m     next_tokens \u001b[38;5;241m=\u001b[39m Categorical(logits\u001b[38;5;241m=\u001b[39mlogits \u001b[38;5;241m/\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtemperature)\u001b[38;5;241m.\u001b[39msample()\n\u001b[1;32m    289\u001b[0m logprobs \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mlog_softmax(logits\u001b[38;5;241m.\u001b[39mfloat(), dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m--> 290\u001b[0m current_logprobs \u001b[38;5;241m=\u001b[39m \u001b[43mlogprobs\u001b[49m\u001b[43m[\u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marange\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlogprobs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mshape\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnext_tokens\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m    291\u001b[0m sum_logprobs \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m current_logprobs \u001b[38;5;241m*\u001b[39m (tokens[:, \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[39meot)\n\u001b[1;32m    293\u001b[0m next_tokens[tokens[:, \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[39meot] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39meot\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=50,language='english')\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=50,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'torgo_train_{speaker_id}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "id": "0f883085-5df8-470f-9480-f78577aa85af",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b3b2d8e-d4da-4847-a692-dcf696a14b0a",
   "metadata": {},
   "outputs": [],
   "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": [
    "result = process_train_data(train_data)\n",
    "torch.save(result,f'Inference/gs_inferences/torgo_{speaker_id}_train.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/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/torgo_{speaker_id}_{split}_en.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c6fa607-072a-4c31-9ee8-1e9d2b3d2e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print()"
   ]
  },
  {
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}