audiolm_hubert_vs_mert.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "peitPTE5J-I9"
      },
      "outputs": [],
      "source": [
        "# pretrained substitute\n",
        "# Encodec as a replacement for SoundStream, and MERT as a replacement for w2v-BERT.\n",
        "# idea from https://github.com/zhvng/open-musiclm"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!rm -rf audiolm_pytorch/ setup.py audiolm-pytorch.zip # clean out any old stuff floating around\n",
        "\n",
        "!pip install torch datasets\n",
        "\n",
        "# download audiolm_pytorch manually so i can inject print statements\n",
        "# !pip uninstall -y audiolm_pytorch\n",
        "\n",
        "# raise AssertionError(\"don't forget to upload the customized version of audiolm_pytorch with print statements\")\n",
        "# !zip -r audiolm_pytorch.zip audiolm_pytorch/\n",
        "\n",
        "import urllib.request\n",
        "import os\n",
        "import zipfile\n",
        "if not os.path.isfile(\"audiolm-pytorch.zip\"):\n",
        "  urllib.request.urlretrieve(\"https://github.com/LWProgramming/audiolm-pytorch/archive/refs/heads/personal_hacks.zip\", \"audiolm-pytorch.zip\")\n",
        "if not os.path.isdir(\"audiolm-pytorch\"):\n",
        "  with zipfile.ZipFile(\"audiolm-pytorch.zip\", 'r') as zip_ref:\n",
        "    zip_ref.extractall(\"audiolm-pytorch\")\n",
        "!mv audiolm-pytorch/audiolm-pytorch-personal_hacks/audiolm_pytorch .\n",
        "\n",
        "# install necessary files for patched audiolm-pytorch\n",
        "!mv audiolm-pytorch/audiolm-pytorch-personal_hacks/setup.py .\n",
        "!pip install . # install requirements from the patched audiolm-pytorch dir\n",
        "!rm -rf audiolm-pytorch # not the one with underscore which is the actual library"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W4q_lA_1KdJv",
        "outputId": "5f0c4f71-80ef-4f0b-a951-013003a73211"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.8/dist-packages (1.13.1+cu116)\n",
            "Requirement already satisfied: datasets in /usr/local/lib/python3.8/dist-packages (2.10.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch) (4.5.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.8/dist-packages (from datasets) (3.8.4)\n",
            "Requirement already satisfied: pyarrow>=6.0.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (9.0.0)\n",
            "Requirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.8/dist-packages (from datasets) (4.64.1)\n",
            "Requirement already satisfied: dill<0.3.7,>=0.3.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (0.3.6)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.8/dist-packages (from datasets) (3.2.0)\n",
            "Requirement already satisfied: multiprocess in /usr/local/lib/python3.8/dist-packages (from datasets) (0.70.14)\n",
            "Requirement already satisfied: fsspec[http]>=2021.11.1 in /usr/local/lib/python3.8/dist-packages (from datasets) (2023.1.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.8/dist-packages (from datasets) (6.0)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.8/dist-packages (from datasets) (1.3.5)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from datasets) (23.0)\n",
            "Requirement already satisfied: responses<0.19 in /usr/local/lib/python3.8/dist-packages (from datasets) (0.18.0)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (2.25.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0.0,>=0.2.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (0.12.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from datasets) (1.22.4)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.8.2)\n",
            "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (3.0.1)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (4.0.2)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (22.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.3.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from huggingface-hub<1.0.0,>=0.2.0->datasets) (3.9.0)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (4.0.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (1.26.14)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (2022.12.7)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (2.10)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2022.7.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.7.3->pandas->datasets) (1.15.0)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Processing /content\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.16.0)\n",
            "Requirement already satisfied: beartype in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.12.0)\n",
            "Requirement already satisfied: einops>=0.6 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.6.0)\n",
            "Requirement already satisfied: ema-pytorch in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.2.1)\n",
            "Requirement already satisfied: fairseq in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.12.2)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (1.2.0)\n",
            "Requirement already satisfied: lion-pytorch in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.0.7)\n",
            "Requirement already satisfied: local-attention>=1.8.4 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (1.8.4)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (1.2.1)\n",
            "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.1.97)\n",
            "Requirement already satisfied: torch>=1.12 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (1.13.1+cu116)\n",
            "Requirement already satisfied: torchaudio in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (0.13.1+cu116)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (4.26.1)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (4.64.1)\n",
            "Requirement already satisfied: vector-quantize-pytorch>=1.0.6 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.2) (1.0.7)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch>=1.12->audiolm-pytorch==0.22.2) (4.5.0)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.2) (6.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.2) (1.22.4)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.2) (5.4.8)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.2) (23.0)\n",
            "Requirement already satisfied: hydra-core<1.1,>=1.0.7 in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.2) (1.0.7)\n",
            "Requirement already satisfied: cffi in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.2) (1.15.1)\n",
            "Requirement already satisfied: sacrebleu>=1.4.12 in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.2) (2.3.1)\n",
            "Requirement already satisfied: bitarray in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.2) (2.7.3)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.2) (2022.6.2)\n",
            "Requirement already satisfied: cython in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.2) (0.29.33)\n",
            "Requirement already satisfied: omegaconf<2.1 in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.2) (2.0.6)\n",
            "Requirement already satisfied: scipy>=1.3.2 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->audiolm-pytorch==0.22.2) (1.10.1)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->audiolm-pytorch==0.22.2) (3.1.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch==0.22.2) (2.25.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch==0.22.2) (3.9.0)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch==0.22.2) (0.12.1)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch==0.22.2) (0.13.2)\n",
            "Requirement already satisfied: antlr4-python3-runtime==4.8 in /usr/local/lib/python3.8/dist-packages (from hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch==0.22.2) (4.8)\n",
            "Requirement already satisfied: importlib-resources in /usr/local/lib/python3.8/dist-packages (from hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch==0.22.2) (5.12.0)\n",
            "Requirement already satisfied: colorama in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch==0.22.2) (0.4.6)\n",
            "Requirement already satisfied: portalocker in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch==0.22.2) (2.7.0)\n",
            "Requirement already satisfied: lxml in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch==0.22.2) (4.9.2)\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch==0.22.2) (0.8.10)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.8/dist-packages (from cffi->fairseq->audiolm-pytorch==0.22.2) (2.21)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch==0.22.2) (2022.12.7)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch==0.22.2) (4.0.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch==0.22.2) (1.26.14)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch==0.22.2) (2.10)\n",
            "Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.8/dist-packages (from importlib-resources->hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch==0.22.2) (3.15.0)\n",
            "Building wheels for collected packages: audiolm-pytorch\n",
            "  Building wheel for audiolm-pytorch (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for audiolm-pytorch: filename=audiolm_pytorch-0.22.2-py3-none-any.whl size=31138 sha256=61861e058ce54dcd8d516719ea35e67587c2907eac006fdce69d3a4609a43f47\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-r001hiu9/wheels/fd/3e/a8/f6cc18082b087ea481047dfbb1056db08877ac1382dcc607b7\n",
            "Successfully built audiolm-pytorch\n",
            "Installing collected packages: audiolm-pytorch\n",
            "  Attempting uninstall: audiolm-pytorch\n",
            "    Found existing installation: audiolm-pytorch 0.22.2\n",
            "    Uninstalling audiolm-pytorch-0.22.2:\n",
            "      Successfully uninstalled audiolm-pytorch-0.22.2\n",
            "Successfully installed audiolm-pytorch-0.22.2\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# # semantic- MERT\n",
        "# # https://huggingface.co/m-a-p/MERT-v0\n",
        "# # MERT-v0 is a completely unsupervised model trained on 1000 hour music audios.\n",
        "\n",
        "# from transformers import Wav2Vec2Processor, HubertModel\n",
        "# import torch\n",
        "# from torch import nn\n",
        "# from datasets import load_dataset\n",
        "\n",
        "# # load demo audio and set processor\n",
        "# dataset = load_dataset(\"hf-internal-testing/librispeech_asr_demo\", \"clean\", split=\"validation\")\n",
        "# dataset = dataset.sort(\"id\")\n",
        "# sampling_rate = dataset.features[\"audio\"].sampling_rate\n",
        "# processor = Wav2Vec2Processor.from_pretrained(\"facebook/hubert-large-ls960-ft\")\n",
        "\n",
        "# # loading our model weights\n",
        "# model = HubertModel.from_pretrained(\"m-a-p/MERT-v0\")\n",
        "\n",
        "# # audio file is decoded on the fly\n",
        "# inputs = processor(dataset[0][\"audio\"][\"array\"], sampling_rate=sampling_rate, return_tensors=\"pt\")\n",
        "# with torch.no_grad():\n",
        "#     outputs = model(**inputs, output_hidden_states=True)\n",
        "\n",
        "# # take a look at the output shape, there are 13 layers of representation\n",
        "# # each layer performs differently in different downstream tasks, you should choose empirically\n",
        "# all_layer_hidden_states = torch.stack(outputs.hidden_states).squeeze()\n",
        "# print(all_layer_hidden_states.shape) # [13 layer, 292 timestep, 768 feature_dim]\n",
        "\n",
        "# # # for utterance level classification tasks, you can simply reduce the representation in time\n",
        "# # time_reduced_hidden_states = all_layer_hidden_states.mean(-2)\n",
        "# # print(time_reduced_hidden_states.shape) # [13, 768]\n",
        "\n",
        "# # # you can even use a learnable weighted average representation\n",
        "# # aggregator = nn.Conv1d(in_channels=13, out_channels=1, kernel_size=1)\n",
        "# # weighted_avg_hidden_states = aggregator(time_reduced_hidden_states.unsqueeze(0)).squeeze()\n",
        "# # print(weighted_avg_hidden_states.shape) # [768]\n"
      ],
      "metadata": {
        "id": "Vo5-Yc1qKMEB"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# original semantic transformer\n",
        "import torch\n",
        "from audiolm_pytorch import HubertWithKmeans, SemanticTransformer, SemanticTransformerTrainer\n",
        "import os\n",
        "import urllib\n",
        "\n",
        "# hubert checkpoints can be downloaded at\n",
        "# https://github.com/facebookresearch/fairseq/tree/main/examples/hubert\n",
        "\n",
        "hubert_ckpt = 'hubert/hubert_base_ls960.pt'\n",
        "hubert_quantizer = f'hubert/hubert_base_ls960_L9_km500.bin' # listed in row \"HuBERT Base (~95M params)\", column Quantizer\n",
        "if not os.path.isdir(\"hubert\"):\n",
        "  os.makedirs(\"hubert\")\n",
        "if not os.path.isfile(hubert_ckpt):\n",
        "  hubert_ckpt_download = f\"https://dl.fbaipublicfiles.com/{hubert_ckpt}\"\n",
        "  urllib.request.urlretrieve(hubert_ckpt_download, f\"./{hubert_ckpt}\")\n",
        "if not os.path.isfile(hubert_quantizer):\n",
        "  hubert_quantizer_download = f\"https://dl.fbaipublicfiles.com/{hubert_quantizer}\"\n",
        "  urllib.request.urlretrieve(hubert_quantizer_download, f\"./{hubert_quantizer}\")\n",
        "\n",
        "wav2vec = HubertWithKmeans(\n",
        "    checkpoint_path = './hubert/hubert_base_ls960.pt',\n",
        "    kmeans_path = './hubert/hubert_base_ls960_L9_km500.bin'\n",
        ")\n",
        "\n",
        "semantic_transformer = SemanticTransformer(\n",
        "    num_semantic_tokens = wav2vec.codebook_size,\n",
        "    dim = 1024,\n",
        "    depth = 6\n",
        ").cuda()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "N_EHyuEwLVVk",
        "outputId": "0ceb8156-d651-4de8-b8d1-4d19a59c421b"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.8/dist-packages/sklearn/base.py:299: UserWarning: Trying to unpickle estimator MiniBatchKMeans from version 0.24.0 when using version 1.2.1. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
            "https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# import wave\n",
        "# import struct\n",
        "\n",
        "# # dataset = load_dataset(\"hf-internal-testing/librispeech_asr_demo\", \"clean\", split=\"validation\")\n",
        "# # dataset[0][\"audio\"][\"array\"]\n",
        "# sampling_rate = dataset.features[\"audio\"].sampling_rate\n",
        "\n",
        "# def save_wav(file_name, audio, sample_rate=sampling_rate):\n",
        "#   # Open up a wav file\n",
        "#   wav_file=wave.open(file_name,\"w\")\n",
        "#   # wav params\n",
        "#   nchannels = 1\n",
        "#   sampwidth = 2\n",
        "#   # 44100 is the industry standard sample rate - CD quality.  If you need to\n",
        "#   # save on file size you can adjust it downwards. The stanard for low quality\n",
        "#   # is 8000 or 8kHz.\n",
        "#   nframes = len(audio)\n",
        "#   comptype = \"NONE\"\n",
        "#   compname = \"not compressed\"\n",
        "#   wav_file.setparams((nchannels, sampwidth, sample_rate, nframes, comptype, compname))\n",
        "#   # WAV files here are using short, 16 bit, signed integers for the \n",
        "#   # sample size.  So we multiply the floating point data we have by 32767, the\n",
        "#   # maximum value for a short integer.  NOTE: It is theortically possible to\n",
        "#   # use the floating point -1.0 to 1.0 data directly in a WAV file but not\n",
        "#   # obvious how to do that using the wave module in python.\n",
        "#   for sample in audio:\n",
        "#     wav_file.writeframes(struct.pack('h', int( sample * 32767.0 )))\n",
        "#   wav_file.close()\n",
        "#   return\n",
        "# save_wav(\"test.wav\", dataset[1][\"audio\"][\"array\"])"
      ],
      "metadata": {
        "id": "jyQ8lrYiLwi7"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from audiolm_pytorch import SemanticTransformerWrapper\n",
        "import numpy as np\n",
        "\n",
        "# in case not already loaded\n",
        "from datasets import load_dataset\n",
        "# load demo audio and set processor\n",
        "dataset = load_dataset(\"hf-internal-testing/librispeech_asr_demo\", \"clean\", split=\"validation\")\n",
        "\n",
        "batch_size = 1\n",
        "# sample data[1] is 77040 samples at 16kHz sampling rate, 93680 for dataset[0]\n",
        "# just reshape it here so batch size for prime_wave is effectively 1\n",
        "samples = np.array([dataset[0][\"audio\"][\"array\"]])\n",
        "prime_wave = torch.tensor(samples).reshape(len(samples), len(samples[0])).cuda() # reshape from samples to batch x num_samples\n",
        "print(prime_wave.shape)\n",
        "# raise AssertionError(prime_wave.shape)\n",
        "max_length = 2048\n",
        "semantic = SemanticTransformerWrapper(\n",
        "            wav2vec = wav2vec,\n",
        "            transformer = semantic_transformer,\n",
        "            audio_conditioner = None,\n",
        "            unique_consecutive = True\n",
        "        ).cuda()\n",
        "semantic_tokens = semantic.generate(\n",
        "            text_embeds = None, # no text, it's not musicLM\n",
        "            batch_size = batch_size,\n",
        "            prime_wave = prime_wave,\n",
        "            max_length = max_length\n",
        "        )\n",
        "semantic_tokens.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "756nwPfnNLw1",
        "outputId": "f6044ec2-6837-483d-b094-3bdf45e150e7"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:datasets.builder:Found cached dataset librispeech_asr_demo (/root/.cache/huggingface/datasets/hf-internal-testing___librispeech_asr_demo/clean/2.1.0/d3bc4c2bc2078fcde3ad0f0f635862e4c0fef78ba94c4a34c4c250a097af240b)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([1, 93680])\n",
            "embed.keys(): dict_keys(['x', 'padding_mask', 'features'])\n",
            "embed['x'] shape: torch.Size([1, 292, 768]), embed['features'].shape: torch.Size([1, 292, 768])\n",
            "wav_input shape: torch.Size([1, 93680]), embed shape: torch.Size([292, 768]), packed_shape: [torch.Size([1, 292])]\n",
            "codebook_indices before unpacking: torch.Size([292])\n",
            "codebook_indices after unpacking: torch.Size([1, 292])\n",
            "ids.shape: torch.Size([1, 292]) and prime_wave True\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "generating semantic:  18%|█▊        | 345/1895 [00:06<00:30, 51.52it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "before masking eos, sample_semantic_ids.shape: torch.Size([1, 499])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 499])"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# ?semantic.wav2vec\n",
        "# torch.tensor(dataset[1][\"audio\"][\"array\"]).cuda().device\n",
        "# semantic_tokens.shape\n",
        "# semantic_tokens[:,-1] # exactly the EOS which is set as w2v.codebook_size == 500 in the case of k-means with k=500\n",
        "# semantic_tokens[:, 0]\n",
        "semantic_tokens[:, -1]\n",
        "# semantic.pad_id"
      ],
      "metadata": {
        "id": "Ov9QvzEfXbhl",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0ae31155-dfdb-461a-a3e2-aa31dc025abd"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([500], device='cuda:0')"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# dataset[1][\"audio\"][\"array\"].shape\n",
        "# # len(dataset[1][\"audio\"][\"array\"]) # 77040\n",
        "# # dataset.features[\"audio\"].sampling_rate # 16000\n",
        "# # so 4.815 seconds of audio"
      ],
      "metadata": {
        "id": "9eJXEn_lYpsR"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.empty((batch_size, 0), dtype = torch.long)"
      ],
      "metadata": {
        "id": "CzzhgEW-pekh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3aaeeccf-212f-418e-e1a5-a474ff67b453"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([], size=(1, 0), dtype=torch.int64)"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "semantic_transformer.eos_id"
      ],
      "metadata": {
        "id": "biq3EnSui7dd",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5670c6ca-afad-4855-a8a5-facfea252e83"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "500"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "TFwgrXS6TA9c"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}