nbs/2B. Whisper quantization (semantic token) model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9158eaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp vq_stoks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cf56fcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ac6d8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import io\n",
    "import sys\n",
    "import time\n",
    "import torch\n",
    "import torchaudio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecbdddfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "from pathlib import Path\n",
    "import json\n",
    "from fastprogress import progress_bar, master_bar\n",
    "import fastprogress\n",
    "import numpy as np\n",
    "import pylab as plt\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "import whisper\n",
    "import whisper.tokenizer\n",
    "from huggingface_hub import hf_hub_download\n",
    "from fastcore.basics import store_attr\n",
    "\n",
    "from torch import nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "import webdataset as wds\n",
    "from whisperspeech import utils, vad_merge\n",
    "\n",
    "from vector_quantize_pytorch import ResidualVQ\n",
    "\n",
    "from fastcore.script import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0ae3314",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "454be05c",
   "metadata": {},
   "source": [
    "# Distill Whisper with a VQ bottleneck\n",
    "\n",
    "Multi-GPU training example:\n",
    "    \n",
    "    python -m whisperspeech.train_multi \\\n",
    "        --task \"vq_stoks base.en-2d-512c-dim64\" --batch-size 32 --epochs 4 \\\n",
    "        --input-dir \"librilight/ librilight-preproc/ --samples 250000\" \\\n",
    "        --tunables=\"--rope --mask_embs --downsample_mean\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f6b51e2",
   "metadata": {},
   "source": [
    "## Prepare the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8f588c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shards = [str(x) for x in Path('/data2/mls-polish/audio//').glob('*.tar')]\n",
    "len(shards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26fa7a5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = vad_merge.chunked_audio_dataset(shards, 'raw').compose(\n",
    "    utils.merge_in(utils.derived_dataset('medium-txt')),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01c64ed7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "609db59e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for x in progress_bar(ds, total=5):\n",
    "    print(f\"{x['__key__']} from {x['__url__']}\")\n",
    "    print()\n",
    "    print(x['txt'])\n",
    "    gain, shift = x['gain_shift.npy']\n",
    "    display(Audio((x['samples'] - shift) * gain, rate=x['sample_rate']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3068df2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Audio\n",
    "\n",
    "def show_samples(shard_spec, filter=lambda x: True):\n",
    "    shards = utils.shard_glob(shard_spec)\n",
    "    ds = vad_merge.chunked_audio_dataset(shards, 'raw').compose(\n",
    "        utils.merge_in(utils.derived_dataset('medium-txt')),\n",
    "        wds.select(filter),\n",
    "    )\n",
    "    for x in progress_bar(ds, total=50):\n",
    "        print(f\"{x['__key__']} from {x['__url__']}\")\n",
    "        print()\n",
    "        print(x['txt'])\n",
    "        display(Audio((x['samples']), rate=x['sample_rate']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9709854",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check short samples for likely VAD errors\n",
    "show_samples('../mls-dutch/audio/*.tar', filter=lambda x: x['tend'] - x['tstart'] < 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "584f36c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def add_masks(samples):\n",
    "    for s in samples:\n",
    "        seconds = s['tend'] - s['tstart']\n",
    "        sr = 16000 #s['sample_rate']\n",
    "        # a mask (downsampled to the Whisper encoder token rate of 50/s) is used\n",
    "        # to teach the model the concept of padding\n",
    "        # this let's us decode shorter sequences later\n",
    "        mask = torch.zeros(30*sr//320, dtype=torch.bool)\n",
    "        mask[:int(seconds * sr) // 320] = 1\n",
    "        s['mask'] = mask\n",
    "        yield s\n",
    "        \n",
    "def get_tokenizer(model, language):\n",
    "    multilingual = not model.endswith(\".en\")\n",
    "    return whisper.tokenizer.get_tokenizer(multilingual, language=language, task=\"transcribe\",\n",
    "                                           num_languages = 100 if model == 'large-v3' else 99)\n",
    "\n",
    "def tokenize_text(samples, ttoks_size=200, model=\"base.en\", language=\"en\"):\n",
    "    tokenizer = get_tokenizer(model, language)\n",
    "    for s in samples:\n",
    "        ttoks = tokenizer.encode(s['txt'])\n",
    "        tokens = list(tokenizer.sot_sequence_including_notimestamps) + ttoks\n",
    "        rpad = ttoks_size - len(tokens)\n",
    "        s['in_ttoks'] = F.pad(torch.tensor(tokens), (0, rpad), value=tokenizer.eot)\n",
    "        s['out_ttoks'] = F.pad(torch.tensor(tokens[1:] + [tokenizer.eot]), (0, rpad), value=-100)\n",
    "        yield s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea1e520d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def load_dataset(\n",
    "        dataset_dir:Path,\n",
    "        txt_label:str=\"base.en-txt\", # the label of the files containing transcriptions\n",
    "        model:str=\"base.en\",\n",
    "        language:str=None,\n",
    "        weight:float=1,\n",
    "        validation:bool=False,\n",
    "        exclude_datasets:str=\"txt-random-valid\", # space separated directory names for validation datasets to exclude\n",
    "    ):\n",
    "    dataset_dir = Path(dataset_dir)\n",
    "    shards = utils.shard_glob(dataset_dir/'audio/*.tar')\n",
    "    with open(dataset_dir/'txt-samples.list') as f: samples = len(f.readlines())\n",
    "    language = utils.readlines(dataset_dir/'language')[0]\n",
    "    \n",
    "    txt_dir = None\n",
    "    for name in ['small.en-txt', 'medium-txt']:\n",
    "        if (dataset_dir/name).exists():\n",
    "            txt_dir = name \n",
    "            break\n",
    "    if txt_dir is None: raise ArgumentError(f\"No transcripts found in {dataset_dir}\")\n",
    "\n",
    "    excludes = {x\n",
    "                for dir in exclude_datasets.split()\n",
    "                for x in utils.readlines(dataset_dir/Path(dir)/\"txt-samples.list\")\n",
    "               } if not validation and exclude_datasets else set()\n",
    "    \n",
    "    if not language and model.endswith('en'): language = 'en'\n",
    "    assert language, \"please provide the dataset language for multilang models\"\n",
    "    \n",
    "    same_on_all_nodes = lambda urls: urls # will only be used for validation\n",
    "    ds = vad_merge.chunked_audio_dataset(shards, 'raw',\n",
    "                                         resampled=not validation, nodesplitter=same_on_all_nodes).compose(\n",
    "        utils.merge_in(utils.derived_dataset(txt_dir)),\n",
    "        wds.select(lambda s: s['__key__'] not in excludes),\n",
    "        utils.resampler(16000, 'samples_16k'),\n",
    "        add_masks,\n",
    "        lambda x: tokenize_text(x, model=model, language=language),\n",
    "        wds.to_tuple('samples_16k', 'mask', 'in_ttoks', 'out_ttoks'),\n",
    "    )\n",
    "    if not validation:\n",
    "        ds = ds.compose(wds.shuffle(500, initial=500))\n",
    "    ds = ds.batched(64)\n",
    "    if validation:\n",
    "        ds = ds.slice(samples // 64)\n",
    "    ds.total_samples = samples\n",
    "    ds.weight = weight\n",
    "    \n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "065814fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = load_dataset('/data2/mls-polish/', model='medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad9e2055",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_ds = load_dataset('/data2/mls-dutch/txt-random-valid/', validation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51e8908f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for x in progress_bar(wds.WebLoader(train_ds, num_workers=4, batch_size=None).unbatched().batched(8), total=100): break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b9fd9ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[torch.Size([8, 480000]),\n",
       " torch.Size([8, 1500]),\n",
       " torch.Size([8, 200]),\n",
       " torch.Size([8, 200])]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[x.shape for x in x]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "148cd5b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = make_model('large-v3-2d-512c-dim64').cuda()\n",
    "model.eval().setup('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1753a22b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    logits, loss = model.forward(*[x.cuda() for x in x])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8be7e7a3",
   "metadata": {},
   "source": [
    "# Training code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebacb9b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "from whisperspeech.train import *\n",
    "from whisperspeech.modules import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61cfaccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "import dataclasses\n",
    "\n",
    "def rand(start, end):\n",
    "    return random.random() * (end - start) + start\n",
    "\n",
    "def logrand(start, end):\n",
    "    return 10**rand(math.log10(start), math.log10(end))\n",
    "\n",
    "@dataclasses.dataclass\n",
    "class Tunables:\n",
    "    init_std :float = 1.5\n",
    "    embeddings_std :float = 4.5e-2\n",
    "    embeddings_lr_scale: float = 1\n",
    "    query_mult :float = 2\n",
    "    rope :bool = True\n",
    "    mask_embs :bool = True # force embeddings corresponding to the input audio padding to a constant value\n",
    "    downsample_conv: bool = False\n",
    "    downsample_mean: bool = True\n",
    "        \n",
    "    codebook_dim: int = 32 # FIXME: unused\n",
    "    codebook_decay: float = 0.9\n",
    "    \n",
    "    lr0 :float = .9e-3\n",
    "    clip_gradient_norm :float = 2\n",
    "    weight_decay :float = 1e-3\n",
    "    warmup_steps :float = 850\n",
    "\n",
    "    random :bool = False\n",
    "\n",
    "    # unused, for backwards compatibility:\n",
    "    output_mult :int = None\n",
    "\n",
    "    def __post_init__(self):\n",
    "        # randomize the hyperparams if requested\n",
    "        if self.random:\n",
    "            self.init_std = logrand(1, 2)\n",
    "            self.embeddings_std = logrand(3e-2,6e-2)\n",
    "            self.embeddings_lr_scale = 2**rand(0,3)\n",
    "            self.query_mult = logrand(1,8)\n",
    "            self.codebook_dim = int(logrand(30,50))\n",
    "            self.codebook_decay = logrand(0.86,0.95)\n",
    "            self.rope = True\n",
    "            self.mask_embs = True\n",
    "            self.downsample_mean = True\n",
    "            \n",
    "            self.lr0 = logrand(.8e-3,1e-3)\n",
    "            self.clip_gradient_norm = 10**rand(-1,1)\n",
    "            self.warmup_steps = logrand(700,1000)\n",
    "            \n",
    "    @staticmethod\n",
    "    def upgrade(args):\n",
    "        args = {k:v for k,v in args.items()}\n",
    "        def old_default(name, value):\n",
    "            if name not in args: args[name] = value\n",
    "        old_default('output_mult', 1)\n",
    "        old_default('query_mult', 1)\n",
    "        old_default('rope', False)\n",
    "        old_default('mask_embs', False)\n",
    "        old_default('downsample_conv', False)\n",
    "        old_default('downsample_mean', False)\n",
    "        if 'encoder_depth_ratio' in args: del args['encoder_depth_ratio']\n",
    "        if 'vq_codes' in args: del args['vq_codes']\n",
    "        return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "957549ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7e0227d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class RQBottleneckTransformer(nn.Module):\n",
    "    def __init__(self, vq_codes=512, q_depth=12, depth=1, n_head=2, head_width=64, ffn_mult=4,\n",
    "                 codebook_dim=2, threshold_ema_dead_code=2, use_cosine_sim = False, kl_loss_mul=1,\n",
    "                 downsample=1, no_quantize=False,\n",
    "                 whisper_model_name='tiny.en', tunables=Tunables()):\n",
    "        super().__init__()\n",
    "        width = n_head * head_width\n",
    "        store_attr(\"codebook_dim,vq_codes,q_depth,n_head,head_width,ffn_mult,depth,use_cosine_sim,downsample,whisper_model_name\")\n",
    "        self.width = width\n",
    "        self.base_width = 3 * head_width\n",
    "        self.vq_codes = vq_codes\n",
    "        self.tunables = tunables\n",
    "        self.stoks_len = 1500//downsample\n",
    "        self.stoks_per_sec = self.stoks_len//30\n",
    "        self.no_quantize = no_quantize\n",
    "                \n",
    "        qk_scale = self.tunables.query_mult * 8 / math.sqrt(head_width)\n",
    "        \n",
    "        self.kl_loss_mul = kl_loss_mul\n",
    "        \n",
    "        if no_quantize:\n",
    "            # a mode to get Whisper baselines into W&B easily, skips all training\n",
    "            self.fake_parameter = nn.Parameter(torch.tensor(0.001))\n",
    "        else:\n",
    "            n_mlp = width * ffn_mult\n",
    "            self.mlp = nn.Sequential(\n",
    "                nn.Linear(width, n_mlp), nn.GELU(), nn.Linear(n_mlp, width)\n",
    "            )\n",
    "            self.mlp_ln = LayerNorm(width)\n",
    "\n",
    "            if tunables.downsample_conv:\n",
    "                self.downsample_conv = nn.Conv1d(width, width, kernel_size=3, stride=downsample, padding=1)\n",
    "            else:\n",
    "                self.downsample_conv = None\n",
    "\n",
    "            if tunables.mask_embs: vq_codes = vq_codes + 1\n",
    "            self.rq = ResidualVQ(\n",
    "                dim = width,\n",
    "                codebook_size = vq_codes, # codebook size\n",
    "                decay = tunables.codebook_decay, # the exponential moving average decay, lower means the dictionary will change faster\n",
    "                commitment_weight = 1.,   # the weight on the commitment loss\n",
    "                threshold_ema_dead_code = threshold_ema_dead_code,\n",
    "                use_cosine_sim = use_cosine_sim,\n",
    "                codebook_dim = codebook_dim,\n",
    "                num_quantizers= 1,\n",
    "            )\n",
    "\n",
    "            self.positional_embedding = nn.Embedding(1500, width) # FIXME: should be self.stoks_len\n",
    "\n",
    "            self._out_blocks = nn.Sequential(*[\n",
    "                ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, rope=tunables.rope) for _ in range(depth)\n",
    "            ])\n",
    "            self.ln_post = LayerNorm(width)\n",
    "\n",
    "        self.positions = torch.arange(0, 1500, dtype=torch.long)\n",
    "        \n",
    "        self.ce_lossf = nn.CrossEntropyLoss(ignore_index=-100)\n",
    "        self.kl_lossf = nn.KLDivLoss(reduction='batchmean')\n",
    "                \n",
    "        self.whmodel = None\n",
    "\n",
    "        self.apply(self.init_transformer)\n",
    "        self.register_buffer('val_true', torch.zeros(1))\n",
    "        self.register_buffer('val_total', torch.zeros(1))\n",
    "    \n",
    "    def setup(self, device):\n",
    "        self.ensure_whisper(device)\n",
    "    \n",
    "    def init_transformer(self, m):\n",
    "        if isinstance(m, LinearHead):\n",
    "            m.no_weight_decay = True\n",
    "            torch.nn.init.constant_(m.weight, 0)\n",
    "        elif isinstance(m, QueryHead):\n",
    "            m.lr_scale = 1/(m.weight.shape[1] / self.base_width)\n",
    "            torch.nn.init.constant_(m.weight, 0)\n",
    "        elif isinstance(m, nn.Embedding):\n",
    "            m.no_weight_decay = True\n",
    "            m.lr_scale = self.tunables.embeddings_lr_scale\n",
    "            std = self.tunables.embeddings_std\n",
    "            torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)\n",
    "        elif isinstance(m, nn.Linear):\n",
    "            m.lr_scale = 1/(m.weight.shape[1] / self.base_width)\n",
    "            std = self.tunables.init_std / m.weight.shape[1]\n",
    "            torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)\n",
    "            if m.bias is not None:\n",
    "                torch.nn.init.trunc_normal_(m.bias, std=std, a=-3*std, b=3*std)\n",
    "        elif isinstance(m, nn.LayerNorm):\n",
    "            m.no_weight_decay = True\n",
    "            torch.nn.init.constant_(m.bias, 0)\n",
    "            torch.nn.init.constant_(m.weight, 1)\n",
    "\n",
    "    @property\n",
    "    def device(self):\n",
    "        return next(self.parameters()).device\n",
    "            \n",
    "    #\n",
    "    # training\n",
    "    #\n",
    "    def log_mel_spectrogram(self, samples):\n",
    "        return whisper.log_mel_spectrogram(samples, 128 if self.whisper_model_name == 'large-v3' else 80)\n",
    "    \n",
    "    @torch.no_grad()\n",
    "    def extract_teacher(self, samples, input_toks, output_toks):\n",
    "        embs = self.whmodel[0].encoder(self.log_mel_spectrogram(samples))\n",
    "        teacher_logits = self.whmodel[0].decoder(input_toks, embs)\n",
    "        # set teacher logits to 0 for padding positions so KLDivLoss ignores them\n",
    "        teacher_logits[output_toks == -100] = 0\n",
    "        return embs, teacher_logits\n",
    "    \n",
    "    def downsample_embeddings(self, x):\n",
    "        if self.downsample_conv is not None:\n",
    "            return x[:,::self.downsample] + self.downsample_conv(x.transpose(-1,-2)).transpose(-2,-1)\n",
    "        elif self.tunables.downsample_mean:\n",
    "            bs,slen,depth = x.shape\n",
    "            return x.reshape(bs,slen//self.downsample,self.downsample,depth).mean(-2)\n",
    "        else:\n",
    "            return x[:,::self.downsample]\n",
    "        \n",
    "    def out_blocks(self, x):\n",
    "        for l in self._out_blocks: x = l(x, self.positions)\n",
    "        return x\n",
    "    \n",
    "    def forward(self, samples, mask, input_toks, output_toks):\n",
    "        embs, teacher_logits = self.extract_teacher(samples, input_toks, output_toks)\n",
    "        \n",
    "        if not self.no_quantize:\n",
    "            x = self.downsample_embeddings(embs)\n",
    "            x = x + self.mlp(self.mlp_ln(x))\n",
    "            # VQ bottleneck\n",
    "            quantized, self.indices, self.commit_loss = self.rq(x)\n",
    "            self.commit_loss = self.commit_loss.mean()\n",
    "\n",
    "            x = quantized.repeat_interleave(self.downsample, -2)\n",
    "            project_out = getattr(self.rq, 'project_out', None) or self.rq.layers[0].project_out\n",
    "            if self.tunables.mask_embs: x[~mask] = project_out(self.rq.layers[0]._codebook.embed[0,self.vq_codes])\n",
    "            x = x + self.positional_embedding(self.positions.to(x.device))\n",
    "            x = self.ln_post(self.out_blocks(x))\n",
    "        \n",
    "        logits = self.whmodel[0].decoder(input_toks, embs if self.no_quantize else x)\n",
    "        self.ce_loss = self.ce_lossf(logits.view(-1,logits.shape[-1]), output_toks.view(-1))\n",
    "        self.kl_loss = self.kl_lossf(F.log_softmax(logits, dim=-1), F.softmax(teacher_logits, dim=-1))\n",
    "        loss = self.ce_loss + self.kl_loss_mul * self.kl_loss\n",
    "        if not self.no_quantize: loss += self.commit_loss\n",
    "        x = None\n",
    "        if self.no_quantize: loss = loss + self.fake_parameter\n",
    "        \n",
    "        if not self.training:\n",
    "            valid_toks = output_toks != -100\n",
    "            self.val_true += (logits.detach().argmax(-1)[valid_toks] == output_toks[valid_toks]).float().sum()\n",
    "            self.val_total += valid_toks.float().sum()\n",
    "\n",
    "        return x, logits, loss\n",
    "\n",
    "    def get_metrics(self):\n",
    "        metrics = {\n",
    "            'acc_0': (self.val_true / self.val_total).item(),\n",
    "        }\n",
    "        self.val_true[:] = 0\n",
    "        self.val_total[:] = 0\n",
    "        return metrics\n",
    "    \n",
    "    #\n",
    "    # inference\n",
    "    #\n",
    "    @classmethod\n",
    "    def load_model(cls, ref=\"collabora/spear-tts-pytorch:whisper-vq-stoks-medium-en+pl.model\",\n",
    "                   repo_id=None, filename=None, local_filename=None):\n",
    "        if repo_id is None and filename is None and local_filename is None:\n",
    "            if \":\" in ref:\n",
    "                repo_id, filename = ref.split(\":\", 1)\n",
    "            else:\n",
    "                local_filename = ref\n",
    "        if not local_filename:\n",
    "            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)\n",
    "        spec = torch.load(local_filename) \n",
    "        vqmodel = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec.get('tunables', {}))))\n",
    "        vqmodel.load_state_dict(spec['state_dict'])\n",
    "        vqmodel.eval()\n",
    "        return vqmodel\n",
    "    \n",
    "    def load_checkpoint(self, local_filename):\n",
    "        spec = torch.load(local_filename, map_location='cpu')\n",
    "        assert 'pytorch-lightning_version' in spec, 'not a valid PyTorch Lightning checkpoint'\n",
    "        state_dict = {k.replace('model.', ''):v\n",
    "                      for k,v in spec['state_dict'].items()}\n",
    "        self.load_state_dict(state_dict)\n",
    "        return self\n",
    "    \n",
    "    def save_model(self, fname, store_parameters=True):\n",
    "        torch.save(dict(config = self.__stored_args__,\n",
    "                        tunables = dataclasses.asdict(self.tunables),\n",
    "                        state_dict = self.state_dict() if store_parameters else None), fname)\n",
    "        \n",
    "    def ensure_whisper(self, device=None):\n",
    "        if self.whmodel is not None: return\n",
    "        device = device or self.device\n",
    "        # the list wrapper is a hack to make sure the whole of Whisper is not sucked into self.parameters()\n",
    "        if self.whmodel is None: self.whmodel = [whisper.load_model(self.whisper_model_name, device=device)]\n",
    "        self.decoding_options = whisper.DecodingOptions()\n",
    "        self.tokenizer = get_tokenizer(self.whisper_model_name, None)\n",
    "    \n",
    "    def quantize(self, embs):\n",
    "        x = self.downsample_embeddings(embs)\n",
    "        x = x + self.mlp(self.mlp_ln(x))\n",
    "        _, stoks, _ = self.rq(x)\n",
    "        if self.q_depth == 1:\n",
    "            stoks = stoks.squeeze(-1)\n",
    "        return stoks\n",
    "\n",
    "    def dequantize(self, stoks):\n",
    "        assert self.q_depth == 1\n",
    "        assert len(stoks.shape) == 1, \"batch processing is not supported\"\n",
    "        if isinstance(stoks, np.ndarray): stoks = torch.tensor(stoks)\n",
    "        # remove padding\n",
    "        padding = torch.nonzero(stoks == self.vq_codes)\n",
    "        if padding.any(): stoks = stoks[:padding[0,0]]\n",
    "        stoks = F.pad(stoks, (0,self.stoks_len - stoks.shape[-1]), value=self.vq_codes if self.tunables.mask_embs else 0)\n",
    "        x = self.rq.layers[0]._codebook.embed[0,stoks.to(torch.long).view(-1)]\n",
    "        x = x.repeat_interleave(self.downsample, -2)\n",
    "        project_out = getattr(self.rq, 'project_out', None) or self.rq.layers[0].project_out\n",
    "        x = project_out(x).unsqueeze(0)\n",
    "        positions = torch.arange(0, x.shape[-2], dtype=torch.long, device=x.device)\n",
    "        x = x + self.positional_embedding(positions)\n",
    "        return self.ln_post(self.out_blocks(x))\n",
    "\n",
    "    def encode_audio(self, audio):\n",
    "        if isinstance(audio, str):\n",
    "            x, sr = torchaudio.load(audio)\n",
    "            x = torchaudio.transforms.Resample(sr, 16000)(x)[0]\n",
    "            audio = x.unsqueeze(0)\n",
    "        return self.encode_mel(self.log_mel_spectrogram(audio).to(self.device))\n",
    "    \n",
    "    def encode_mel(self, mel):\n",
    "        assert len(mel.shape) == 3, \"invalid mel spectrogram shape, expect (batch,chn,time)\"\n",
    "        self.ensure_whisper()\n",
    "        n = mel.shape[-1]\n",
    "        if n > whisper.audio.N_FRAMES:\n",
    "            padding = 0\n",
    "            padded = mel[:,:,:whisper.audio.N_FRAMES]\n",
    "        else:\n",
    "            padding = -n % whisper.audio.N_FRAMES\n",
    "            padded = F.pad(mel, (0, padding), value=-1.5)\n",
    "        embs = self.whmodel[0].encoder(padded)#.to(self.whmodel[0].device))#[:,:n//2]\n",
    "        stoks = self.quantize(embs)\n",
    "        if self.tunables.mask_embs:\n",
    "            return stoks[:,:n//2//self.downsample]\n",
    "        else:\n",
    "            return stoks\n",
    "    \n",
    "    def decode_text(self, stoks, decoding_options=None):\n",
    "        self.ensure_whisper(self.device)\n",
    "        if decoding_options is None: decoding_options = self.decoding_options\n",
    "        embs = self.dequantize(stoks).to(self.whmodel[0].device)\n",
    "        return self.whmodel[0].decode(embs, decoding_options)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "768f49f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pylab as plt\n",
    "import fastprogress\n",
    "import IPython\n",
    "import numpy as np\n",
    "\n",
    "class RQVisual:\n",
    "    def __init__ (self, model, masterbar, total_steps):\n",
    "        self.model = model\n",
    "        self.masterbar = masterbar\n",
    "        self.total_steps = total_steps\n",
    "        self.epochs = total_steps // masterbar.main_bar.total\n",
    "        \n",
    "        gs = plt.GridSpec(3, 1, height_ratios=[2,2,1])\n",
    "        graph_fig = plt.figure(figsize=(10,6))\n",
    "        self.graph_fig = graph_fig\n",
    "        self.loss_p = graph_fig.add_subplot(gs[0])\n",
    "        self.acc_p = graph_fig.add_subplot(gs[1], sharex=self.loss_p)\n",
    "        self.acc_p.tick_params('x', labelbottom=False)\n",
    "        self.lr_p = graph_fig.add_subplot(gs[2], sharex=self.loss_p)\n",
    "        self.lr_p.tick_params('x', labelbottom=False)\n",
    "        self.graph_out = None\n",
    "        \n",
    "        self.its = []\n",
    "        self.train_losses = []\n",
    "        self.val_losses = []\n",
    "        self.lr_history = []\n",
    "        self.entropy = np.nan\n",
    "        self.entropy_history = []\n",
    "            \n",
    "    def show(self):\n",
    "        self.start_t = time.time()\n",
    "        self.masterbar.write([\"samples\", \"train\", \"val\", \"codebook entropy\", \"time\"], table=True)\n",
    "        self.graph_out = display(self.graph_fig, display_id=True)\n",
    "        self.entropy_out = display(IPython.display.HTML(''), display_id=True)\n",
    "    \n",
    "    def hide(self):\n",
    "        if self.graph_out is not None:\n",
    "            self.graph_out.update(IPython.display.HTML(''))\n",
    "    \n",
    "    def plot(self):\n",
    "        loss_p, acc_p, lr_p = self.loss_p, self.acc_p, self.lr_p\n",
    "        loss_p.clear()\n",
    "        loss_p.plot(self.its, self.train_losses)\n",
    "        loss_p.plot(self.its, self.val_losses)\n",
    "        loss_p.set_xlim(10000, self.total_steps)\n",
    "        loss_p.set_xscale('log')\n",
    "        loss_p.set_yscale('log')\n",
    "        acc_p.clear()\n",
    "        acc_p.plot(self.its, np.stack(self.entropy_history), ':')\n",
    "        lr_p.clear()\n",
    "        lrs = np.array(self.lr_history)\n",
    "        lr_p.plot(self.its, lrs)\n",
    "        self.graph_out.update(self.graph_fig)\n",
    "    \n",
    "    def add_data(self, it, lr, train_loss, val_los):\n",
    "        self.its.append(it)\n",
    "        self.train_losses.append(train_loss)\n",
    "        self.val_losses.append(val_los)\n",
    "        self.lr_history.append(lr)\n",
    "        with torch.no_grad():\n",
    "            cls = vqmodel.rq.layers[0]._codebook.cluster_size\n",
    "            pdf = cls / cls.sum()\n",
    "            entropy = -torch.nansum(pdf * pdf.log2())\n",
    "        self.entropy_history.append(entropy.cpu().numpy())\n",
    "        self.entropy_out.update(f\"Entropy: {self.entropy_history[-1]:.2f}\")\n",
    "        self.plot()\n",
    "\n",
    "    def add_table_row(self, it, avg_train_loss, val_loss):\n",
    "        elapsed_t = time.time() - self.start_t\n",
    "        self.masterbar.write([it, f\"{avg_train_loss:.5f}\", f\"{val_loss:.5f}\", f\"{self.entropy_history[-1]:.2f}\", fastprogress.core.format_time(elapsed_t)], table=True)\n",
    "    \n",
    "    def on_iter(self, bar, it, avg_train_loss, val_loss):\n",
    "        epoch = math.ceil(it / self.total_steps * self.epochs)\n",
    "        bar.comment = f\"#{epoch}/{self.epochs} loss: {avg_train_loss:.3f} / {val_loss:.3f}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e38acb1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def make_model(size:str, no_quantize=False, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):\n",
    "    common = dict(\n",
    "        q_depth=1, depth=1, threshold_ema_dead_code=0, use_cosine_sim=True, tunables=tunables,\n",
    "        no_quantize = no_quantize,\n",
    "    )\n",
    "    if size == 'base.en-2d-4096c':\n",
    "        model = RQBottleneckTransformer(codebook_dim=32, vq_codes=4096, n_head=8, downsample=2, \n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'base.en-2d-512c':\n",
    "        model = RQBottleneckTransformer(codebook_dim=32, vq_codes=512, n_head=8, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'base.en-2d-512c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=8, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'base-2d-512c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=8, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'base-2d-1024c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=1024, n_head=8, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'medium-2d-256c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=256, n_head=16, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'medium-2d-256c-dim128':\n",
    "        model = RQBottleneckTransformer(codebook_dim=128, vq_codes=256, n_head=16, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'medium-2d-512c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=16, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'medium-2d-512c-dim128':\n",
    "        model = RQBottleneckTransformer(codebook_dim=128, vq_codes=512, n_head=16, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'medium-2d-512c-dim256':\n",
    "        model = RQBottleneckTransformer(codebook_dim=256, vq_codes=512, n_head=16, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'medium-2d-1024c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=1024, n_head=16, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'medium-2d-2048c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=2048, n_head=16, downsample=2,\n",
    "                                        whisper_model_name=size.split(\"-\")[0], **common)\n",
    "        return model\n",
    "    if size == 'large-v2-2d-512c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=20, downsample=2,\n",
    "                                        whisper_model_name='large-v2', **common)\n",
    "        return model\n",
    "    if size == 'large-v3-2d-512c-dim64':\n",
    "        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=20, downsample=2,\n",
    "                                        whisper_model_name='large-v3', **common)\n",
    "        return model\n",
    "    raise ArgumentError(f\"invalid model size: {size}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}