{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72bf1b45-66fd-450d-8d5c-bec9e0b3d08f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from data2vec_feature_reader import Data2vecFeatureReader\n",
    "\n",
    "reader = Data2vecFeatureReader(\"./../../models/vox_pretrained.pt\", 18, device=\"cuda:0\", max_chunk=1600000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "84a9d238-048a-4772-a47b-5aadc50f36df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "490421d1c2f54cca9855f1a5397185f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading dataset shards:   0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be44942581b34d5388b0264e7b40d472",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading dataset shards:   0%|          | 0/60 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "\n",
    "cache_dir = \"./../../../cache\"\n",
    "\n",
    "dataset = load_dataset(\"openslr/librispeech_asr\", cache_dir=cache_dir, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cffd49ca-3524-4ac4-8ba5-bc4fcc9e0f53",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RepCodec(\n",
       "  (encoder): Encoder(\n",
       "    (conv): Conv1d(\n",
       "      (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
       "    )\n",
       "    (conv_blocks): ModuleList(\n",
       "      (0-1): 2 x EncoderBlock(\n",
       "        (res_units): ModuleList(\n",
       "          (0-1): 2 x ResidualUnit(\n",
       "            (activation): ELU(alpha=1.0)\n",
       "            (conv1): Conv1d(\n",
       "              (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
       "            )\n",
       "            (conv2): Conv1d1x1(1024, 1024, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "        )\n",
       "        (conv): Conv1d(\n",
       "          (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (conv1): Conv1d(\n",
       "      (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
       "    )\n",
       "    (conv_blocks): ModuleList(\n",
       "      (0-1): 2 x DecoderBlock(\n",
       "        (conv): Conv1d(\n",
       "          (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "        )\n",
       "        (res_units): ModuleList(\n",
       "          (0-1): 2 x ResidualUnit(\n",
       "            (activation): ELU(alpha=1.0)\n",
       "            (conv1): Conv1d(\n",
       "              (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
       "            )\n",
       "            (conv2): Conv1d1x1(1024, 1024, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (conv2): Conv1d(\n",
       "      (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
       "    )\n",
       "  )\n",
       "  (projector): Projector(\n",
       "    (project): Conv1d(\n",
       "      (conv): Conv1d(1024, 1024, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
       "    )\n",
       "  )\n",
       "  (quantizer): Quantizer(\n",
       "    (codebook): ResidualVQ(\n",
       "      (layers): ModuleList(\n",
       "        (0): VectorQuantize()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from repcodec.RepCodec import RepCodec\n",
    "import torch\n",
    "import yaml\n",
    "\n",
    "config = \"./../repcodec/configs/repcodec_dim1024.yaml\"\n",
    "with open(config) as fp:\n",
    "    conf = yaml.load(fp, Loader=yaml.FullLoader)\n",
    "\n",
    "model = RepCodec(**conf)\n",
    "model.load_state_dict(torch.load(\"./../../models/data2vec_large_l18.pkl\", map_location=\"cuda:0\")[\"model\"][\"repcodec\"])\n",
    "model.quantizer.initial()\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a9a1731e-052c-4af0-a29c-b171a988b300",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[22], line 27\u001b[0m\n\u001b[1;32m     23\u001b[0m     feat\u001b[38;5;241m.\u001b[39mappend(feat_chunk)\n\u001b[1;32m     25\u001b[0m features \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mcat(feat, \u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39mpermute(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m---> 27\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoder\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfeatures\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfloat32\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     28\u001b[0m z \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mprojector(x)\n\u001b[1;32m     29\u001b[0m _, idx \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mquantizer\u001b[38;5;241m.\u001b[39mcodebook\u001b[38;5;241m.\u001b[39mforward_index(z\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m1\u001b[39m))\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\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/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\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\u001b[1;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/jupyter_workspace/users/Darshan/RepCodec/repcodec/modules/encoder.py:86\u001b[0m, in \u001b[0;36mEncoder.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     85\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x):\n\u001b[0;32m---> 86\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     87\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_blocks):\n\u001b[1;32m     88\u001b[0m         x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv_blocks[i](x)\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\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/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\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\u001b[1;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/jupyter_workspace/users/Darshan/RepCodec/repcodec/layers/conv_layer.py:55\u001b[0m, in \u001b[0;36mConv1d.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x):\n\u001b[1;32m     49\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     50\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[1;32m     51\u001b[0m \u001b[38;5;124;03m        x (Tensor): Float tensor variable with the shape  (B, C, T).\u001b[39;00m\n\u001b[1;32m     52\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     53\u001b[0m \u001b[38;5;124;03m        Tensor: Float tensor variable with the shape (B, C, T).\u001b[39;00m\n\u001b[1;32m     54\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 55\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     56\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m x\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\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/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\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\u001b[1;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/conv.py:310\u001b[0m, in \u001b[0;36mConv1d.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    309\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m--> 310\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_conv_forward\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/conv.py:306\u001b[0m, in \u001b[0;36mConv1d._conv_forward\u001b[0;34m(self, input, weight, bias)\u001b[0m\n\u001b[1;32m    302\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpadding_mode \u001b[38;5;241m!=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mzeros\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[1;32m    303\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m F\u001b[38;5;241m.\u001b[39mconv1d(F\u001b[38;5;241m.\u001b[39mpad(\u001b[38;5;28minput\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reversed_padding_repeated_twice, mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpadding_mode),\n\u001b[1;32m    304\u001b[0m                     weight, bias, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstride,\n\u001b[1;32m    305\u001b[0m                     _single(\u001b[38;5;241m0\u001b[39m), \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdilation, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgroups)\n\u001b[0;32m--> 306\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv1d\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstride\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    307\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpadding\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdilation\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgroups\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same"
     ]
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "sample = dataset[\"train.clean.100\"][1]\n",
    "\n",
    "x = sample[\"audio\"][\"array\"]\n",
    "\n",
    "with torch.no_grad():\n",
    "    x = torch.from_numpy(x).float().to(reader.device)\n",
    "    if reader.task.cfg.normalize:\n",
    "        x = F.layer_norm(x, x.shape)\n",
    "    x = x.view(1, -1)\n",
    "\n",
    "    feat = []\n",
    "    for start in range(0, x.size(1), reader.max_chunk):\n",
    "        x_chunk = x[:, start: start + reader.max_chunk]\n",
    "        res = reader.model.extract_features(\n",
    "            source=x_chunk,\n",
    "            padding_mask=None,\n",
    "            mask=False,\n",
    "            layer=reader.layer,\n",
    "        )\n",
    "        feat_chunk = res[\"x\"]\n",
    "        feat.append(feat_chunk)\n",
    "        \n",
    "    features = torch.cat(feat, 1).permute(0, 2, 1)\n",
    "\n",
    "    x = model.encoder(features)\n",
    "    z = model.projector(x)\n",
    "    _, idx = model.quantizer.codebook.forward_index(z.transpose(2, 1))\n",
    "    tokens = idx.cpu().data.numpy().tolist()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d51709a9-6fb3-450b-a517-005367095663",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 804, 1024])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dfc977d7-f27c-40d7-b545-fbdf26728cbe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([726, 1024])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1810e6dc-2ece-4aca-a29a-e1933b8ce82a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import tqdm\n",
    "from npy_append_array import NpyAppendArray\n",
    "\n",
    "def get_shard_range(tot, nshard, rank):\n",
    "    assert rank < nshard and rank >= 0, f\"invaid rank/nshard {rank}/{nshard}\"\n",
    "    start = round(tot / nshard * rank)\n",
    "    end = round(tot / nshard * (rank + 1))\n",
    "    assert start < end, f\"start={start}, end={end}\"\n",
    "    logger.info(\n",
    "        f\"rank {rank} of {nshard}, process {end-start} \"\n",
    "        f\"({start}-{end}) out of {tot}\"\n",
    "    )\n",
    "    return start, end\n",
    "\n",
    "def get_path_iterator(tsv, nshard, rank):\n",
    "    with open(tsv, \"r\") as f:\n",
    "        root = f.readline().rstrip()\n",
    "        lines = [line.rstrip() for line in f]\n",
    "        start, end = get_shard_range(len(lines), nshard, rank)\n",
    "        lines = lines[start:end]\n",
    "        def iterate():\n",
    "            for line in lines:\n",
    "                subpath, nsample = line.split(\"\\t\")\n",
    "                yield f\"{root}/{subpath}\", int(nsample)\n",
    "    return iterate, len(lines)\n",
    "\n",
    "def dump_feature(reader, generator, num, nshard, rank, feat_dir):\n",
    "    iterator = generator()\n",
    "\n",
    "    feat_path = f\"{feat_dir}/{rank}_{nshard}.npy\"\n",
    "    leng_path = f\"{feat_dir}/{rank}_{nshard}.len\"\n",
    "\n",
    "    os.makedirs(feat_dir, exist_ok=True)\n",
    "    if os.path.exists(feat_path):\n",
    "        os.remove(feat_path)\n",
    "\n",
    "    feat_f = NpyAppendArray(feat_path)\n",
    "    with open(leng_path, \"w\") as leng_f:\n",
    "        for path, nsample in tqdm.tqdm(iterator, total=num):\n",
    "            feat = reader.get_feats(path, nsample)\n",
    "            feat_f.append(feat.cpu().numpy())\n",
    "            leng_f.write(f\"{len(feat)}\\n\")\n",
    "    logger.info(\"finished successfully\")\n",
    "\n",
    "generator, num = get_path_iterator(tsv_path, nshard, rank)\n",
    "dump_feature(reader, generator, num, nshard, rank, feat_dir)"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}