{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "aa1d00f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch version: 2.8.0+cu128\n",
      "Torchaudio version: 2.8.0+cu128\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchaudio\n",
    "from torchaudio.transforms import MelSpectrogram, AmplitudeToDB\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from IPython.display import Audio\n",
    "from pathlib import Path\n",
    "\n",
    "print(f\"PyTorch version: {torch.__version__}\")\n",
    "print(f\"Torchaudio version: {torchaudio.__version__}\")\n",
    "\n",
    "workspace_dir = Path.cwd().parent.parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a6d351e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16000\n",
      "torch.Size([1, 132800])\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 128060 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 354192 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(8.6255)\n",
      "tensor(8.6255) tensor(-9.5740)\n",
      "tensor(0.9009) tensor(-1.)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio id=\"scaled\" controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 354192 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<DisplayHandle display_id=scaled>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "waveform, sample_rate = torchaudio.load_with_torchcodec(workspace_dir.joinpath(\"data/test/F001_001.wav\").as_posix(), channels_first=True, )\n",
    "print(sample_rate)\n",
    "print(waveform.shape)\n",
    "if waveform.shape[0] > 1:\n",
    "    waveform = waveform.mean(dim=0, keepdim=True)\n",
    "\n",
    "display(Audio(waveform[:, sample_rate * 2:sample_rate * 2 + sample_rate * 3], rate=sample_rate))\n",
    "display(Audio(waveform, rate=sample_rate))\n",
    "waveform = waveform * 20.0\n",
    "print(waveform.max())\n",
    "# display(Audio(waveform, rate=sample_rate, normalize=False))\n",
    "\n",
    "print(waveform.max(), waveform.min())\n",
    "waveform = waveform / waveform.abs().max()\n",
    "print(waveform.max(), waveform.min())\n",
    "display(Audio(waveform, rate=sample_rate, element_id=\"scaled\", normalize=False), display_id=\"scaled\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cb24ba9",
   "metadata": {},
   "source": [
    "## Mel spectrogram based approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "59557614",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MelSpectrogram 配置:\n",
      "  sample_rate: 16000\n",
      "  n_fft: 400\n",
      "  win_length: 400\n",
      "  hop_length: 160\n",
      "  n_mels: 320\n",
      "  f_min: 0\n",
      "  f_max: 8000\n",
      "  power: 3.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/blacksheep/projekts/study_asr/.venv/lib/python3.11/site-packages/torchaudio/functional/functional.py:585: UserWarning: At least one mel filterbank has all zero values. The value for `n_mels` (320) may be set too high. Or, the value for `n_freqs` (201) may be set too low.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# Cell 3: 配置 MelSpectrogram 参数\n",
    "# 可交互式修改这些参数，实时查看效果\n",
    "CONFIG = {\n",
    "    \"sample_rate\": 16000,\n",
    "    \"n_fft\": 400,           # FFT 窗口大小\n",
    "    \"win_length\": 400,      # 窗口长度\n",
    "    \"hop_length\": 160,      # 帧移 (10ms @ 16kHz)\n",
    "    \"n_mels\": 80 * 4,           # Mel 滤波器数量\n",
    "    \"f_min\": 0,\n",
    "    \"f_max\": 8000,\n",
    "    \"power\": 3.0,           # 2=功率谱, 1=幅度谱\n",
    "}\n",
    "\n",
    "# 创建变换\n",
    "mel_transform = MelSpectrogram(**CONFIG)\n",
    "print(\"MelSpectrogram 配置:\")\n",
    "for k, v in CONFIG.items():\n",
    "    print(f\"  {k}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7b9f9e7f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 487352 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mel 频谱图形状: torch.Size([1, 320, 831])\n"
     ]
    }
   ],
   "source": [
    "# Cell 4: 计算并可视化 Mel 频谱图\n",
    "# 计算 Mel 频谱图\n",
    "mel_spec = mel_transform(waveform)\n",
    "\n",
    "log_mel_spec = AmplitudeToDB()(mel_spec)\n",
    "\n",
    "# 创建可视化\n",
    "fig, axes = plt.subplots(2, 1, figsize=(12, 8))\n",
    "\n",
    "# 原始波形\n",
    "axes[0].plot(waveform[0].numpy())\n",
    "axes[0].set_title(\"Original waveform\")\n",
    "axes[0].set_xlabel(\"sample\")\n",
    "axes[0].set_ylabel(\"Amplitude\")\n",
    "axes[0].grid(True)\n",
    "\n",
    "# Mel 频谱图\n",
    "im = axes[1].imshow(\n",
    "    log_mel_spec[0].numpy(),\n",
    "    cmap='gray',\n",
    "    aspect='auto',\n",
    "    origin='lower',\n",
    "    extent=[0, waveform.shape[-1] / sample_rate, 0, CONFIG[\"n_mels\"]]\n",
    ")\n",
    "axes[1].set_title(\"Log Mel Spectrogram\")\n",
    "axes[1].set_xlabel(\"Time (second)\")\n",
    "axes[1].set_ylabel(\"Mels\")\n",
    "plt.colorbar(im, ax=axes[1], format='%+2.0f dB')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"Mel 频谱图形状: {mel_spec.shape}\")  # (channel, n_mels, time_frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "56db5576",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 320, 831])\n"
     ]
    }
   ],
   "source": [
    "print(log_mel_spec.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "cfa5546a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 320, 831])\n",
      "CNN OUT: torch.Size([1, 32, 40, 208])\n",
      "transformer input:  torch.Size([1, 208, 640])\n",
      "tensor([208])\n",
      "conformer out:  torch.Size([1, 208, 200])\n",
      "conformer length out:  tensor([208])\n",
      "loss:  tensor(-3.9094, grad_fn=<MeanBackward0>)\n"
     ]
    }
   ],
   "source": [
    "from torch.nn import Conv2d, Linear, CTCLoss\n",
    "from torchaudio.models import Conformer\n",
    "from torch import Tensor\n",
    "import math\n",
    "\n",
    "print(log_mel_spec.shape)\n",
    "#torch.Size([1, 320, 831])\n",
    "\n",
    "filter_1 = Conv2d(in_channels=1, out_channels=16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
    "filter_2 = Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
    "filter_3 = Conv2d(in_channels=32, out_channels=32, kernel_size=(3, 3), stride=(2, 1), padding=(1, 1))\n",
    "proj = Linear(in_features=1280, out_features=640)\n",
    "ctc_head = Linear(in_features=640, out_features=200, bias=False)\n",
    "encoder = Conformer(\n",
    "    input_dim=640,\n",
    "    num_heads=8,\n",
    "    ffn_dim=2048,\n",
    "    num_layers=6,\n",
    "    depthwise_conv_kernel_size=31,\n",
    "    dropout=0.1,\n",
    "    \n",
    ")\n",
    "\n",
    "log_mel_spec = log_mel_spec.unsqueeze(1)\n",
    "\n",
    "out: Tensor = filter_1(log_mel_spec)\n",
    "out = filter_2(out)\n",
    "out = filter_3(out)\n",
    "print(\"CNN OUT:\", out.shape)\n",
    "# torch.Size([1, 32, 20, 208])\n",
    "\n",
    "transformer_input: Tensor = out.reshape([out.shape[0], out.shape[3], out.shape[1] * out.shape[2]])\n",
    "transformer_input = proj(transformer_input)\n",
    "print(\"transformer input: \", transformer_input.shape)\n",
    "# transformer_input = proj(transformer_input)\n",
    "# print(transformer_input.shape)\n",
    "#[N, T, d_model] where N = batch, T = seq len, d_model = feature\n",
    "lengths = torch.tensor([transformer_input.shape[1]], dtype=torch.long)\n",
    "print(lengths)\n",
    "out_latent_space, out_lengths = encoder(input=transformer_input, lengths=lengths)\n",
    "out_latent_space: Tensor = ctc_head(out_latent_space)\n",
    "out_latent_space: Tensor\n",
    "out_lengths:Tensor\n",
    "\n",
    "print(\"conformer out: \", out_latent_space.shape)\n",
    "print(\"conformer length out: \", out_lengths)\n",
    "\n",
    "criterion = CTCLoss(blank=0, zero_infinity=True, )\n",
    "dummy_target = torch.randint(0, 200, [1, math.floor(out_latent_space.shape[1] / 4)])\n",
    "loss: Tensor = criterion(log_probs=out_latent_space.reshape(out_latent_space.shape[0] * out_latent_space.shape[1], out_latent_space.shape[2]), targets=dummy_target, input_lengths=out_lengths, target_lengths=torch.tensor([dummy_target.shape[1]], dtype=torch.long))\n",
    "print(\"loss: \", loss)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "421d0749",
   "metadata": {},
   "source": [
    "## Conv1D over waveform approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "83fc7c1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Linear, Conv1d, ModuleList, GELU, GroupNorm\n",
    "from torchaudio.models import Conformer\n",
    "waveform = waveform.unsqueeze(0)[:, :, :16000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a76e8e18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of waveform: torch.Size([1, 1, 16000])\n",
      "Shape after 0th: layer: torch.Size([1, 512, 3199]), layer_instance: Conv1d(1, 512, kernel_size=(10,), stride=(5,))\n",
      "Shape after 1th: layer: torch.Size([1, 512, 3199]), layer_instance: GroupNorm(32, 512, eps=1e-05, affine=True)\n",
      "Shape after 2th: layer: torch.Size([1, 512, 3199]), layer_instance: GELU(approximate='none')\n",
      "Shape after 3th: layer: torch.Size([1, 512, 1599]), layer_instance: Conv1d(512, 512, kernel_size=(3,), stride=(2,))\n",
      "Shape after 4th: layer: torch.Size([1, 512, 1599]), layer_instance: GroupNorm(32, 512, eps=1e-05, affine=True)\n",
      "Shape after 5th: layer: torch.Size([1, 512, 1599]), layer_instance: GELU(approximate='none')\n",
      "Shape after 6th: layer: torch.Size([1, 512, 799]), layer_instance: Conv1d(512, 512, kernel_size=(3,), stride=(2,))\n",
      "Shape after 7th: layer: torch.Size([1, 512, 799]), layer_instance: GroupNorm(32, 512, eps=1e-05, affine=True)\n",
      "Shape after 8th: layer: torch.Size([1, 512, 799]), layer_instance: GELU(approximate='none')\n",
      "Shape after 9th: layer: torch.Size([1, 512, 399]), layer_instance: Conv1d(512, 512, kernel_size=(3,), stride=(2,))\n",
      "Shape after 10th: layer: torch.Size([1, 512, 399]), layer_instance: GroupNorm(32, 512, eps=1e-05, affine=True)\n",
      "Shape after 11th: layer: torch.Size([1, 512, 399]), layer_instance: GELU(approximate='none')\n",
      "Shape after 12th: layer: torch.Size([1, 512, 199]), layer_instance: Conv1d(512, 512, kernel_size=(3,), stride=(2,))\n",
      "Shape after 13th: layer: torch.Size([1, 512, 199]), layer_instance: GroupNorm(32, 512, eps=1e-05, affine=True)\n",
      "Shape after 14th: layer: torch.Size([1, 512, 199]), layer_instance: GELU(approximate='none')\n",
      "Shape after 15th: layer: torch.Size([1, 512, 99]), layer_instance: Conv1d(512, 512, kernel_size=(2,), stride=(2,))\n",
      "Shape after 16th: layer: torch.Size([1, 512, 99]), layer_instance: GroupNorm(32, 512, eps=1e-05, affine=True)\n",
      "Shape after 17th: layer: torch.Size([1, 512, 99]), layer_instance: GELU(approximate='none')\n",
      "Shape after 18th: layer: torch.Size([1, 512, 49]), layer_instance: Conv1d(512, 512, kernel_size=(2,), stride=(2,))\n",
      "Shape after 19th: layer: torch.Size([1, 512, 49]), layer_instance: GroupNorm(32, 512, eps=1e-05, affine=True)\n",
      "Shape after 20th: layer: torch.Size([1, 512, 49]), layer_instance: GELU(approximate='none')\n",
      "torch.Size([1, 49, 256])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "print(f\"Shape of waveform: {waveform.shape}\")\n",
    "\n",
    "filters = ModuleList([\n",
    "    Conv1d(in_channels=1, out_channels=512, kernel_size=10, stride=5),\n",
    "    GroupNorm(num_channels=512, num_groups=32),\n",
    "    GELU(),\n",
    "    Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),\n",
    "    GroupNorm(num_channels=512, num_groups=32),\n",
    "    GELU(),\n",
    "    Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),\n",
    "    GroupNorm(num_channels=512, num_groups=32),\n",
    "    GELU(),\n",
    "    Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),\n",
    "    GroupNorm(num_channels=512, num_groups=32),\n",
    "    GELU(),\n",
    "    Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2),\n",
    "    GroupNorm(num_channels=512, num_groups=32),\n",
    "    GELU(),\n",
    "    Conv1d(in_channels=512, out_channels=512, kernel_size=2, stride=2),\n",
    "    GroupNorm(num_channels=512, num_groups=32),\n",
    "    GELU(),\n",
    "    Conv1d(in_channels=512, out_channels=512, kernel_size=2, stride=2),\n",
    "    GroupNorm(num_channels=512, num_groups=32),\n",
    "    GELU(),\n",
    "])\n",
    "\n",
    "projector = Linear(in_features=512, out_features=256)\n",
    "\n",
    "filter_out = waveform\n",
    "for i, filter in enumerate(filters):\n",
    "    filter_out: Tensor = filter(filter_out)\n",
    "    print(f\"Shape after {i}th: layer: {filter_out.shape}, layer_instance: {filter}\")\n",
    "filter_out = filter_out.permute(0, 2, 1)\n",
    "filter_out = projector(filter_out)\n",
    "print(filter_out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "567df2f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "49\n"
     ]
    }
   ],
   "source": [
    "dummy_waveform = torch.zeros(1, 1, 16000)\n",
    "for i, filter in enumerate(filters):\n",
    "    # print(i)\n",
    "    dummy_waveform = filter(dummy_waveform)\n",
    "shink_factor = dummy_waveform.shape[2]\n",
    "print(shink_factor)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "study-asr (3.11.12)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
