coarse_fine_explore.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JNhddkiJiW8p",
        "outputId": "6542ba18-966b-4bb9-c455-4a21b18c4966"
      },
      "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",
            "Collecting datasets\n",
            "  Downloading datasets-2.10.1-py3-none-any.whl (469 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m469.0/469.0 KB\u001b[0m \u001b[31m9.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: boto3 in /usr/local/lib/python3.8/dist-packages (1.26.84)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch) (4.5.0)\n",
            "Collecting responses<0.19\n",
            "  Downloading responses-0.18.0-py3-none-any.whl (38 kB)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.8/dist-packages (from datasets) (3.8.4)\n",
            "Collecting xxhash\n",
            "  Downloading xxhash-3.2.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (213 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m213.0/213.0 KB\u001b[0m \u001b[31m24.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: pyarrow>=6.0.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (9.0.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",
            "Collecting dill<0.3.7,>=0.3.0\n",
            "  Downloading dill-0.3.6-py3-none-any.whl (110 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m110.5/110.5 KB\u001b[0m \u001b[31m15.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.8/dist-packages (from datasets) (2.25.1)\n",
            "Requirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.8/dist-packages (from datasets) (4.64.1)\n",
            "Collecting multiprocess\n",
            "  Downloading multiprocess-0.70.14-py38-none-any.whl (132 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m132.0/132.0 KB\u001b[0m \u001b[31m17.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from datasets) (23.0)\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: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from datasets) (1.22.4)\n",
            "Collecting huggingface-hub<1.0.0,>=0.2.0\n",
            "  Downloading huggingface_hub-0.12.1-py3-none-any.whl (190 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.3/190.3 KB\u001b[0m \u001b[31m23.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: botocore<1.30.0,>=1.29.84 in /usr/local/lib/python3.8/dist-packages (from boto3) (1.29.84)\n",
            "Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /usr/local/lib/python3.8/dist-packages (from boto3) (1.0.1)\n",
            "Requirement already satisfied: s3transfer<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from boto3) (0.6.0)\n",
            "Requirement already satisfied: python-dateutil<3.0.0,>=2.1 in /usr/local/lib/python3.8/dist-packages (from botocore<1.30.0,>=1.29.84->boto3) (2.8.2)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.25.4 in /usr/local/lib/python3.8/dist-packages (from botocore<1.30.0,>=1.29.84->boto3) (1.26.14)\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: frozenlist>=1.1.1 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.3.3)\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: 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: aiosignal>=1.1.2 in /usr/local/lib/python3.8/dist-packages (from aiohttp->datasets) (1.3.1)\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: 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: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->datasets) (2.10)\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: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas->datasets) (2022.7.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil<3.0.0,>=2.1->botocore<1.30.0,>=1.29.84->boto3) (1.15.0)\n",
            "Installing collected packages: xxhash, dill, responses, multiprocess, huggingface-hub, datasets\n",
            "Successfully installed datasets-2.10.1 dill-0.3.6 huggingface-hub-0.12.1 multiprocess-0.70.14 responses-0.18.0 xxhash-3.2.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",
            "Collecting accelerate\n",
            "  Downloading accelerate-0.16.0-py3-none-any.whl (199 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m199.7/199.7 KB\u001b[0m \u001b[31m5.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting beartype\n",
            "  Downloading beartype-0.12.0-py3-none-any.whl (754 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m754.5/754.5 KB\u001b[0m \u001b[31m24.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting einops>=0.6\n",
            "  Downloading einops-0.6.0-py3-none-any.whl (41 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.6/41.6 KB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting ema-pytorch\n",
            "  Downloading ema_pytorch-0.2.1-py3-none-any.whl (4.4 kB)\n",
            "Collecting fairseq\n",
            "  Downloading fairseq-0.12.2-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (11.0 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.0/11.0 MB\u001b[0m \u001b[31m95.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: joblib in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.3) (1.2.0)\n",
            "Collecting lion-pytorch\n",
            "  Downloading lion_pytorch-0.0.7-py3-none-any.whl (4.3 kB)\n",
            "Collecting local-attention>=1.8.4\n",
            "  Downloading local_attention-1.8.4-py3-none-any.whl (8.0 kB)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.3) (1.2.1)\n",
            "Collecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.97-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m67.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: torch>=1.12 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.3) (1.13.1+cu116)\n",
            "Requirement already satisfied: torchaudio in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.3) (0.13.1+cu116)\n",
            "Collecting transformers\n",
            "  Downloading transformers-4.26.1-py3-none-any.whl (6.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.3/6.3 MB\u001b[0m \u001b[31m102.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch==0.22.3) (4.64.1)\n",
            "Collecting vector-quantize-pytorch>=1.0.6\n",
            "  Downloading vector_quantize_pytorch-1.0.7-py3-none-any.whl (9.0 kB)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch>=1.12->audiolm-pytorch==0.22.3) (4.5.0)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.3) (23.0)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.3) (6.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.3) (1.22.4)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch==0.22.3) (5.4.8)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.3) (2022.6.2)\n",
            "Collecting sacrebleu>=1.4.12\n",
            "  Downloading sacrebleu-2.3.1-py3-none-any.whl (118 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m118.9/118.9 KB\u001b[0m \u001b[31m16.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting hydra-core<1.1,>=1.0.7\n",
            "  Downloading hydra_core-1.0.7-py3-none-any.whl (123 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m123.8/123.8 KB\u001b[0m \u001b[31m17.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: cffi in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.3) (1.15.1)\n",
            "Requirement already satisfied: cython in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch==0.22.3) (0.29.33)\n",
            "Collecting bitarray\n",
            "  Downloading bitarray-2.7.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (271 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m271.8/271.8 KB\u001b[0m \u001b[31m32.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting omegaconf<2.1\n",
            "  Downloading omegaconf-2.0.6-py3-none-any.whl (36 kB)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->audiolm-pytorch==0.22.3) (3.1.0)\n",
            "Requirement already satisfied: scipy>=1.3.2 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->audiolm-pytorch==0.22.3) (1.10.1)\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.3) (0.12.1)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch==0.22.3) (2.25.1)\n",
            "Collecting tokenizers!=0.11.3,<0.14,>=0.11.1\n",
            "  Downloading tokenizers-0.13.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (7.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.6/7.6 MB\u001b[0m \u001b[31m31.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch==0.22.3) (3.9.0)\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.3) (5.12.0)\n",
            "Collecting antlr4-python3-runtime==4.8\n",
            "  Downloading antlr4-python3-runtime-4.8.tar.gz (112 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m112.4/112.4 KB\u001b[0m \u001b[31m15.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: lxml in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch==0.22.3) (4.9.2)\n",
            "Collecting portalocker\n",
            "  Downloading portalocker-2.7.0-py2.py3-none-any.whl (15 kB)\n",
            "Collecting colorama\n",
            "  Downloading colorama-0.4.6-py2.py3-none-any.whl (25 kB)\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.3) (0.8.10)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.8/dist-packages (from cffi->fairseq->audiolm-pytorch==0.22.3) (2.21)\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.3) (1.26.14)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch==0.22.3) (4.0.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch==0.22.3) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch==0.22.3) (2022.12.7)\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.3) (3.15.0)\n",
            "Building wheels for collected packages: audiolm-pytorch, antlr4-python3-runtime\n",
            "  Building wheel for audiolm-pytorch (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for audiolm-pytorch: filename=audiolm_pytorch-0.22.3-py3-none-any.whl size=31177 sha256=49fc60738d50202cfa8059e0a8247979148178ac98a7685b0ca051c3321b9114\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-scgvqrr3/wheels/fd/3e/a8/f6cc18082b087ea481047dfbb1056db08877ac1382dcc607b7\n",
            "  Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.8-py3-none-any.whl size=141231 sha256=221b163cf13de901dc145900fa3d72e87309f4f90664d093ca3680ab45d82b3a\n",
            "  Stored in directory: /root/.cache/pip/wheels/c8/d0/ab/d43c02eaddc5b9004db86950802442ad9a26f279c619e28da0\n",
            "Successfully built audiolm-pytorch antlr4-python3-runtime\n",
            "Installing collected packages: tokenizers, sentencepiece, bitarray, antlr4-python3-runtime, portalocker, omegaconf, einops, colorama, beartype, vector-quantize-pytorch, sacrebleu, local-attention, lion-pytorch, hydra-core, ema-pytorch, accelerate, transformers, fairseq, audiolm-pytorch\n",
            "Successfully installed accelerate-0.16.0 antlr4-python3-runtime-4.8 audiolm-pytorch-0.22.3 beartype-0.12.0 bitarray-2.7.3 colorama-0.4.6 einops-0.6.0 ema-pytorch-0.2.1 fairseq-0.12.2 hydra-core-1.0.7 lion-pytorch-0.0.7 local-attention-1.8.4 omegaconf-2.0.6 portalocker-2.7.0 sacrebleu-2.3.1 sentencepiece-0.1.97 tokenizers-0.13.2 transformers-4.26.1 vector-quantize-pytorch-1.0.7\n"
          ]
        }
      ],
      "source": [
        "#exploring the coarse and fine transformers to see if i understand it right\n",
        "!rm -rf audiolm_pytorch/ setup.py audiolm-pytorch.zip # clean out any old stuff floating around\n",
        "\n",
        "!pip install torch datasets boto3\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",
        "\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\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')\n",
        "\n",
        "%set_env AWS_SHARED_CREDENTIALS_FILE=drive/MyDrive/Colab Notebooks/AWS-SECRET-credentials\n",
        "%set_env AWS_CONFIG_FILE=drive/MyDrive/Colab Notebooks/AWS-config"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZuNrP98dv4ik",
        "outputId": "2432761c-9112-458a-c0a6-d6a98ece6653"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive/\n",
            "env: AWS_SHARED_CREDENTIALS_FILE=drive/MyDrive/Colab Notebooks/AWS-SECRET-credentials\n",
            "env: AWS_CONFIG_FILE=drive/MyDrive/Colab Notebooks/AWS-config\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# # load soundstream checkpoint from version 0.11.1-- never mind, doesn't work\n",
        "# import boto3\n",
        "\n",
        "# boto3.setup_default_session(profile_name='mltraining')\n",
        "# s3_resource = boto3.resource('s3')\n",
        "# bucket_name = \"itsleonwu-paperspace\"\n",
        "# bucket = s3_resource.Bucket(name=bucket_name)\n",
        "\n",
        "# # for obj in bucket.objects.all():\n",
        "# #     print(obj.key)\n",
        "# soundstream_ckpt = \"actual_soundstream_results/20000.pt\" # name of file in bucket\n",
        "# bucket.download_file(soundstream_ckpt, \"soundstream_ckpt.pt\")"
      ],
      "metadata": {
        "id": "v9lgu_Rqwwub"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # 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",
        "# !rm -rf placeholder_dataset\n",
        "\n",
        "import math\n",
        "import wave\n",
        "import struct\n",
        "# Placeholder data generation\n",
        "def get_sinewave(freq=440.0, duration_ms=200, volume=1.0, sample_rate=44100.0):\n",
        "  # code adapted from https://stackoverflow.com/a/33913403\n",
        "  audio = []\n",
        "  num_samples = duration_ms * (sample_rate / 1000.0)\n",
        "  for x in range(int(num_samples)):\n",
        "    audio.append(volume * math.sin(2 * math.pi * freq * (x / sample_rate)))\n",
        "  return audio\n",
        "\n",
        "def save_wav(file_name, audio, sample_rate=44100.0):\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",
        "\n",
        "def make_placeholder_dataset(dataset_folder=\"placeholder_dataset\"):\n",
        "  # Make a placeholder dataset with a few .wav files that you can \"train\" on, just to verify things work e2e\n",
        "  if os.path.isdir(dataset_folder):\n",
        "    return\n",
        "  os.makedirs(dataset_folder)\n",
        "  save_wav(f\"{dataset_folder}/example.wav\", get_sinewave(6000))\n",
        "  save_wav(f\"{dataset_folder}/example2.wav\", get_sinewave(duration_ms=5000))\n",
        "  os.makedirs(f\"{dataset_folder}/subdirectory\")\n",
        "  save_wav(f\"{dataset_folder}/subdirectory/example.wav\", get_sinewave(freq=330.0))\n",
        "\n",
        "make_placeholder_dataset()"
      ],
      "metadata": {
        "id": "Ea6GWZMo55n4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# original semantic transformer\n",
        "import torch\n",
        "from audiolm_pytorch import HubertWithKmeans, SemanticTransformer, SemanticTransformerWrapper, 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()\n",
        "\n",
        "semantic = SemanticTransformerWrapper(\n",
        "            wav2vec = wav2vec,\n",
        "            transformer = semantic_transformer,\n",
        "            audio_conditioner = None,\n",
        "            unique_consecutive = True\n",
        "        ).cuda()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 92
        },
        "id": "quTnRsftihZo",
        "outputId": "c9dbe32a-ec74-4c87-c8c3-22878fc1cd24"
      },
      "execution_count": null,
      "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 torch\n",
        "from audiolm_pytorch import HubertWithKmeans, AudioLMSoundStream, CoarseTransformer, CoarseTransformerWrapper, CoarseTransformerTrainer\n",
        "\n",
        "soundstream = AudioLMSoundStream(\n",
        "    # rq_num_quantizers=8\n",
        ")\n",
        "# soundstream.load('soundstream_ckpt.pt')\n",
        "\n",
        "coarse_transformer = CoarseTransformer(\n",
        "    num_semantic_tokens = wav2vec.codebook_size,\n",
        "    codebook_size = 1024,\n",
        "    num_coarse_quantizers = 3,\n",
        "    dim = 512,\n",
        "    depth = 6\n",
        ")\n",
        "\n",
        "# coarse = CoarseTransformerWrapper(transformer=coarse_transformer, soundstream=soundstream, wav2vec=wav2vec)\n",
        "\n",
        "trainer = CoarseTransformerTrainer(\n",
        "    transformer = coarse_transformer,\n",
        "    soundstream = soundstream,\n",
        "    wav2vec = wav2vec,\n",
        "    folder = 'placeholder_dataset',\n",
        "    batch_size = 1,\n",
        "    data_max_length = 320 * 32,\n",
        "    num_train_steps = 3\n",
        ")\n",
        "\n",
        "trainer.train()\n",
        "\n",
        "# TODO: continue from here.\n",
        "# - why doesn't generate end up calling forward via forward_with_cond_scale?\n",
        "# should then hit part where we get sem and coarse token id shapes"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "grSTTam0ikEW",
        "outputId": "ccd2a9de-1b76-4a15-c2f3-0c4279ce1c1b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "training with dataset of 2 samples and validating with randomly splitted 1 samples\n",
            "do you want to clear previous experiment checkpoints and results? (y/n) y\n",
            "embed.keys(): dict_keys(['x', 'padding_mask', 'features'])\n",
            "embed['x'] shape: torch.Size([1, 31, 768]), embed['features'].shape: torch.Size([1, 31, 768])\n",
            "wav_input shape: torch.Size([1, 10240]), embed shape: torch.Size([31, 768]), packed_shape: [torch.Size([1, 31])]\n",
            "codebook_indices before unpacking: torch.Size([31])\n",
            "codebook_indices after unpacking: torch.Size([1, 31])\n",
            "before rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 32, 3])\n",
            "after rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 96])\n",
            "0: loss: 56.759849548339844\n",
            "embed.keys(): dict_keys(['x', 'padding_mask', 'features'])\n",
            "embed['x'] shape: torch.Size([1, 31, 768]), embed['features'].shape: torch.Size([1, 31, 768])\n",
            "wav_input shape: torch.Size([1, 10240]), embed shape: torch.Size([31, 768]), packed_shape: [torch.Size([1, 31])]\n",
            "codebook_indices before unpacking: torch.Size([31])\n",
            "codebook_indices after unpacking: torch.Size([1, 31])\n",
            "before rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 32, 3])\n",
            "after rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 96])\n",
            "0: valid loss 15.861306190490723\n",
            "0: saving model to results\n",
            "embed.keys(): dict_keys(['x', 'padding_mask', 'features'])\n",
            "embed['x'] shape: torch.Size([1, 31, 768]), embed['features'].shape: torch.Size([1, 31, 768])\n",
            "wav_input shape: torch.Size([1, 10240]), embed shape: torch.Size([31, 768]), packed_shape: [torch.Size([1, 31])]\n",
            "codebook_indices before unpacking: torch.Size([31])\n",
            "codebook_indices after unpacking: torch.Size([1, 31])\n",
            "before rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 32, 3])\n",
            "after rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 96])\n",
            "1: loss: 16.717803955078125\n",
            "embed.keys(): dict_keys(['x', 'padding_mask', 'features'])\n",
            "embed['x'] shape: torch.Size([1, 31, 768]), embed['features'].shape: torch.Size([1, 31, 768])\n",
            "wav_input shape: torch.Size([1, 10240]), embed shape: torch.Size([31, 768]), packed_shape: [torch.Size([1, 31])]\n",
            "codebook_indices before unpacking: torch.Size([31])\n",
            "codebook_indices after unpacking: torch.Size([1, 31])\n",
            "before rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 32, 3])\n",
            "after rearranging: semantic_token_ids.shape torch.Size([1, 31]), coarse_token_ids.shape torch.Size([1, 96])\n",
            "2: loss: 39.311344146728516\n",
            "training complete\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\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",
        "\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",
        "print(f\"semantic_tokens.shape {semantic_tokens.shape}\")\n",
        "\n",
        "coarse_tokens = coarse.generate(\n",
        "    text_embeds = None,\n",
        "    semantic_token_ids = semantic_tokens,\n",
        "    reconstruct_wave = False # just want the tokens\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "KLvADzsTyKY0",
        "outputId": "dde2bf60-73dd-487f-d005-d680eb829cc4"
      },
      "execution_count": null,
      "outputs": [
        {
          "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": [
            "\n",
            "generating semantic:   0%|          | 0/1895 [00:00<?, ?it/s]\u001b[A\n",
            "generating semantic:   0%|          | 8/1895 [00:00<00:26, 71.71it/s]\u001b[A\n",
            "generating semantic:   1%|          | 16/1895 [00:00<00:25, 73.37it/s]\u001b[A\n",
            "generating semantic:   1%|▏         | 24/1895 [00:00<00:24, 75.52it/s]\u001b[A\n",
            "generating semantic:   2%|▏         | 32/1895 [00:00<00:25, 74.08it/s]\u001b[A\n",
            "generating semantic:   2%|▏         | 40/1895 [00:00<00:27, 68.60it/s]\u001b[A\n",
            "generating semantic:   2%|▏         | 47/1895 [00:00<00:28, 65.72it/s]\u001b[A\n",
            "generating semantic:   3%|▎         | 54/1895 [00:00<00:28, 65.30it/s]\u001b[A\n",
            "generating semantic:   3%|▎         | 61/1895 [00:00<00:29, 62.76it/s]\u001b[A\n",
            "generating semantic:   4%|▎         | 68/1895 [00:01<00:28, 63.23it/s]\u001b[A\n",
            "generating semantic:   4%|▍         | 75/1895 [00:01<00:28, 63.74it/s]\u001b[A\n",
            "generating semantic:   4%|▍         | 82/1895 [00:01<00:28, 63.99it/s]\u001b[A\n",
            "generating semantic:   5%|▍         | 89/1895 [00:01<00:28, 63.70it/s]\u001b[A\n",
            "generating semantic:   5%|▌         | 96/1895 [00:01<00:28, 62.45it/s]\u001b[A\n",
            "generating semantic:   5%|▌         | 103/1895 [00:01<00:28, 62.08it/s]\u001b[A\n",
            "generating semantic:   6%|▌         | 110/1895 [00:01<00:28, 62.02it/s]\u001b[A\n",
            "generating semantic:   6%|▌         | 117/1895 [00:01<00:28, 61.84it/s]\u001b[A\n",
            "generating semantic:   7%|▋         | 124/1895 [00:01<00:28, 61.24it/s]\u001b[A\n",
            "generating semantic:   7%|▋         | 131/1895 [00:02<00:28, 61.36it/s]\u001b[A\n",
            "generating semantic:   7%|▋         | 138/1895 [00:02<00:29, 60.27it/s]\u001b[A\n",
            "generating semantic:   8%|▊         | 145/1895 [00:02<00:29, 58.85it/s]\u001b[A\n",
            "generating semantic:   8%|▊         | 151/1895 [00:02<00:30, 58.06it/s]\u001b[A\n",
            "generating semantic:   8%|▊         | 157/1895 [00:02<00:30, 57.35it/s]\u001b[A\n",
            "generating semantic:   9%|▊         | 163/1895 [00:02<00:30, 56.71it/s]\u001b[A\n",
            "generating semantic:   9%|▉         | 169/1895 [00:02<00:30, 55.98it/s]\u001b[A\n",
            "generating semantic:   9%|▉         | 175/1895 [00:02<00:31, 53.82it/s]\u001b[A\n",
            "generating semantic:  10%|▉         | 181/1895 [00:02<00:32, 52.53it/s]\u001b[A\n",
            "generating semantic:  10%|▉         | 187/1895 [00:03<00:33, 51.33it/s]\u001b[A\n",
            "generating semantic:  10%|█         | 193/1895 [00:03<00:33, 50.22it/s]\u001b[A\n",
            "generating semantic:  11%|█         | 199/1895 [00:03<00:34, 49.39it/s]\u001b[A\n",
            "generating semantic:  11%|█         | 204/1895 [00:03<00:34, 49.18it/s]\u001b[A\n",
            "generating semantic:  11%|█         | 209/1895 [00:03<00:34, 49.04it/s]\u001b[A\n",
            "generating semantic:  11%|█▏        | 214/1895 [00:03<00:35, 47.50it/s]\u001b[A\n",
            "generating semantic:  12%|█▏        | 219/1895 [00:03<00:35, 47.31it/s]\u001b[A\n",
            "generating semantic:  12%|█▏        | 224/1895 [00:03<00:35, 47.53it/s]\u001b[A\n",
            "generating semantic:  12%|█▏        | 229/1895 [00:03<00:35, 47.20it/s]\u001b[A\n",
            "generating semantic:  12%|█▏        | 234/1895 [00:04<00:36, 45.95it/s]\u001b[A\n",
            "generating semantic:  13%|█▎        | 239/1895 [00:04<00:36, 44.94it/s]\u001b[A\n",
            "generating semantic:  13%|█▎        | 244/1895 [00:04<00:37, 43.77it/s]\u001b[A\n",
            "generating semantic:  13%|█▎        | 249/1895 [00:04<00:38, 42.79it/s]\u001b[A\n",
            "generating semantic:  13%|█▎        | 254/1895 [00:04<00:39, 42.06it/s]\u001b[A\n",
            "generating semantic:  14%|█▎        | 259/1895 [00:04<00:39, 41.32it/s]\u001b[A\n",
            "generating semantic:  14%|█▍        | 264/1895 [00:04<00:39, 41.21it/s]\u001b[A\n",
            "generating semantic:  14%|█▍        | 269/1895 [00:04<00:39, 41.26it/s]\u001b[A\n",
            "generating semantic:  14%|█▍        | 274/1895 [00:05<00:39, 40.90it/s]\u001b[A\n",
            "generating semantic:  15%|█▍        | 279/1895 [00:05<00:39, 40.49it/s]\u001b[A\n",
            "generating semantic:  15%|█▍        | 284/1895 [00:05<00:40, 40.22it/s]\u001b[A\n",
            "generating semantic:  15%|█▌        | 289/1895 [00:05<00:40, 40.08it/s]\u001b[A\n",
            "generating semantic:  16%|█▌        | 294/1895 [00:05<00:40, 39.56it/s]\u001b[A\n",
            "generating semantic:  16%|█▌        | 298/1895 [00:05<00:41, 38.78it/s]\u001b[A\n",
            "generating semantic:  16%|█▌        | 302/1895 [00:05<00:41, 38.73it/s]\u001b[A\n",
            "generating semantic:  16%|█▌        | 306/1895 [00:05<00:41, 38.65it/s]\u001b[A\n",
            "generating semantic:  16%|█▋        | 310/1895 [00:05<00:41, 38.62it/s]\u001b[A\n",
            "generating semantic:  17%|█▋        | 314/1895 [00:06<00:41, 38.22it/s]\u001b[A\n",
            "generating semantic:  17%|█▋        | 318/1895 [00:06<00:41, 37.88it/s]\u001b[A\n",
            "generating semantic:  17%|█▋        | 322/1895 [00:06<00:41, 37.84it/s]\u001b[A\n",
            "generating semantic:  17%|█▋        | 326/1895 [00:06<00:41, 37.78it/s]\u001b[A\n",
            "generating semantic:  17%|█▋        | 330/1895 [00:06<00:42, 37.14it/s]\u001b[A\n",
            "generating semantic:  18%|█▊        | 334/1895 [00:06<00:42, 36.95it/s]\u001b[A\n",
            "generating semantic:  18%|█▊        | 338/1895 [00:06<00:42, 36.81it/s]\u001b[A\n",
            "generating semantic:  18%|█▊        | 342/1895 [00:06<00:42, 36.64it/s]\u001b[A\n",
            "generating semantic:  18%|█▊        | 346/1895 [00:06<00:42, 36.45it/s]\u001b[A\n",
            "generating semantic:  18%|█▊        | 350/1895 [00:07<00:42, 36.37it/s]\u001b[A\n",
            "generating semantic:  19%|█▊        | 354/1895 [00:07<00:42, 36.41it/s]\u001b[A\n",
            "generating semantic:  19%|█▉        | 358/1895 [00:07<00:42, 36.31it/s]\u001b[A\n",
            "generating semantic:  19%|█▉        | 362/1895 [00:07<00:43, 35.42it/s]\u001b[A\n",
            "generating semantic:  19%|█▉        | 366/1895 [00:07<00:44, 34.54it/s]\u001b[A\n",
            "generating semantic:  20%|█▉        | 370/1895 [00:07<00:44, 34.02it/s]\u001b[A\n",
            "generating semantic:  20%|█▉        | 374/1895 [00:07<00:45, 33.41it/s]\u001b[A\n",
            "generating semantic:  20%|█▉        | 378/1895 [00:07<00:46, 32.90it/s]\u001b[A\n",
            "generating semantic:  20%|██        | 382/1895 [00:08<00:46, 32.67it/s]\u001b[A\n",
            "generating semantic:  20%|██        | 386/1895 [00:08<00:46, 32.50it/s]\u001b[A\n",
            "generating semantic:  21%|██        | 390/1895 [00:08<00:46, 32.34it/s]\u001b[A\n",
            "generating semantic:  21%|██        | 394/1895 [00:08<00:46, 32.16it/s]\u001b[A\n",
            "generating semantic:  21%|██        | 398/1895 [00:08<00:46, 32.09it/s]\u001b[A\n",
            "generating semantic:  21%|██        | 402/1895 [00:08<00:47, 31.51it/s]\u001b[A\n",
            "generating semantic:  21%|██▏       | 406/1895 [00:08<00:47, 31.44it/s]\u001b[A\n",
            "generating semantic:  22%|██▏       | 410/1895 [00:08<00:47, 31.37it/s]\u001b[A\n",
            "generating semantic:  22%|██▏       | 414/1895 [00:09<00:47, 31.08it/s]\u001b[A\n",
            "generating semantic:  22%|██▏       | 418/1895 [00:09<00:48, 30.60it/s]\u001b[A\n",
            "generating semantic:  22%|██▏       | 422/1895 [00:09<00:48, 30.60it/s]\u001b[A\n",
            "generating semantic:  22%|██▏       | 426/1895 [00:09<00:47, 30.80it/s]\u001b[A\n",
            "generating semantic:  23%|██▎       | 430/1895 [00:09<00:48, 30.50it/s]\u001b[A\n",
            "generating semantic:  23%|██▎       | 434/1895 [00:09<00:48, 30.31it/s]\u001b[A\n",
            "generating semantic:  23%|██▎       | 438/1895 [00:09<00:48, 30.28it/s]\u001b[A\n",
            "generating semantic:  23%|██▎       | 442/1895 [00:09<00:48, 30.16it/s]\u001b[A\n",
            "generating semantic:  24%|██▎       | 446/1895 [00:10<00:48, 30.09it/s]\u001b[A\n",
            "generating semantic:  24%|██▎       | 450/1895 [00:10<00:48, 29.97it/s]\u001b[A\n",
            "generating semantic:  24%|██▍       | 453/1895 [00:10<00:48, 29.75it/s]\u001b[A\n",
            "generating semantic:  24%|██▍       | 456/1895 [00:10<00:48, 29.53it/s]\u001b[A\n",
            "generating semantic:  24%|██▍       | 459/1895 [00:10<00:49, 29.21it/s]\u001b[A\n",
            "generating semantic:  24%|██▍       | 462/1895 [00:10<00:49, 29.12it/s]\u001b[A\n",
            "generating semantic:  25%|██▍       | 465/1895 [00:10<00:49, 29.18it/s]\u001b[A\n",
            "generating semantic:  25%|██▍       | 468/1895 [00:10<00:48, 29.16it/s]\u001b[A\n",
            "generating semantic:  25%|██▍       | 471/1895 [00:10<00:49, 28.97it/s]\u001b[A\n",
            "generating semantic:  25%|██▌       | 474/1895 [00:11<00:49, 28.69it/s]\u001b[A\n",
            "generating semantic:  25%|██▌       | 477/1895 [00:11<00:49, 28.64it/s]\u001b[A\n",
            "generating semantic:  25%|██▌       | 480/1895 [00:11<00:49, 28.49it/s]\u001b[A\n",
            "generating semantic:  25%|██▌       | 483/1895 [00:11<00:49, 28.53it/s]\u001b[A\n",
            "generating semantic:  26%|██▌       | 486/1895 [00:11<00:49, 28.55it/s]\u001b[A\n",
            "generating semantic:  26%|██▌       | 489/1895 [00:11<00:50, 28.01it/s]\u001b[A\n",
            "generating semantic:  26%|██▌       | 492/1895 [00:11<00:51, 27.17it/s]\u001b[A\n",
            "generating semantic:  26%|██▌       | 495/1895 [00:11<00:52, 26.57it/s]\u001b[A\n",
            "generating semantic:  26%|██▋       | 498/1895 [00:11<00:53, 25.91it/s]\u001b[A\n",
            "generating semantic:  26%|██▋       | 501/1895 [00:12<00:54, 25.56it/s]\u001b[A\n",
            "generating semantic:  27%|██▋       | 504/1895 [00:12<00:54, 25.41it/s]\u001b[A\n",
            "generating semantic:  27%|██▋       | 507/1895 [00:12<00:55, 25.17it/s]\u001b[A\n",
            "generating semantic:  27%|██▋       | 510/1895 [00:12<00:55, 24.91it/s]\u001b[A\n",
            "generating semantic:  27%|██▋       | 513/1895 [00:12<00:55, 24.93it/s]\u001b[A\n",
            "generating semantic:  27%|██▋       | 516/1895 [00:12<00:55, 24.77it/s]\u001b[A\n",
            "generating semantic:  27%|██▋       | 519/1895 [00:12<00:55, 24.69it/s]\u001b[A\n",
            "generating semantic:  28%|██▊       | 522/1895 [00:12<00:56, 24.43it/s]\u001b[A\n",
            "generating semantic:  28%|██▊       | 525/1895 [00:13<00:55, 24.49it/s]\u001b[A\n",
            "generating semantic:  28%|██▊       | 528/1895 [00:13<00:55, 24.47it/s]\u001b[A\n",
            "generating semantic:  28%|██▊       | 531/1895 [00:13<00:55, 24.49it/s]\u001b[A\n",
            "generating semantic:  28%|██▊       | 534/1895 [00:13<00:55, 24.54it/s]\u001b[A\n",
            "generating semantic:  28%|██▊       | 537/1895 [00:13<00:55, 24.61it/s]\u001b[A\n",
            "generating semantic:  28%|██▊       | 540/1895 [00:13<00:55, 24.57it/s]\u001b[A\n",
            "generating semantic:  29%|██▊       | 543/1895 [00:13<00:54, 24.60it/s]\u001b[A\n",
            "generating semantic:  29%|██▉       | 546/1895 [00:13<00:55, 24.48it/s]\u001b[A\n",
            "generating semantic:  29%|██▉       | 549/1895 [00:14<00:55, 24.43it/s]\u001b[A\n",
            "generating semantic:  29%|██▉       | 552/1895 [00:14<00:55, 24.41it/s]\u001b[A\n",
            "generating semantic:  29%|██▉       | 555/1895 [00:14<00:55, 24.19it/s]\u001b[A\n",
            "generating semantic:  29%|██▉       | 558/1895 [00:14<00:55, 23.91it/s]\u001b[A\n",
            "generating semantic:  30%|██▉       | 561/1895 [00:14<00:56, 23.66it/s]\u001b[A\n",
            "generating semantic:  30%|██▉       | 564/1895 [00:14<00:56, 23.49it/s]\u001b[A\n",
            "generating semantic:  30%|██▉       | 567/1895 [00:14<00:56, 23.44it/s]\u001b[A\n",
            "generating semantic:  30%|███       | 570/1895 [00:14<00:57, 23.24it/s]\u001b[A\n",
            "generating semantic:  30%|███       | 573/1895 [00:15<00:57, 23.16it/s]\u001b[A\n",
            "generating semantic:  30%|███       | 576/1895 [00:15<00:56, 23.25it/s]\u001b[A\n",
            "generating semantic:  31%|███       | 579/1895 [00:15<00:57, 22.93it/s]\u001b[A\n",
            "generating semantic:  31%|███       | 582/1895 [00:15<00:57, 22.86it/s]\u001b[A\n",
            "generating semantic:  31%|███       | 585/1895 [00:15<00:57, 22.81it/s]\u001b[A\n",
            "generating semantic:  31%|███       | 588/1895 [00:15<00:56, 23.02it/s]\u001b[A\n",
            "generating semantic:  31%|███       | 591/1895 [00:15<00:57, 22.76it/s]\u001b[A\n",
            "generating semantic:  31%|███▏      | 594/1895 [00:16<00:57, 22.77it/s]\u001b[A\n",
            "generating semantic:  32%|███▏      | 597/1895 [00:16<00:57, 22.68it/s]\u001b[A\n",
            "generating semantic:  32%|███▏      | 600/1895 [00:16<00:57, 22.66it/s]\u001b[A\n",
            "generating semantic:  32%|███▏      | 603/1895 [00:16<00:57, 22.62it/s]\u001b[A\n",
            "generating semantic:  32%|███▏      | 606/1895 [00:16<00:57, 22.49it/s]\u001b[A\n",
            "generating semantic:  32%|███▏      | 609/1895 [00:16<00:57, 22.50it/s]\u001b[A\n",
            "generating semantic:  32%|███▏      | 612/1895 [00:16<00:57, 22.31it/s]\u001b[A\n",
            "generating semantic:  32%|███▏      | 615/1895 [00:16<00:57, 22.23it/s]\u001b[A\n",
            "generating semantic:  33%|███▎      | 618/1895 [00:17<00:58, 21.80it/s]\u001b[A\n",
            "generating semantic:  33%|███▎      | 621/1895 [00:17<00:59, 21.47it/s]\u001b[A\n",
            "generating semantic:  33%|███▎      | 624/1895 [00:17<00:59, 21.19it/s]\u001b[A\n",
            "generating semantic:  33%|███▎      | 627/1895 [00:17<01:00, 20.84it/s]\u001b[A\n",
            "generating semantic:  33%|███▎      | 630/1895 [00:17<01:00, 20.82it/s]\u001b[A\n",
            "generating semantic:  33%|███▎      | 633/1895 [00:17<01:00, 20.72it/s]\u001b[A\n",
            "generating semantic:  34%|███▎      | 636/1895 [00:17<01:01, 20.64it/s]\u001b[A\n",
            "generating semantic:  34%|███▎      | 639/1895 [00:18<01:01, 20.43it/s]\u001b[A\n",
            "generating semantic:  34%|███▍      | 642/1895 [00:18<01:01, 20.35it/s]\u001b[A\n",
            "generating semantic:  34%|███▍      | 645/1895 [00:18<01:01, 20.21it/s]\u001b[A\n",
            "generating semantic:  34%|███▍      | 648/1895 [00:18<01:01, 20.21it/s]\u001b[A\n",
            "generating semantic:  34%|███▍      | 651/1895 [00:18<01:02, 19.94it/s]\u001b[A\n",
            "generating semantic:  34%|███▍      | 653/1895 [00:18<01:02, 19.94it/s]\u001b[A\n",
            "generating semantic:  35%|███▍      | 655/1895 [00:18<01:02, 19.89it/s]\u001b[A\n",
            "generating semantic:  35%|███▍      | 657/1895 [00:19<01:02, 19.82it/s]\u001b[A\n",
            "generating semantic:  35%|███▍      | 659/1895 [00:19<01:02, 19.81it/s]\u001b[A\n",
            "generating semantic:  35%|███▍      | 661/1895 [00:19<01:02, 19.73it/s]\u001b[A\n",
            "generating semantic:  35%|███▍      | 663/1895 [00:19<01:02, 19.70it/s]\u001b[A\n",
            "generating semantic:  35%|███▌      | 665/1895 [00:19<01:02, 19.64it/s]\u001b[A\n",
            "generating semantic:  35%|███▌      | 667/1895 [00:19<01:02, 19.56it/s]\u001b[A\n",
            "generating semantic:  35%|███▌      | 669/1895 [00:19<01:02, 19.61it/s]\u001b[A\n",
            "generating semantic:  35%|███▌      | 671/1895 [00:19<01:02, 19.58it/s]\u001b[A\n",
            "generating semantic:  36%|███▌      | 673/1895 [00:19<01:02, 19.61it/s]\u001b[A\n",
            "generating semantic:  36%|███▌      | 675/1895 [00:19<01:02, 19.63it/s]\u001b[A\n",
            "generating semantic:  36%|███▌      | 677/1895 [00:20<01:02, 19.62it/s]\u001b[A\n",
            "generating semantic:  36%|███▌      | 679/1895 [00:20<01:02, 19.57it/s]\u001b[A\n",
            "generating semantic:  36%|███▌      | 681/1895 [00:20<01:02, 19.50it/s]\u001b[A\n",
            "generating semantic:  36%|███▌      | 683/1895 [00:20<01:02, 19.44it/s]\u001b[A\n",
            "generating semantic:  36%|███▌      | 685/1895 [00:20<01:02, 19.36it/s]\u001b[A\n",
            "generating semantic:  36%|███▋      | 687/1895 [00:20<01:02, 19.33it/s]\u001b[A\n",
            "generating semantic:  36%|███▋      | 689/1895 [00:20<01:02, 19.16it/s]\u001b[A\n",
            "generating semantic:  36%|███▋      | 691/1895 [00:20<01:03, 19.10it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 693/1895 [00:20<01:02, 19.12it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 695/1895 [00:20<01:02, 19.15it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 697/1895 [00:21<01:02, 19.18it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 699/1895 [00:21<01:02, 19.13it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 701/1895 [00:21<01:02, 19.01it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 703/1895 [00:21<01:02, 18.98it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 705/1895 [00:21<01:03, 18.81it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 707/1895 [00:21<01:03, 18.77it/s]\u001b[A\n",
            "generating semantic:  37%|███▋      | 709/1895 [00:21<01:03, 18.71it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 711/1895 [00:21<01:03, 18.74it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 713/1895 [00:21<01:02, 18.80it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 715/1895 [00:22<01:02, 18.78it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 717/1895 [00:22<01:02, 18.77it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 719/1895 [00:22<01:03, 18.62it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 721/1895 [00:22<01:03, 18.54it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 723/1895 [00:22<01:03, 18.52it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 725/1895 [00:22<01:03, 18.29it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 727/1895 [00:22<01:03, 18.46it/s]\u001b[A\n",
            "generating semantic:  38%|███▊      | 729/1895 [00:22<01:03, 18.41it/s]\u001b[A\n",
            "generating semantic:  39%|███▊      | 731/1895 [00:22<01:03, 18.36it/s]\u001b[A\n",
            "generating semantic:  39%|███▊      | 733/1895 [00:23<01:02, 18.48it/s]\u001b[A\n",
            "generating semantic:  39%|███▉      | 735/1895 [00:23<01:02, 18.42it/s]\u001b[A\n",
            "generating semantic:  39%|███▉      | 737/1895 [00:23<01:03, 18.35it/s]\u001b[A\n",
            "generating semantic:  39%|███▉      | 739/1895 [00:23<01:02, 18.38it/s]\u001b[A\n",
            "generating semantic:  39%|███▉      | 741/1895 [00:23<01:03, 18.28it/s]\u001b[A\n",
            "generating semantic:  39%|███▉      | 743/1895 [00:23<01:03, 18.15it/s]\u001b[A\n",
            "generating semantic:  39%|███▉      | 745/1895 [00:23<01:04, 17.76it/s]\u001b[A\n",
            "generating semantic:  39%|███▉      | 747/1895 [00:23<01:05, 17.53it/s]\u001b[A\n",
            "generating semantic:  40%|███▉      | 749/1895 [00:23<01:06, 17.25it/s]\u001b[A\n",
            "generating semantic:  40%|███▉      | 751/1895 [00:24<01:06, 17.28it/s]\u001b[A\n",
            "generating semantic:  40%|███▉      | 753/1895 [00:24<01:06, 17.22it/s]\u001b[A\n",
            "generating semantic:  40%|███▉      | 755/1895 [00:24<01:06, 17.18it/s]\u001b[A\n",
            "generating semantic:  40%|███▉      | 757/1895 [00:24<01:06, 17.01it/s]\u001b[A\n",
            "generating semantic:  40%|████      | 759/1895 [00:24<01:07, 16.88it/s]\u001b[A\n",
            "generating semantic:  40%|████      | 761/1895 [00:24<01:07, 16.82it/s]\u001b[A\n",
            "generating semantic:  40%|████      | 763/1895 [00:24<01:07, 16.72it/s]\u001b[A\n",
            "generating semantic:  40%|████      | 765/1895 [00:24<01:07, 16.67it/s]\u001b[A\n",
            "generating semantic:  40%|████      | 767/1895 [00:25<01:07, 16.67it/s]\u001b[A\n",
            "generating semantic:  41%|████      | 769/1895 [00:25<01:07, 16.80it/s]\u001b[A\n",
            "generating semantic:  41%|████      | 771/1895 [00:25<01:07, 16.77it/s]\u001b[A\n",
            "generating semantic:  41%|████      | 773/1895 [00:25<00:36, 30.39it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "before masking eos, sample_semantic_ids.shape: torch.Size([1, 927])\n",
            "semantic_tokens.shape torch.Size([1, 927])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "generating coarse:   0%|          | 0/512 [00:00<?, ?it/s]\u001b[A\n",
            "generating coarse:   0%|          | 1/512 [00:03<28:04,  3.30s/it]\u001b[A\n",
            "generating coarse:   0%|          | 2/512 [00:06<27:28,  3.23s/it]\u001b[A\n",
            "generating coarse:   1%|          | 3/512 [00:10<30:04,  3.55s/it]\u001b[A\n",
            "generating coarse:   1%|          | 4/512 [00:14<32:56,  3.89s/it]\u001b[A\n",
            "generating coarse:   1%|          | 5/512 [00:17<30:26,  3.60s/it]\u001b[A\n",
            "generating coarse:   1%|          | 6/512 [00:21<28:57,  3.43s/it]\u001b[A\n",
            "generating coarse:   1%|▏         | 7/512 [00:24<28:04,  3.33s/it]\u001b[A\n",
            "generating coarse:   2%|▏         | 8/512 [00:28<31:21,  3.73s/it]\u001b[A\n",
            "generating coarse:   2%|▏         | 9/512 [00:32<31:41,  3.78s/it]\u001b[A\n",
            "generating coarse:   2%|▏         | 10/512 [00:35<30:14,  3.61s/it]\u001b[A\n",
            "generating coarse:   2%|▏         | 11/512 [00:39<29:09,  3.49s/it]\u001b[A\n",
            "generating coarse:   2%|▏         | 12/512 [00:43<30:25,  3.65s/it]\u001b[A\n",
            "generating coarse:   3%|▎         | 13/512 [00:47<32:40,  3.93s/it]\u001b[A\n",
            "generating coarse:   3%|▎         | 14/512 [00:50<30:58,  3.73s/it]\u001b[A\n",
            "generating coarse:   3%|▎         | 15/512 [00:54<29:45,  3.59s/it]\u001b[A\n",
            "generating coarse:   3%|▎         | 16/512 [00:57<29:25,  3.56s/it]\u001b[A\n",
            "generating coarse:   3%|▎         | 17/512 [01:02<32:58,  4.00s/it]\u001b[A\n",
            "generating coarse:   4%|▎         | 18/512 [01:06<32:10,  3.91s/it]\u001b[A\n",
            "generating coarse:   4%|▎         | 19/512 [01:09<30:55,  3.76s/it]\u001b[A\n",
            "generating coarse:   4%|▍         | 20/512 [01:13<30:01,  3.66s/it]\u001b[A\n",
            "generating coarse:   4%|▍         | 21/512 [01:18<33:15,  4.06s/it]\u001b[A\n",
            "generating coarse:   4%|▍         | 22/512 [01:22<32:47,  4.02s/it]\u001b[A\n",
            "generating coarse:   4%|▍         | 23/512 [01:25<31:15,  3.84s/it]\u001b[A\n",
            "generating coarse:   5%|▍         | 24/512 [01:29<30:11,  3.71s/it]\u001b[A\n",
            "generating coarse:   5%|▍         | 25/512 [01:33<32:46,  4.04s/it]\u001b[A\n",
            "generating coarse:   5%|▌         | 26/512 [01:38<33:10,  4.10s/it]\u001b[A\n",
            "generating coarse:   5%|▌         | 27/512 [01:41<31:43,  3.92s/it]\u001b[A\n",
            "generating coarse:   5%|▌         | 28/512 [01:45<30:37,  3.80s/it]\u001b[A\n",
            "generating coarse:   6%|▌         | 29/512 [01:49<33:14,  4.13s/it]\u001b[A\n",
            "generating coarse:   6%|▌         | 30/512 [01:54<33:34,  4.18s/it]\u001b[A\n",
            "generating coarse:   6%|▌         | 31/512 [01:57<32:02,  4.00s/it]\u001b[A\n",
            "generating coarse:   6%|▋         | 32/512 [02:01<30:59,  3.87s/it]\u001b[A\n",
            "generating coarse:   6%|▋         | 33/512 [02:06<33:29,  4.20s/it]\u001b[A\n",
            "generating coarse:   7%|▋         | 34/512 [02:10<33:57,  4.26s/it]\u001b[A\n",
            "generating coarse:   7%|▋         | 35/512 [02:14<32:31,  4.09s/it]\u001b[A\n",
            "generating coarse:   7%|▋         | 36/512 [02:18<31:27,  3.96s/it]\u001b[A\n",
            "generating coarse:   7%|▋         | 37/512 [02:23<34:19,  4.34s/it]\u001b[A\n",
            "generating coarse:   7%|▋         | 38/512 [02:27<33:57,  4.30s/it]\u001b[A\n",
            "generating coarse:   8%|▊         | 39/512 [02:31<32:27,  4.12s/it]\u001b[A\n",
            "generating coarse:   8%|▊         | 40/512 [02:35<31:32,  4.01s/it]\u001b[A\n",
            "generating coarse:   8%|▊         | 41/512 [02:40<35:14,  4.49s/it]\u001b[A\n",
            "generating coarse:   8%|▊         | 42/512 [02:44<33:58,  4.34s/it]\u001b[A\n",
            "generating coarse:   8%|▊         | 43/512 [02:48<32:41,  4.18s/it]\u001b[A\n",
            "generating coarse:   9%|▊         | 44/512 [02:52<33:08,  4.25s/it]\u001b[A\n",
            "generating coarse:   9%|▉         | 45/512 [02:58<35:45,  4.59s/it]\u001b[A\n",
            "generating coarse:   9%|▉         | 46/512 [03:02<33:59,  4.38s/it]\u001b[A\n",
            "generating coarse:   9%|▉         | 47/512 [03:06<32:52,  4.24s/it]\u001b[A\n",
            "generating coarse:   9%|▉         | 48/512 [03:11<34:56,  4.52s/it]\u001b[A\n",
            "generating coarse:  10%|▉         | 49/512 [03:16<35:33,  4.61s/it]\u001b[A\n",
            "generating coarse:  10%|▉         | 50/512 [03:19<33:52,  4.40s/it]\u001b[A\n",
            "generating coarse:  10%|▉         | 51/512 [03:23<32:46,  4.26s/it]\u001b[A\n",
            "generating coarse:  10%|█         | 52/512 [03:29<36:15,  4.73s/it]\u001b[A\n",
            "generating coarse:  10%|█         | 53/512 [03:33<34:54,  4.56s/it]\u001b[A\n",
            "generating coarse:  11%|█         | 54/512 [03:37<33:35,  4.40s/it]\u001b[A\n",
            "generating coarse:  11%|█         | 55/512 [03:42<34:32,  4.53s/it]\u001b[A\n",
            "generating coarse:  11%|█         | 56/512 [03:47<36:04,  4.75s/it]\u001b[A\n",
            "generating coarse:  11%|█         | 57/512 [03:52<34:26,  4.54s/it]\u001b[A\n",
            "generating coarse:  11%|█▏        | 58/512 [03:56<33:21,  4.41s/it]\u001b[A\n",
            "generating coarse:  12%|█▏        | 59/512 [04:01<36:32,  4.84s/it]\u001b[A\n",
            "generating coarse:  12%|█▏        | 60/512 [04:06<35:34,  4.72s/it]\u001b[A\n",
            "generating coarse:  12%|█▏        | 61/512 [04:10<34:19,  4.57s/it]\u001b[A\n",
            "generating coarse:  12%|█▏        | 62/512 [04:16<36:05,  4.81s/it]\u001b[A\n",
            "generating coarse:  12%|█▏        | 63/512 [04:21<36:49,  4.92s/it]\u001b[A\n",
            "generating coarse:  12%|█▎        | 64/512 [04:25<35:08,  4.71s/it]\u001b[A\n",
            "generating coarse:  13%|█▎        | 65/512 [04:29<34:27,  4.63s/it]\u001b[A\n",
            "generating coarse:  13%|█▎        | 66/512 [04:35<37:35,  5.06s/it]\u001b[A\n",
            "generating coarse:  13%|█▎        | 67/512 [04:40<36:01,  4.86s/it]\u001b[A\n",
            "generating coarse:  13%|█▎        | 68/512 [04:44<34:46,  4.70s/it]\u001b[A\n",
            "generating coarse:  13%|█▎        | 69/512 [04:50<37:41,  5.10s/it]\u001b[A\n",
            "generating coarse:  14%|█▎        | 70/512 [04:55<36:37,  4.97s/it]\u001b[A\n",
            "generating coarse:  14%|█▍        | 71/512 [04:59<35:06,  4.78s/it]\u001b[A\n",
            "generating coarse:  14%|█▍        | 72/512 [05:05<36:55,  5.04s/it]\u001b[A\n",
            "generating coarse:  14%|█▍        | 73/512 [05:10<37:21,  5.11s/it]\u001b[A\n",
            "generating coarse:  14%|█▍        | 74/512 [05:15<35:51,  4.91s/it]\u001b[A\n",
            "generating coarse:  15%|█▍        | 75/512 [05:20<36:17,  4.98s/it]\u001b[A\n",
            "generating coarse:  15%|█▍        | 76/512 [05:25<37:52,  5.21s/it]\u001b[A\n",
            "generating coarse:  15%|█▌        | 77/512 [05:30<36:21,  5.01s/it]\u001b[A\n",
            "generating coarse:  15%|█▌        | 78/512 [05:35<35:55,  4.97s/it]\u001b[A\n",
            "generating coarse:  15%|█▌        | 79/512 [05:41<38:31,  5.34s/it]\u001b[A\n",
            "generating coarse:  16%|█▌        | 80/512 [05:46<36:50,  5.12s/it]\u001b[A\n",
            "generating coarse:  16%|█▌        | 81/512 [05:50<35:47,  4.98s/it]\u001b[A\n",
            "generating coarse:  16%|█▌        | 82/512 [05:57<39:00,  5.44s/it]\u001b[A\n",
            "generating coarse:  16%|█▌        | 83/512 [06:01<37:12,  5.20s/it]\u001b[A\n",
            "generating coarse:  16%|█▋        | 84/512 [06:06<36:04,  5.06s/it]\u001b[A\n",
            "generating coarse:  17%|█▋        | 85/512 [06:13<39:26,  5.54s/it]\u001b[A\n",
            "generating coarse:  17%|█▋        | 86/512 [06:18<37:49,  5.33s/it]\u001b[A\n",
            "generating coarse:  17%|█▋        | 87/512 [06:22<36:23,  5.14s/it]\u001b[A\n",
            "generating coarse:  17%|█▋        | 88/512 [06:29<39:39,  5.61s/it]\u001b[A\n",
            "generating coarse:  17%|█▋        | 89/512 [06:34<37:42,  5.35s/it]\u001b[A\n",
            "generating coarse:  18%|█▊        | 90/512 [06:39<36:30,  5.19s/it]\u001b[A\n",
            "generating coarse:  18%|█▊        | 91/512 [06:45<39:49,  5.68s/it]\u001b[A\n",
            "generating coarse:  18%|█▊        | 92/512 [06:50<37:56,  5.42s/it]\u001b[A\n",
            "generating coarse:  18%|█▊        | 93/512 [06:55<36:39,  5.25s/it]\u001b[A\n",
            "generating coarse:  18%|█▊        | 94/512 [07:02<39:46,  5.71s/it]\u001b[A\n",
            "generating coarse:  19%|█▊        | 95/512 [07:07<38:08,  5.49s/it]\u001b[A\n",
            "generating coarse:  19%|█▉        | 96/512 [07:12<37:27,  5.40s/it]\u001b[A\n",
            "generating coarse:  19%|█▉        | 97/512 [07:19<40:06,  5.80s/it]\u001b[A\n",
            "generating coarse:  19%|█▉        | 98/512 [07:24<38:20,  5.56s/it]\u001b[A\n",
            "generating coarse:  19%|█▉        | 99/512 [07:30<38:36,  5.61s/it]\u001b[A\n",
            "generating coarse:  20%|█▉        | 100/512 [07:36<39:43,  5.79s/it]\u001b[A\n",
            "generating coarse:  20%|█▉        | 101/512 [07:41<38:02,  5.55s/it]\u001b[A\n",
            "generating coarse:  20%|█▉        | 102/512 [07:47<39:04,  5.72s/it]\u001b[A\n",
            "generating coarse:  20%|██        | 103/512 [07:53<39:29,  5.79s/it]\u001b[A\n",
            "generating coarse:  20%|██        | 104/512 [07:58<37:58,  5.58s/it]\u001b[A\n",
            "generating coarse:  21%|██        | 105/512 [08:05<40:05,  5.91s/it]\u001b[A\n",
            "generating coarse:  21%|██        | 106/512 [08:10<39:20,  5.81s/it]\u001b[A\n",
            "generating coarse:  21%|██        | 107/512 [08:15<37:55,  5.62s/it]\u001b[A\n",
            "generating coarse:  21%|██        | 108/512 [08:22<40:47,  6.06s/it]\u001b[A\n",
            "generating coarse:  21%|██▏       | 109/512 [08:28<38:48,  5.78s/it]\u001b[A\n",
            "generating coarse:  21%|██▏       | 110/512 [08:33<38:05,  5.69s/it]\u001b[A\n",
            "generating coarse:  22%|██▏       | 111/512 [08:42<31:28,  4.71s/it]\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-25-cc833545ab81>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"semantic_tokens.shape {semantic_tokens.shape}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 20\u001b[0;31m coarse_tokens = coarse.generate(\n\u001b[0m\u001b[1;32m     21\u001b[0m     \u001b[0mtext_embeds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m     \u001b[0msemantic_token_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msemantic_tokens\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/content/audiolm_pytorch/audiolm_pytorch.py\u001b[0m in \u001b[0;36minner\u001b[0;34m(model, *args, **kwargs)\u001b[0m\n\u001b[1;32m     63\u001b[0m         \u001b[0mwas_training\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m         \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 65\u001b[0;31m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     66\u001b[0m         \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwas_training\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     67\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<@beartype(audiolm_pytorch.audiolm_pytorch.CoarseTransformerWrapper.generate) at 0x7f3bd4986c10>\u001b[0m in \u001b[0;36mgenerate\u001b[0;34m(__beartype_func, __beartype_conf, __beartype_get_violation, __beartype_object_9474080, __beartype_getrandbits, *args, **kwargs)\u001b[0m\n",
            "\u001b[0;32m/content/audiolm_pytorch/audiolm_pytorch.py\u001b[0m in \u001b[0;36mgenerate\u001b[0;34m(self, semantic_token_ids, text, text_embeds, max_time_steps, cond_scale, filter_thres, temperature, reconstruct_wave, **kwargs)\u001b[0m\n\u001b[1;32m   1279\u001b[0m                 \u001b[0mis_last_step\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_coarse_quantizers\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1280\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1281\u001b[0;31m                 _, coarse_logits = self.transformer.forward_with_cond_scale(\n\u001b[0m\u001b[1;32m   1282\u001b[0m                     \u001b[0mcoarse_token_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msampled_coarse_token_ids\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1283\u001b[0m                     \u001b[0msemantic_token_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msemantic_token_ids\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/content/audiolm_pytorch/audiolm_pytorch.py\u001b[0m in \u001b[0;36mforward_with_cond_scale\u001b[0;34m(self, cond_scale, *args, **kwargs)\u001b[0m\n\u001b[1;32m    652\u001b[0m         \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    653\u001b[0m     ):\n\u001b[0;32m--> 654\u001b[0;31m         \u001b[0msemantic_logits\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcoarse_logits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcond_drop_prob\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0.\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    655\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    656\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mcond_scale\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhas_condition\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<@beartype(audiolm_pytorch.audiolm_pytorch.CoarseTransformer.forward) at 0x7f3bd4986160>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(__beartype_func, __beartype_conf, __beartype_get_violation, __beartype_object_9474080, __beartype_getrandbits, *args, **kwargs)\u001b[0m\n",
            "\u001b[0;32m/content/audiolm_pytorch/audiolm_pytorch.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, semantic_token_ids, coarse_token_ids, self_attn_mask, text, text_embeds, cond_drop_prob, return_only_coarse_logits)\u001b[0m\n\u001b[1;32m    724\u001b[0m         ), dim = 1)\n\u001b[1;32m    725\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 726\u001b[0;31m         \u001b[0mtokens\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransformer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtokens\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcontext\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtext_embeds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself_attn_mask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself_attn_mask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcontext_mask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtext_mask\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    727\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    728\u001b[0m         \u001b[0mpred_semantic_tokens\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpred_coarse_tokens\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtokens\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0msemantic_seq_len\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtokens\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0msemantic_seq_len\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/content/audiolm_pytorch/audiolm_pytorch.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, self_attn_mask, context, context_mask, attn_bias)\u001b[0m\n\u001b[1;32m    443\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    444\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mattn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcross_attn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mff\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayers\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 445\u001b[0;31m             \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mattn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattn_bias\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrel_pos_bias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself_attn_mask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mself_attn_kwargs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    446\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    447\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcross_attn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/content/audiolm_pytorch/audiolm_pytorch.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, context, mask, attn_bias, prefix_context, prefix_context_mask)\u001b[0m\n\u001b[1;32m    369\u001b[0m         \u001b[0;31m# aggregate\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    370\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 371\u001b[0;31m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meinsum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'b h i j, b j d -> b h i d'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    372\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    373\u001b[0m         \u001b[0;31m# merge heads\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/functional.py\u001b[0m in \u001b[0;36meinsum\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    376\u001b[0m         \u001b[0;31m# the path for contracting 0 or 1 time(s) is already optimized\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    377\u001b[0m         \u001b[0;31m# or the user has disabled using opt_einsum\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 378\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_VF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meinsum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mequation\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moperands\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[attr-defined]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    379\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    380\u001b[0m     \u001b[0mpath\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torch \n",
        "import torch\n",
        "x = 50 * torch.arange(12)\n",
        "# x = torch.tensor([[1,2,3],[4,5,6]])\n",
        "# print(x)\n",
        "x"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1CUiFWa17bQs",
        "outputId": "aa2bfab4-c9f3-4091-9d10-afb091085117"
      },
      "execution_count": 6,
      "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: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch) (4.5.0)\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([  0,  50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 550])"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install einops\n",
        "import einops\n",
        "einops.repeat(x, 'q -> 1 (3 q)')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OoN1zkdt9N2O",
        "outputId": "bacc7aff-3500-47ed-e257-a7e6f5305410"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting einops\n",
            "  Downloading einops-0.6.0-py3-none-any.whl (41 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.6/41.6 KB\u001b[0m \u001b[31m3.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: einops\n",
            "Successfully installed einops-0.6.0\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[  0,  50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 550,   0,  50,\n",
              "         100, 150, 200, 250, 300, 350, 400, 450, 500, 550,   0,  50, 100, 150,\n",
              "         200, 250, 300, 350, 400, 450, 500, 550]])"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "sSgaYi1fGBop"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}