lj_speech_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01d1f516-d1fc-4bff-b3dd-083984a81a91",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2599a77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = Path('./data/lj_speech')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e4d3362",
   "metadata": {},
   "source": [
    "# Feature Extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5028542d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import soundfile as sf\n",
    "from datasets import load_dataset\n",
    "from torchaudio.transforms import Resample\n",
    "from transformers import Wav2Vec2Model, Wav2Vec2FeatureExtractor\n",
    "from gpt2_asr.speech_feature_extraction import extract_features_to_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c678e39c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not DATA_PATH.exists():\n",
    "    \n",
    "    lj_speech = load_dataset('lj_speech', split='train')  # Is only train.\n",
    "    \n",
    "    # Load Model for feature extraction.\n",
    "    wave2vec_name = 'facebook/wav2vec2-large-960h-lv60-self'\n",
    "    wave2vec_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wave2vec_name)\n",
    "    wave2vec = Wav2Vec2Model.from_pretrained(wave2vec_name)\n",
    "    wave2vec.eval().cuda()\n",
    "    resampler = Resample(22050, 16_000)\n",
    "    \n",
    "    # Extract audio and transcriptions.\n",
    "    examples = []\n",
    "    for i, eg in enumerate(lj_speech):\n",
    "        audio = torch.tensor(sf.read(eg['file'])[0])\n",
    "        eg['audio'] = resampler(audio.float()).numpy()\n",
    "\n",
    "        # TODO: Temporary (?) Helper for generation since using empty `input_ids` lead to errors.\n",
    "        eg['transcription'] = 'Transcription: ' + eg['normalized_text']\n",
    "\n",
    "        examples.append(eg)\n",
    "        \n",
    "    # `max_len` is just the longest sample in the dataset (determined in advance).\n",
    "    extract_features_to_files(wave2vec, wave2vec_extractor, examples, \n",
    "                              batch_size=8, max_len=161540, output_path=DATA_PATH, val_pct=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b766206-420c-4023-aeb6-e82cf57d5108",
   "metadata": {},
   "outputs": [],
   "source": [
    "wave2vec.cpu()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "969fa3a2",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2c5c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim import AdamW\n",
    "from transformers import AutoModelForCausalLM, GPT2Tokenizer\n",
    "from gpt2_s2t.model import S2TModel\n",
    "from gpt2_s2t.data_loading import S2TDataset, make_collate_fn\n",
    "from gpt2_s2t.evaluation_utils import get_predictions, calculate_mean_loss, calculate_wer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80e5a6e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "accelerator = Accelerator(fp16=True)\n",
    "print(f'Using {accelerator.device}.')\n",
    "\n",
    "tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
    "tokenizer.add_special_tokens({'pad_token': '_'})\n",
    "gpt2_model = AutoModelForCausalLM.from_pretrained('gpt2', add_cross_attention=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e22a69",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = S2TModel(gpt2_model)\n",
    "best_val_loss = 10\n",
    "\n",
    "# # Sometimes loss suddenly increases by a lot. Then the best saved model can be loaded here and trained further.\n",
    "# # However, training also seems to converge again, so may not be necessary.\n",
    "# model.load_state_dict(torch.load('test.pt'))\n",
    "# best_val_loss = 1.56"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a1e67e",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_grouped_parameters = [\n",
    "    {\n",
    "        \"params\": [p for n, p in model.named_parameters() if ('crossattention' in n or 'projection' in n)],\n",
    "        \"lr\": 5e-4,\n",
    "    },\n",
    "    {\n",
    "        \"params\": [p for n, p in model.named_parameters() if ('crossattention' not in n and 'projection' not in n)],\n",
    "        \"lr\": 1e-6,\n",
    "    },\n",
    "]\n",
    "optimizer = AdamW(optimizer_grouped_parameters, weight_decay=0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68dbe280",
   "metadata": {},
   "outputs": [],
   "source": [
    "collate_fn = make_collate_fn(tokenizer)\n",
    "\n",
    "train_ds = S2TDataset(DATA_PATH / 'train')\n",
    "train_dl = DataLoader(train_ds, batch_size=16, shuffle=True, collate_fn=collate_fn, num_workers=8)\n",
    "\n",
    "val_ds = S2TDataset(DATA_PATH / 'val')\n",
    "val_dl = DataLoader(val_ds, batch_size=16, shuffle=False, collate_fn=collate_fn, num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6efc5ac0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model, optimizer, train_dl, val_dl = accelerator.prepare(model, optimizer, train_dl, val_dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bb85e60",
   "metadata": {},
   "source": [
    "## Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b566fb-ba93-4a53-8f0f-214823d2fe0b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "epochs = 40\n",
    "accumulate_gradients = 2\n",
    "\n",
    "for i in range(epochs):\n",
    "    \n",
    "    model.train()\n",
    "    for step, (encoder_hidden_states, input_ids) in enumerate(train_dl):\n",
    "        out = model(encoder_hidden_states, input_ids)\n",
    "        accelerator.backward(out.loss)\n",
    "        if step % 100 == 0:\n",
    "            print(out.loss.item())\n",
    "        if (step + 1) % accumulate_gradients == 0:\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "    model.eval()\n",
    "    val_loss = calculate_mean_loss(model, val_dl)\n",
    "    print('VAL: ', val_loss)\n",
    "    if val_loss < best_val_loss:\n",
    "        torch.save(model.state_dict(), 'test.pt')\n",
    "        best_val_loss = val_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da0afce2",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "346659a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_state_dict(torch.load('test.pt'))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "430064b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Validation loss:', calculate_mean_loss(model, val_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9514495d-6287-429d-96b1-c8d62dbeb2f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prediction on random validation example.\n",
    "eg = val_ds[np.random.randint(len(val_ds))]\n",
    "predicted_text = get_predictions(\n",
    "    eg['wave2vec_features'][None, ...].cuda(), model, tokenizer\n",
    ")[0]\n",
    "\n",
    "prefix_len = len('Transcription: ')\n",
    "print('PRED:', predicted_text[prefix_len:])\n",
    "print('GOLD:', eg['transcription'][prefix_len:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46f2f29e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Word error rate.\n",
    "val_preds = []\n",
    "for batch in val_dl:\n",
    "    encoder_hidden_states, _ = batch\n",
    "    val_preds += get_predictions(encoder_hidden_states, model, tokenizer)\n",
    "val_preds = [pred[prefix_len:] for pred in val_preds]\n",
    "val_golds = [eg['transcription'][prefix_len:] for eg in val_ds]\n",
    "\n",
    "print('WER:', calculate_wer(val_preds, val_golds))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}