nbs/1C. VAD merging.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fdbe3b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp vad_merge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cf56fcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecbdddfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from fastprogress import progress_bar\n",
    "from fastcore.script import *\n",
    "\n",
    "from whisperspeech import utils\n",
    "import webdataset as wds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1d80d3b",
   "metadata": {},
   "source": [
    "# VAD merging\n",
    "\n",
    "We merge the VAD segments into longer chunks to make training more efficient (otherwise we'll spend a lot of time calculating padding)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4eb5cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = wds.WebDataset(['../wolnelektury-wds2/wolnelektury-audio-000000.tar']).compose(\n",
    "    wds.decode(wds.torch_audio),\n",
    "    utils.merge_in(utils.derived_dataset('vad')),\n",
    "    utils.find_audio,\n",
    "    utils.split_to_chunks,\n",
    "    utils.merge_in(utils.derived_dataset('spk_emb')),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71a85cd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd83c05",
   "metadata": {},
   "outputs": [],
   "source": [
    "prev = None\n",
    "for s in progress_bar(ds, total=20):\n",
    "    sim = F.cosine_similarity(torch.tensor(s['spk_emb.npy']), torch.tensor((prev if prev is not None else s)['spk_emb.npy']), dim=0)\n",
    "    secs = s['tend'] - s['tstart']\n",
    "    same = sim > 0.6 if secs > 2 else sim > 0.1\n",
    "    if not same: print(\"new\")\n",
    "    print(s['__key__'], sim, secs)\n",
    "    display(IPython.display.Audio(s['samples'], rate=s['sample_rate']))\n",
    "    if secs > 2:\n",
    "        prev = s\n",
    "    time.sleep(.5)\n",
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "875a4bba",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "# we need to split first to merge in the spk_emb.npy data\n",
    "# this is similar to utils.split_to_chunks but works without the audio data\n",
    "def split(stream, ikey='vad.npy', copy_keys=[], split_keys=[]):\n",
    "    for s in stream:\n",
    "        imax = len(s[ikey]) - 1\n",
    "        if len(s[ikey]) == 0:\n",
    "            # Preserve info about audio files without any speech.\n",
    "            # We need to push this info through a weird side-channel \n",
    "            # because we want to be able to a merge with naively\n",
    "            # splitted data.\n",
    "            new = {\"__key__\": s['__key__'] + \"_none\",\n",
    "                   \"src_key\": s['__key__'],\n",
    "                   \"__url__\": s['__url__'],\n",
    "                   \"__skip_merge__\": True}\n",
    "            for k in copy_keys:  new[k] = np.array([])\n",
    "            for k in split_keys: new[k] = np.array([])\n",
    "            new[ikey] = s[ikey]\n",
    "            yield new\n",
    "        for i,(ts,te) in enumerate(s[ikey]):\n",
    "            new = {\"__key__\": s['__key__'] + f\"_{i:03d}\",\n",
    "                   \"src_key\": s['__key__'],\n",
    "                   \"__url__\": s['__url__'],\n",
    "                   \"i\": i, \"imax\": imax}\n",
    "            for k in copy_keys:  new[k] = s[k]\n",
    "            for k in split_keys: new[k] = s[k][i]\n",
    "            new[ikey] = s[ikey][i]\n",
    "            yield new\n",
    "\n",
    "def merge_by_src_key(stream, copy_keys=[], merge_keys=['vad.npy']):\n",
    "    def make_record(src):\n",
    "        s = {\n",
    "            \"__url__\": src['__url__'],\n",
    "            \"__key__\": src['src_key'],\n",
    "        }\n",
    "        for k in copy_keys: s[k] = src[k]\n",
    "        for k in merge_keys: s[k] = []\n",
    "        return s\n",
    "    def finish_record(s):\n",
    "        for k in merge_keys: s[k] = np.array(s[k])\n",
    "        return s\n",
    "    ms = None\n",
    "    for s in stream:\n",
    "        try:\n",
    "            # push accumulated data\n",
    "            if ms and s['src_key'] != ms['__key__']:\n",
    "                yield finish_record(ms)\n",
    "                ms = None\n",
    "            # prepare a merged record for the new data\n",
    "            if ms is None:\n",
    "                ms = make_record(s)\n",
    "            for k in merge_keys:\n",
    "                if k in s: ms[k].append(s[k])\n",
    "        except:\n",
    "            print(f\"Error processing {s['__key__']}:\")\n",
    "            print(s)\n",
    "            raise\n",
    "    yield finish_record(ms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9de4aa8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = wds.WebDataset(['../librilight/vad/librilight-small-flac-000000.tar.gz']).compose(\n",
    "    wds.decode(),\n",
    "    lambda x: split(x, copy_keys=['gain_shift.npy'], split_keys=['powers.npy']),\n",
    "    utils.merge_in(utils.derived_dataset('spk_emb')),\n",
    "    lambda x: merge_by_src_key(x, copy_keys=['gain_shift.npy'], merge_keys=['powers.npy', 'vad.npy']),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5635dda",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Automatic pdb calling has been turned ON\n"
     ]
    }
   ],
   "source": [
    "%pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13cea1b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['__url__', '__key__', 'gain_shift.npy', 'powers.npy', 'vad.npy'])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for s in ds: break\n",
    "s.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393b7c84",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def random_cutter(dur):\n",
    "    if random.random() < 0.5:\n",
    "        return dur > 30 * (random.random()*0.95+0.05)\n",
    "    else:\n",
    "        return dur > 30\n",
    "\n",
    "def random_cutter2(dur):\n",
    "    if random.random() < 0.25:\n",
    "        return True\n",
    "    else:\n",
    "        return dur > 30 * (random.random()*0.95+0.05)\n",
    "    \n",
    "def chunk_merger(prefix, should_cut=lambda x: x > 30):\n",
    "    def _merger(stream):\n",
    "        for s in stream:\n",
    "            segments, speakers = s['vad.npy'], s['spk_emb.npy']\n",
    "            if segments.size == 0:\n",
    "                s[prefix+'.vad.npy'], s[prefix+'.spk_emb.npy'] = np.array([]), np.array([])\n",
    "                s[prefix+'.subvads.pyd'] = []\n",
    "                yield s\n",
    "                continue\n",
    "            curr_start = segments[0][0]\n",
    "            curr_end = 0\n",
    "            curr_spk = None\n",
    "            curr_chunks = []\n",
    "            spk_acc = torch.tensor(speakers[0])\n",
    "            spk_acc_N = 1\n",
    "            merged = []\n",
    "            merged_chunks = []\n",
    "            merged_spk = []\n",
    "\n",
    "            for (ts,te),new_spk in zip(segments, speakers):\n",
    "                secs = te - ts\n",
    "                new_spk = torch.tensor(new_spk)\n",
    "                spk_change = False\n",
    "                if curr_spk is not None:\n",
    "                    sim = F.cosine_similarity(curr_spk, new_spk, dim=0)\n",
    "                    spk_change = sim < 0.5 if secs > 2 else sim < 0.1\n",
    "                if (spk_change or should_cut(te - curr_start)) and curr_end - curr_start > 0:\n",
    "                    merged.append((curr_start, curr_end))\n",
    "                    merged_spk.append(spk_acc / spk_acc_N)\n",
    "                    merged_chunks.append(curr_chunks)\n",
    "                    curr_start = ts\n",
    "                    spk_acc = new_spk\n",
    "                    curr_chunks = []\n",
    "                curr_spk = new_spk\n",
    "                if secs > 2:\n",
    "                    spk_acc += new_spk\n",
    "                    spk_acc_N += 1\n",
    "                curr_end = te\n",
    "                curr_chunks.append((ts, te))\n",
    "            merged.append((curr_start, curr_end))\n",
    "            merged_spk.append(spk_acc / spk_acc_N)\n",
    "            merged_chunks.append(curr_chunks)\n",
    "            s[prefix+'.vad.npy'], s[prefix+'.spk_emb.npy'] = np.array(merged), torch.stack(merged_spk).numpy()\n",
    "            s[prefix+'.subvads.pyd'] = merged_chunks\n",
    "            yield s\n",
    "    return _merger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a60e399e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = wds.WebDataset([utils.derived_name('../wolnelektury-wds2/wolnelektury-audio-000000.tar', 'vad')]).compose(\n",
    "    wds.decode(),\n",
    "    split,\n",
    "    utils.merge_in(utils.derived_dataset('spk_emb', base='vad', suffix='')),\n",
    "    merge_by_src_key,\n",
    "    chunk_merger,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5b171d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for s in ds: break\n",
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11ddaf1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = wds.WebDataset(['../wolnelektury-wds2/wolnelektury-audio-000000.tar']).compose(\n",
    "    wds.decode(wds.torch_audio),\n",
    "    utils.merge_in(utils.derived_dataset('vad')),\n",
    "    utils.find_audio,\n",
    "    utils.split_to_chunks,\n",
    "    utils.merge_in(utils.derived_dataset('spk_emb')),\n",
    "    merge_by_src_key,\n",
    "    chunk_merger,\n",
    "    utils.merge_in(utils.derived_dataset('audio', suffix='', decoders=[wds.torch_audio])),\n",
    "    utils.find_audio,\n",
    "    lambda x: utils.split_to_chunks(x, metakeys=['spk_emb.npy']),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2803c4c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for s in ds: break\n",
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd15ddbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "prev = None\n",
    "for s in progress_bar(ds, total=20):\n",
    "    sim = F.cosine_similarity(torch.tensor(s['spk_emb.npy']), torch.tensor((prev if prev is not None else s)['spk_emb.npy']), dim=0)\n",
    "    secs = s['tend'] - s['tstart']\n",
    "    same = sim > 0.6 if secs > 2 else sim > 0.1\n",
    "    if not same: print(\"new\")\n",
    "    print(s['__key__'], sim, secs, sum([e-s for s,e in s['orig_s']['subvads.pyd'][s['i']]]))\n",
    "    display(IPython.display.Audio(s['samples'], rate=s['sample_rate']))\n",
    "    if secs > 2:\n",
    "        prev = s\n",
    "    time.sleep(.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b38b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "# we filter before splitting to keep empty merged samples even if we filter out everything\n",
    "def filter_bad_samples(stream):\n",
    "    for s in stream:\n",
    "        if 'librilight' in s['__url__'] or 'test-shard.tar' in s['__url__']:\n",
    "            for k in ['vad.npy', 'spk_emb.npy', 'powers.npy']:\n",
    "                s[k] = s[k][1:-1]\n",
    "\n",
    "        if s['vad.npy'].size > 0:\n",
    "            lengths = s['vad.npy'][:,1] - s['vad.npy'][:,0]\n",
    "            mask = (lengths < 1) & (s['powers.npy'] < -6)\n",
    "            for k in ['vad.npy', 'spk_emb.npy', 'powers.npy']:\n",
    "                s[k] = s[k][~mask]\n",
    "        yield s\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14261153",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_samples(stream):\n",
    "    for s in stream:\n",
    "        print(s['__key__'], s['vad.npy'])\n",
    "        yield s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f271d55",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "@call_parse\n",
    "def prepare_mvad(\n",
    "    input:str,  # input VAD shard path\n",
    "    output:str, # output shard path\n",
    "    eqvad:bool=False, # make the chunk length distribution more uniform\n",
    "    ignore_spk_emb:bool=False,\n",
    "):    \n",
    "    if ignore_spk_emb:\n",
    "        def chg_spk_emb(stream):\n",
    "            for s in stream:\n",
    "                for x in s['spk_emb.npy']: x[:] = 1\n",
    "                yield s\n",
    "    else:\n",
    "        def chg_spk_emb(stream):\n",
    "            for s in stream: yield s\n",
    "    \n",
    "    ds = wds.WebDataset([input]).compose(\n",
    "        wds.decode(),\n",
    "        lambda x: split(x, copy_keys=['gain_shift.npy'], split_keys=['powers.npy']),\n",
    "        utils.merge_in(utils.derived_dataset('spk_emb')),\n",
    "        lambda x: merge_by_src_key(x, copy_keys=['gain_shift.npy'], merge_keys=['powers.npy', 'vad.npy', 'spk_emb.npy']),\n",
    "        filter_bad_samples,\n",
    "        chg_spk_emb,\n",
    "        chunk_merger('raw', lambda x: True),\n",
    "        chunk_merger('eq', random_cutter),\n",
    "        chunk_merger('max')\n",
    "    )\n",
    "\n",
    "    with utils.AtomicTarWriter(output) as sink:\n",
    "        for s in progress_bar(ds, total='noinfer'):\n",
    "#             if len(s['vad.npy']) > 1:\n",
    "#                 print(s)\n",
    "            del s['vad.npy'], s['spk_emb.npy'], s['powers.npy']\n",
    "            sink.write(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9239505",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Automatic pdb calling has been turned OFF\n"
     ]
    }
   ],
   "source": [
    "%pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86024335",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    progress:not([value]), progress:not([value])::-webkit-progress-bar {\n",
       "        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px);\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      <progress value='5' class='' max='5' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      100.00% [5/5 00:00&lt;00:00]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prepare_mvad('../test-dataset/vad/test-shard.tar.gz', '../test-dataset/mvad/test-shard.tar.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "737a49aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def find_vad_kind(kind):\n",
    "    def _finder(stream):\n",
    "        for s in stream:\n",
    "            for k in ['vad.npy', 'spk_emb.npy']:\n",
    "                s[k] = s[f'{kind}.{k}']\n",
    "            yield s\n",
    "    return _finder\n",
    "\n",
    "def chunked_audio_dataset(shards, kind='max', copy_keys=['gain_shift.npy'], split_keys=['spk_emb.npy'],\n",
    "                          resampled=False, nodesplitter=wds.shardlists.single_node_only):\n",
    "    return wds.WebDataset(shards, resampled=resampled, nodesplitter=nodesplitter).compose(\n",
    "        wds.decode(utils.torch_audio_opus),\n",
    "        utils.find_audio,\n",
    "        utils.merge_in(utils.derived_dataset('mvad')),\n",
    "        find_vad_kind(kind),\n",
    "        lambda x: utils.split_to_chunks(x, copy_keys=copy_keys, split_keys=split_keys),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00b924c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = chunked_audio_dataset(['../wolnelektury-wds2/wolnelektury-audio-000000.tar'])\n",
    "prev = None\n",
    "for s in progress_bar(ds, total=6):\n",
    "    sim = F.cosine_similarity(torch.tensor(s['spk_emb.npy']), torch.tensor((prev if prev is not None else s)['spk_emb.npy']), dim=0)\n",
    "    if sim < 0.5: print(\"new\")\n",
    "    print(s['__key__'], sim, s['tend'] - s['tstart'], sum([e-s for s,e in s['orig_s']['subvads.pyd'][s['i']]]))\n",
    "    display(IPython.display.Audio(s['samples'], rate=s['sample_rate']))\n",
    "    time.sleep(.5)\n",
    "    prev = s"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}