{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "104d435d-3b18-48d8-9571-d98ff2da94db",
   "metadata": {},
   "source": [
    "# Forced Alignment with Wav2Vec2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b721fbcd-bfe0-4ff3-8cb0-e4d1658d65bb",
   "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",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    "import IPython\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c05a95f8-14a3-42ba-97e1-74009573e286",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchaudio.transforms import Resample\n",
    "waveform, sample_rate = torchaudio.load(\"data/sample3.mp3\")\n",
    "\n",
    "waveform = waveform[0:1]\n",
    "resampler = Resample(orig_freq=sample_rate, new_freq=16000)\n",
    "waveform = resampler(waveform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a33c5447-319b-462f-b5bc-df6b4cee2bc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchaudio.pipelines import MMS_FA as bundle\n",
    "\n",
    "model = bundle.get_model()\n",
    "model.to(device)\n",
    "\n",
    "labels = bundle.get_labels()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e581a85f-6860-4b34-8552-42e388adc8b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "emissions, _ = model(waveform.to(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "091729e7-9287-42ef-9a82-2ffee8c2f1e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "emissions = torch.log_softmax(emissions, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d1a58f3e-5061-4be1-8bd4-9a419036dcee",
   "metadata": {},
   "outputs": [],
   "source": [
    "emission = emissions[0].cpu().detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3ffb9d66-655f-4b76-96b3-451570eea5bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 73304 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch.nn.functional as F\n",
    "\n",
    "probs = F.softmax(emission, dim=1)\n",
    "\n",
    "# 计算 ratio\n",
    "ratio = waveform.size(1) / emission.size()[0] / 16000\n",
    "\n",
    "# 获取 x 轴数据\n",
    "x_values = [i * ratio for i in range(probs.size()[0])]\n",
    "\n",
    "# 获取 y 轴数据\n",
    "y_values_8 = [probs[i][8] for i in range(probs.size()[0])]\n",
    "y_values_1 = [probs[i][1] for i in range(probs.size()[0])]\n",
    "y_values_4 = [probs[i][4] for i in range(probs.size()[0])]\n",
    "y_values_3 = [probs[i][0] for i in range(probs.size()[0])]\n",
    "\n",
    "# 创建图表\n",
    "plt.figure(figsize=(10, 6))\n",
    "\n",
    "# 绘制四条线\n",
    "plt.plot(x_values, y_values_8, label='s', color='green')\n",
    "plt.plot(x_values, y_values_1, label='a', color='red')\n",
    "plt.plot(x_values, y_values_3, label='e', color='orange')\n",
    "plt.plot(x_values, y_values_4, label='n', color='purple')\n",
    "\n",
    "# 添加标题和标签\n",
    "plt.title('Probabilities Over Time')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel('Emission Value')\n",
    "import numpy as np\n",
    "x_ticks = np.arange(0, 2.70, 0.2)\n",
    "plt.xticks(x_ticks)\n",
    "plt.xlim(0,2.7)\n",
    "\n",
    "# 添加图例\n",
    "plt.legend()\n",
    "\n",
    "# 显示图表\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa7021d-876f-4302-9b11-24e1046bf9b6",
   "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
}
