{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Run WannabeTokenizer\n",
        "\n",
        "Set the three paths below, then run the cells in order. Start Jupyter from this package folder so the relative checkpoint path works."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "import subprocess\n",
        "import sys\n",
        "\n",
        "PACKAGE_DIR = Path.cwd()\n",
        "TOKENIZER = PACKAGE_DIR / 'checkpoints' / 'WannabeTokenizer.safetensors'\n",
        "YUE2_CHECKPOINT = Path('/path/to/yue2_3b_bf16.safetensors')\n",
        "OUTPUT_DIR = PACKAGE_DIR / 'output'\n",
        "\n",
        "DEVICE = 'auto'  # 'auto', 'cuda', 'cuda:0', or 'cpu'\n",
        "SEED = 4  # Shout out to eMuhdra\n",
        "STEPS = 32\n",
        "OVERWRITE = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "INPUT_AUDIO = Path('/path/to/input.wav')\n",
        "command = [\n",
        "    sys.executable,\n",
        "    str(PACKAGE_DIR / 'run_tokenizer.py'),\n",
        "    str(INPUT_AUDIO),\n",
        "    '--tokenizer', str(TOKENIZER),\n",
        "    '--yue2-checkpoint', str(YUE2_CHECKPOINT),\n",
        "    '--output-dir', str(OUTPUT_DIR),\n",
        "    '--device', DEVICE,\n",
        "    '--seed', str(SEED),\n",
        "    '--steps', str(STEPS),\n",
        "]\n",
        "if OVERWRITE:\n",
        "    command.append('--overwrite')\n",
        "\n",
        "subprocess.run(command, cwd=PACKAGE_DIR, check=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from IPython.display import Audio, display\n",
        "\n",
        "tokens = np.load(OUTPUT_DIR / 'tokens.npy')\n",
        "print(f'{len(tokens)} tokens at 25 Hz; dtype={tokens.dtype}; range={tokens.min()}..{tokens.max()}')\n",
        "display(Audio(filename=str(OUTPUT_DIR / 'reconstructed.wav')))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}