{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "037ddb26",
   "metadata": {},
   "source": [
    "### Quick tutorial on how to perform forced alignment on a piece of a Chinese music"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88ea490d",
   "metadata": {},
   "source": [
    "**Check the environments**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bddc096d-4567-44cd-bc8e-52121f653039",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.5.1\n",
      "2.5.1\n",
      "cpu\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchaudio\n",
    "\n",
    "print(torch.__version__)\n",
    "print(torchaudio.__version__)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8969ecd8",
   "metadata": {},
   "source": [
    "**Import necessary libraires**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0ef77f11-0df6-44fe-955a-5da18af3f3cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "import IPython\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eec5c071",
   "metadata": {},
   "source": [
    "**Load the model**\n",
    "\n",
    "This may take some time, depending on your internet speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d52293b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchaudio.pipelines import MMS_FA as bundle\n",
    "\n",
    "model = bundle.get_model()\n",
    "model.to(device)\n",
    "\n",
    "tokenizer = bundle.get_tokenizer()\n",
    "aligner = bundle.get_aligner()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba94178f",
   "metadata": {},
   "source": [
    "**Preparing text scripts**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "05da7399-1244-43b4-b1db-b1e80eb04be5",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_normalized = \"yan yu mi san aaaaaaaaaaaa\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e4ab5d7",
   "metadata": {},
   "source": [
    "**Load & check the audio file**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "63e02aba",
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform, sample_rate = torchaudio.load(\"./data/sample3.mp3\")\n",
    "waveform = waveform[0:1]\n",
    "from torchaudio.transforms import Resample\n",
    "resampler = Resample(orig_freq=sample_rate, new_freq=16000)\n",
    "waveform = resampler(waveform)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "404c0d38",
   "metadata": {},
   "source": [
    "**Define some helper functions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "27065420",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_alignments(waveform: torch.Tensor, transcript: List[str]):\n",
    "    with torch.inference_mode():\n",
    "        emission, _ = model(waveform.to(device))\n",
    "        token_spans = aligner(emission[0], tokenizer(transcript))\n",
    "    return emission, token_spans\n",
    "\n",
    "# Compute average score weighted by the span length\n",
    "def _score(spans):\n",
    "    return sum(s.score * len(s) for s in spans) / sum(len(s) for s in spans)\n",
    "\n",
    "\n",
    "def plot_alignments(waveform, token_spans, emission, transcript, sample_rate=bundle.sample_rate):\n",
    "    ratio = waveform.size(1) / emission.size(1) / sample_rate\n",
    "\n",
    "    fig, axes = plt.subplots(2, 1)\n",
    "    axes[0].imshow(emission[0].detach().cpu().T, aspect=\"auto\")\n",
    "    axes[0].set_title(\"Emission\")\n",
    "    axes[0].set_xticks([])\n",
    "\n",
    "    axes[1].specgram(waveform[0], Fs=sample_rate)\n",
    "    for t_spans, chars in zip(token_spans, transcript):\n",
    "        t0, t1 = t_spans[0].start, t_spans[-1].end\n",
    "        axes[0].axvspan(t0 - 0.5, t1 - 0.5, facecolor=\"None\", hatch=\"/\", edgecolor=\"white\")\n",
    "        axes[1].axvspan(ratio * t0, ratio * t1, facecolor=\"None\", hatch=\"/\", edgecolor=\"white\")\n",
    "        axes[1].annotate(f\"{_score(t_spans):.2f}\", (ratio * t0, sample_rate * 0.51), annotation_clip=False)\n",
    "\n",
    "        for span, char in zip(t_spans, chars):\n",
    "            t0 = span.start * ratio\n",
    "            axes[1].annotate(char, (t0, sample_rate * 0.55), annotation_clip=False)\n",
    "\n",
    "    axes[1].set_xlabel(\"time [second]\")\n",
    "    fig.tight_layout()\n",
    "    \n",
    "def preview_word(waveform, spans, num_frames, idx, transcript, sample_rate=bundle.sample_rate):\n",
    "    ratio = waveform.size(1) / num_frames\n",
    "    x0 = int(ratio * spans[idx][0].start)\n",
    "    if idx+1<len(spans):\n",
    "        x1 = int(ratio * spans[idx+1][0].start)\n",
    "    else:\n",
    "        x1 = int(ratio * spans[idx][-1].end)\n",
    "    print(f\"{transcript[idx]} ({_score(spans[idx]):.2f}): {x0 / sample_rate:.3f} - {x1 / sample_rate:.3f} sec\")\n",
    "    segment = waveform[:, x0:x1]\n",
    "    return IPython.display.Audio(segment.numpy(), rate=sample_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0d7b3e5",
   "metadata": {},
   "source": [
    "**Preview the audio and the aligned result**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9aaaf40f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Normalized Transcript:  yan yu mi san aaaaaaaaaaaa\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 114556 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 301692 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "transcript = text_normalized.split()\n",
    "emission, token_spans = compute_alignments(waveform, transcript)\n",
    "num_frames = emission.size(1)\n",
    "\n",
    "plot_alignments(waveform, token_spans, emission, transcript)\n",
    "\n",
    "print(\"Normalized Transcript: \", text_normalized)\n",
    "IPython.display.Audio(waveform, rate=16000)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a28865a1",
   "metadata": {},
   "source": [
    "**Preview a single word**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "044467c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aaaaaaaaaaaa (0.02): 1.453 - 1.997 sec\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 23304 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preview_word(waveform, token_spans, num_frames, 4, transcript)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e5bf64d-6123-4555-a227-c3047d3564f5",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
