{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "a0863295",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
]
},
{
"cell_type": "code",
"execution_count": 35,
"id": "cccf4b4c",
"metadata": {},
"outputs": [],
"source": [
"import contextlib\n",
"import time\n",
"import random\n",
"\n",
"import tqdm\n",
"import numpy as np\n",
"import torch\n",
"import nemo.collections.asr as nemo_asr\n",
"\n",
"from suno_utils.audio import Audio\n",
"from suno_utils.utils.display import suppress_logging\n",
"from suno_utils.utils.text import normalize_whitespace\n",
"from suno_utils.ctcdecode.decoder import build_ctcdecoder\n",
"\n",
"if (\n",
" torch.cuda.is_available()\n",
" and hasattr(torch.cuda, 'amp')\n",
" and hasattr(torch.cuda.amp, 'autocast')\n",
"):\n",
" autocast = torch.cuda.amp.autocast\n",
"else:\n",
" @contextlib.contextmanager\n",
" def autocast():\n",
" yield\n",
" \n",
"SAMPLE_RATE = 16_000\n",
" \n",
"def gpu_stats(clear_cache=True):\n",
" if clear_cache:\n",
" torch.cuda.empty_cache()\n",
" from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo\n",
" t = torch.cuda.get_device_properties(0).total_memory\n",
" r = torch.cuda.memory_reserved(0)\n",
" a = torch.cuda.memory_allocated(0)\n",
" print(\"torch:\")\n",
" print(round(t / 1e9, 1), \"Gb total\")\n",
" print(round(r / 1e9, 1), \"Gb reserved\")\n",
" print(round(a / 1e9, 1), \"Gb allocated\")\n",
" print()\n",
" nvmlInit()\n",
" h = nvmlDeviceGetHandleByIndex(0)\n",
" info = nvmlDeviceGetMemoryInfo(h)\n",
" print(\"nvidia:\")\n",
" print(round(info.total / 1e9, 1), \"Gb total\")\n",
" print(round(info.used / 1e9, 1), \"Gb used\")\n",
" print(round(info.free / 1e9, 1), \"Gb free\") \n",
" \n",
"def load_model(model_name=\"stt_en_conformer_ctc_large\"):\n",
" with suppress_logging():\n",
" model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(model_name=model_name)\n",
" model.preprocessor.featurizer.dither = 0.0\n",
" model.preprocessor.featurizer.pad_to = 0\n",
" model.eval()\n",
" model.encoder.freeze()\n",
" model.decoder.freeze()\n",
" torch.cuda.synchronize()\n",
" torch.cuda.empty_cache()\n",
" return model\n",
"\n",
"def get_tokens_from_logits(model, logits):\n",
" vocab = {n: c for n, c in enumerate(list(model.decoder.vocabulary) + [\"\"])}\n",
" probs = np.exp(logits)\n",
" argmax_preds = probs.argmax(axis=1)\n",
" tokens = [vocab[idx] for idx in argmax_preds]\n",
" return tokens\n",
"\n",
"def get_text_from_tokens(tokens):\n",
" squashed_tokens = []\n",
" prev_token = None\n",
" for token in tokens:\n",
" if token != prev_token:\n",
" squashed_tokens.append(token)\n",
" prev_token = token\n",
" text = \"\".join(squashed_tokens).replace(\"▁\", \" \")\n",
" text = normalize_whitespace(text)\n",
" return text\n",
"\n",
"def decode_logits(model, logits):\n",
" tokens = get_tokens_from_logits(model, logits)\n",
" text = get_text_from_tokens(tokens)\n",
" return text\n",
"\n",
"def _collate_features(batch, pad_id=0):\n",
" \"\"\"collate batch of audio sig, audio len\"\"\"\n",
" packed_batch = list(zip(*batch))\n",
" if len(packed_batch) == 2:\n",
" _, audio_lengths = packed_batch\n",
" else:\n",
" raise ValueError(\"Expects 2 tensors in the batch!\")\n",
" max_audio_len = 0\n",
" has_audio = audio_lengths[0] is not None\n",
" if not has_audio:\n",
" return None, None\n",
" max_audio_len = max(audio_lengths).item()\n",
" audio_signal = []\n",
" for sig, sig_len in batch:\n",
" sig_len = sig_len.item()\n",
" if sig_len < max_audio_len:\n",
" pad = (0, int(max_audio_len - sig_len))\n",
" sig = torch.nn.functional.pad(sig, pad)\n",
" audio_signal.append(sig)\n",
" audio_signal = torch.stack(audio_signal)\n",
" audio_lengths = torch.stack(audio_lengths)\n",
" return audio_signal, audio_lengths\n",
"\n",
"def _get_features(model, audio_signal, audio_lengths):\n",
" processed_signal, processed_signal_length = model.preprocessor(\n",
" input_signal=audio_signal, length=audio_lengths,\n",
" )\n",
" return processed_signal, processed_signal_length\n",
"\n",
"def _infer(model, processed_signal, processed_signal_length):\n",
" logits, logits_len, labels = model.forward(\n",
" processed_signal=processed_signal, processed_signal_length=processed_signal_length,\n",
" )\n",
" return logits, logits_len, labels\n",
"\n",
"MIN_PREDICT_DURATION_MS = 10 # this is to avoid model errors\n",
"MAX_PREDICT_DURATION_MS = 30_100 # this is to avoid gpu oom\n",
"\n",
"def predict_logits(model, audio_arr_list):\n",
" if len(audio_arr_list) == 0:\n",
" return []\n",
" # alert if contains an array that is too short or too long\n",
" if any([\n",
" (\n",
" arr.shape[0] < int(MIN_PREDICT_DURATION_MS / 1_000 * SAMPLE_RATE) or \n",
" arr.shape[0] > int(MAX_PREDICT_DURATION_MS / 1_000 * SAMPLE_RATE)\n",
" ) for arr in audio_arr_list\n",
" ]):\n",
" raise ValueError(\"array size error\")\n",
" with torch.inference_mode(), torch.no_grad(), autocast():\n",
" arr_list_gpu = [torch.from_numpy(arr).to(model.device) for arr in audio_arr_list]\n",
" arr_len_list_gpu = [torch.tensor(arr.shape[-1]).to(model.device) for arr in audio_arr_list]\n",
" batch = list(zip(arr_list_gpu, arr_len_list_gpu))\n",
" audio_signal, audio_lengths = _collate_features(batch)\n",
" processed_signal, processed_signal_length = _get_features(model, audio_signal, audio_lengths)\n",
" logits, logits_len, labels = _infer(model, processed_signal, processed_signal_length)\n",
" logits_list = [l[:idx].detach().cpu().numpy().squeeze() for idx, l in zip(logits_len, logits)]\n",
" # labels_list = [l[:idx].detach().cpu().numpy().squeeze() for idx, l in zip(logits_len, labels)]\n",
" del (\n",
" arr_list_gpu, arr_len_list_gpu, batch, audio_signal, audio_lengths, \n",
" processed_signal, processed_signal_length, logits, logits_len, labels\n",
" )\n",
" torch.cuda.synchronize()\n",
" return logits_list\n",
"\n",
"def _find_break_idx(tokens):\n",
" best_idx = None\n",
" # break at last new word\n",
" for n, t in enumerate(tokens[::-1]):\n",
" if t.startswith(\"▁\"):\n",
" best_idx = len(tokens) - n - 1\n",
" break\n",
" # if nothing found then break at an early blank\n",
" for n, t in enumerate(tokens[-5:]):\n",
" if t == \"\":\n",
" best_idx = len(tokens) + n - 5\n",
" break\n",
" # if nothing found then hard break\n",
" if best_idx is None:\n",
" best_idx = max(0, len(tokens) - 5)\n",
" return best_idx\n",
"\n",
"def _find_break_idx_from_guess(tokens, idx_guess, decoded_words=None):\n",
" best_idx = idx_guess\n",
" all_break_idx = [n for n, t in enumerate(tokens) if t.startswith(\"▁\")]\n",
" if len(all_break_idx) == 0:\n",
" return best_idx\n",
" # find break index that is closest to use as a best guess\n",
" best_idx = all_break_idx[np.argsort([np.abs(idx - best_idx) for idx in all_break_idx])[0]]\n",
" # refine index in word space\n",
" if decoded_words is None:\n",
" decoded_words = []\n",
" if len(decoded_words) == 0:\n",
" best_idx = 0\n",
" return best_idx\n",
" # user words to finetune selection\n",
" # if we already did a good job then return\n",
" discarded_words = get_text_from_tokens(tokens[:best_idx]).split()\n",
" if len(discarded_words) > 0 and discarded_words[-1] == decoded_words[-1]:\n",
" return best_idx\n",
" # check if there is word overlap (up to 2)\n",
" later_break_idx = [idx for idx in all_break_idx if idx > best_idx][:2]\n",
" for break_idx in later_break_idx[::-1]:\n",
" extra_words = get_text_from_tokens(tokens[best_idx:break_idx]).split()\n",
" if extra_words == decoded_words[-len(extra_words):]:\n",
" return break_idx\n",
" # check earlier break index incase we missed words (up to 2)\n",
" earlier_break_idx = [idx for idx in all_break_idx if idx < best_idx][-2:]\n",
" for break_idx in earlier_break_idx[::-1]:\n",
" discarded_words = get_text_from_tokens(tokens[:break_idx]).split()\n",
" if discarded_words == decoded_words[-len(discarded_words):]:\n",
" return break_idx\n",
" return best_idx"
]
},
{
"cell_type": "code",
"execution_count": 28,
"id": "2af4ba6d",
"metadata": {},
"outputs": [],
"source": [
"# notebook visualization stuff\n",
"import html\n",
"from IPython.display import display, HTML, Javascript\n",
"\n",
"def display_boxes(n_users):\n",
" for n_user in range(n_users):\n",
" display(HTML(\"User \" + str(n_user) + \":\"))\n",
" \n",
"def display_text(n_user, text):\n",
" display(Javascript(\n",
" \"var el = document.getElementsByClassName('asr_user_\" + str(n_user) + \"');\"\n",
" \"for (var i = 0; i < el.length; ++i) {el[i].innerHTML = '\" + html.escape(text) + \"';}\"\n",
" ))"
]
},
{
"cell_type": "code",
"execution_count": 29,
"id": "ffd965dd",
"metadata": {},
"outputs": [],
"source": [
"model = load_model()\n",
"decoder = build_ctcdecoder(model.decoder.vocabulary)\n",
"audio = Audio.from_file(\"/home/georg/data/sample_audio/russia.wav\", sample_rate=16_000, byte_width=2)"
]
},
{
"cell_type": "code",
"execution_count": 30,
"id": "84bf3809",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 108 ms, sys: 25.3 ms, total: 133 ms\n",
"Wall time: 133 ms\n"
]
}
],
"source": [
"%%time\n",
"audio_arr_list = [audio.array_float] * 16\n",
"out = predict_logits(model, audio_arr_list)\n",
"# torch.cuda.synchronize()\n",
"# torch.cuda.empty_cache()\n",
"# this should take ~130ms"
]
},
{
"cell_type": "code",
"execution_count": 33,
"id": "5a79a6b9",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 1.14 s, sys: 181 ms, total: 1.32 s\n",
"Wall time: 1.32 s\n"
]
}
],
"source": [
"%%time\n",
"for _ in range(10):\n",
" audio_arr_list = [audio.array_float] * 16\n",
" out = predict_logits(model, audio_arr_list)\n",
"# should be ~1.4s"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "034aee4c",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "110b1ce9",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "1764caab",
"metadata": {},
"source": [
"## Set up data"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "431b8e1a",
"metadata": {},
"outputs": [],
"source": [
"# this type of array list could come from e.g. multiple concurrent users/sessions\n",
"# USER_AUDIO_ARRAYS = [\n",
"# audio.array_float, \n",
"# audio.array_float[:100000],\n",
"# ]\n",
"\n",
"# get a few random ones\n",
"n_items = 6\n",
"d = \"/mnt/data-ssd-1/data/librispeech/LibriSpeech/test-clean-processed/\"\n",
"fns = os.listdir(d) \n",
"random.seed(6006)\n",
"random.shuffle(fns)\n",
"USER_AUDIO_ARRAYS = [\n",
" Audio.from_file(os.path.join(d, fn), sample_rate=16_000, byte_width=2).array_float\n",
" for fn in fns[:n_items]\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "8c5768a5",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'you are my all the world and i must strive to know my shames and praises from your tongue none else to me nor i to none alive that my steeled sense or changes right or wrong'"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"logits_list = predict_logits(model, USER_AUDIO_ARRAYS)\n",
"decode_logits(model, logits_list[0])"
]
},
{
"cell_type": "markdown",
"id": "929bcabc",
"metadata": {},
"source": [
"## Sliding window inference (single user)"
]
},
{
"cell_type": "code",
"execution_count": 364,
"id": "572f62e0",
"metadata": {
"scrolled": false
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"president joe biden backed by the full symbolic power of the western alliance is locked in a showdown with the russian president of vladimir putin who is using ukraine as a hostage to try to force the u s to renegotiate the settled outcome of the cold war\r"
]
}
],
"source": [
"# we are pretending that user upload the audio in USER_AUDIO_ARRAYS\n",
"# they do so continuously such that we can realtime process chunks of STEP_DURATION_MS\n",
"# end of stream is known here but needs to be communicated separately in a real scenario\n",
"STEP_DURATION_MS = 500\n",
"MAX_CONTEXT_DURATION_S = 5\n",
"RESPECT_REALTIME = True\n",
"\n",
"USER_AUDIO_ARRAY = USER_AUDIO_ARRAYS[0]\n",
"\n",
"# caluclate some basics\n",
"step_n_array = int(STEP_DURATION_MS / 1_000 * SAMPLE_RATE)\n",
"n_array_per_logit = 8000 / 13 # TODO: is there a better estimate of this?\n",
"max_context_n_array = int(MAX_CONTEXT_DURATION_S * SAMPLE_RATE)\n",
"\n",
"# store user info\n",
"user_words = []\n",
"user_prev_audio_start_idx = 0\n",
"user_prev_audio_end_idx = 0\n",
"user_prev_token_end_idx = 0\n",
"\n",
"n_step = 0\n",
"t0 = time.time()\n",
"while True:\n",
" # sleep if necessary to simulate realtime\n",
" if RESPECT_REALTIME:\n",
" sleep_duration_s = STEP_DURATION_MS / 1_000 * (n_step + 1) - (time.time() - t0)\n",
" if sleep_duration_s > 0.01:\n",
" time.sleep(sleep_duration_s)\n",
" # get audio chunk to predict\n",
" audio_end_idx = min(USER_AUDIO_ARRAY.shape[0], (n_step + 1) * step_n_array)\n",
" audio_start_idx = max(0, audio_end_idx - max_context_n_array)\n",
" audio_array_segment = USER_AUDIO_ARRAY[audio_start_idx:audio_end_idx]\n",
" # we are done if audio too small for prediction\n",
" if audio_array_segment.shape[0] < int(MIN_PREDICT_DURATION_MS / 1_000 * SAMPLE_RATE):\n",
" break\n",
" # do prediction\n",
" logits = predict_logits(model, [audio_array_segment])[0]\n",
" tokens = get_tokens_from_logits(model, logits)\n",
" # get best guess for token_start_idx based on what has already been decoded\n",
" token_start_idx = user_prev_token_end_idx - int(round(\n",
" (audio_start_idx - user_prev_audio_start_idx) / n_array_per_logit\n",
" ))\n",
" token_start_idx = _find_break_idx_from_guess(tokens, token_start_idx, decoded_words=user_words)\n",
" # end-of-stream signal implicitly give through end of array\n",
" if audio_end_idx == USER_AUDIO_ARRAY.shape[0]:\n",
" user_words.extend(get_text_from_tokens(tokens[token_start_idx:]).split())\n",
" break\n",
" # find reliable token end index\n",
" token_end_idx = token_start_idx + _find_break_idx(tokens[token_start_idx:])\n",
" # add tokens to user stack\n",
" user_words.extend(get_text_from_tokens(tokens[token_start_idx:token_end_idx]).split()) \n",
" # display results for user\n",
" output_text = \" \".join([\n",
" w \n",
" for w in user_words + get_text_from_tokens(tokens[token_end_idx:]).split() \n",
" if len(w) > 0\n",
" ])\n",
" if len(output_text) == 0:\n",
" output_text = \" \"\n",
" print(output_text, end=\"\\r\")\n",
" # prepare for next step\n",
" user_prev_audio_start_idx = audio_start_idx\n",
" user_prev_audio_end_idx = audio_end_idx\n",
" user_prev_token_end_idx = token_end_idx\n",
" n_step += 1\n",
"# TODO: improve the above with 'partial' vs 'final' results"
]
},
{
"cell_type": "markdown",
"id": "b4758f4a",
"metadata": {},
"source": [
"## Sliding window inference (multi user)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac4f1a48",
"metadata": {},
"outputs": [],
"source": [
"# we are pretending that user upload the audio in USER_AUDIO_ARRAYS\n",
"# they do so continuously such that we can realtime process chunks of STEP_DURATION_MS\n",
"# end of stream is known here but needs to be communicated separately in a real scenario\n",
"STEP_DURATION_MS = 500\n",
"MAX_CONTEXT_DURATION_S = 5\n",
"RESPECT_REALTIME = True\n",
"\n",
"# caluclate some basics\n",
"step_n_array = int(STEP_DURATION_MS / 1_000 * SAMPLE_RATE)\n",
"n_array_per_logit = 8000 / 13 # TODO: is there a better estimate of this?\n",
"max_context_n_array = int(MAX_CONTEXT_DURATION_S * SAMPLE_RATE)\n",
"\n",
"# store user info\n",
"user_data = [\n",
" {\n",
" \"words\": [],\n",
" \"prev_audio_start_idx\": 0,\n",
" \"prev_audio_end_idx\": 0,\n",
" \"prev_token_end_idx\": 0,\n",
" \"is_done\": False\n",
" } for _ in range(len(USER_AUDIO_ARRAYS))\n",
"]\n",
"\n",
"display_boxes(len(USER_AUDIO_ARRAYS))\n",
"\n",
"n_step = 0\n",
"t0 = time.time()\n",
"while True:\n",
" # sleep if necessary to simulate realtime\n",
" if RESPECT_REALTIME:\n",
" sleep_duration_s = STEP_DURATION_MS / 1_000 * (n_step + 1) - (time.time() - t0)\n",
" if sleep_duration_s > 0.01:\n",
" time.sleep(sleep_duration_s)\n",
" # get audio chunks to predict in this batch\n",
" # TODO: limit this to a certain batchsize and do triaging for which users to serve\n",
" batch = []\n",
" for user_idx, audio_array in enumerate(USER_AUDIO_ARRAYS):\n",
" if user_data[user_idx][\"is_done\"]:\n",
" continue\n",
" audio_end_idx = min(audio_array.shape[0], (n_step + 1) * step_n_array)\n",
" # end-of-stream signal implicitly give through end of array\n",
" is_final_pred = audio_end_idx == audio_array.shape[0]\n",
" audio_start_idx = min(audio_array.shape[0], max(0, n_step * step_n_array - max_context_n_array))\n",
" audio_array_segment = audio_array[audio_start_idx:audio_end_idx]\n",
" # we are done if audio too small for prediction\n",
" if audio_array_segment.shape[0] < int(MIN_PREDICT_DURATION_MS / 1_000 * SAMPLE_RATE):\n",
" continue\n",
" batch.append({\n",
" \"audio_array\": audio_array_segment,\n",
" \"user_idx\": user_idx,\n",
" \"audio_start_idx\": audio_start_idx,\n",
" \"audio_end_idx\": audio_end_idx,\n",
" \"is_final_pred\": is_final_pred,\n",
" })\n",
" if len(batch) == 0:\n",
" break\n",
" # do prediction\n",
" logits_list = predict_logits(model, [m[\"audio_array\"] for m in batch])\n",
" for m, logits in zip(batch, logits_list):\n",
" user_idx = m[\"user_idx\"]\n",
" audio_start_idx = m[\"audio_start_idx\"]\n",
" audio_end_idx = m[\"audio_end_idx\"]\n",
" is_final_pred = m[\"is_final_pred\"]\n",
" prev_audio_start_idx = user_data[user_idx][\"prev_audio_start_idx\"]\n",
" prev_token_end_idx = user_data[user_idx][\"prev_token_end_idx\"]\n",
" # decode tokens\n",
" tokens = get_tokens_from_logits(model, logits)\n",
" # get best guess for token_start_idx based on what has already been decoded\n",
" token_start_idx = prev_token_end_idx - int(round(\n",
" (audio_start_idx - prev_audio_start_idx) / n_array_per_logit\n",
" ))\n",
" token_start_idx = _find_break_idx_from_guess(\n",
" tokens, token_start_idx, \n",
" decoded_words=user_data[user_idx][\"words\"],\n",
" )\n",
" # end-of-stream signal implicitly give through end of array\n",
" if is_final_pred:\n",
" user_data[user_idx][\"words\"].extend(get_text_from_tokens(tokens[token_start_idx:]).split())\n",
" user_data[user_idx][\"is_done\"] = True\n",
" continue\n",
" # find reliable token end index\n",
" token_end_idx = token_start_idx + _find_break_idx(tokens[token_start_idx:])\n",
" # add tokens to user stack\n",
" user_data[user_idx][\"words\"].extend(get_text_from_tokens(tokens[token_start_idx:token_end_idx]).split()) \n",
" # prepare for next step\n",
" user_data[user_idx][\"prev_audio_start_idx\"] = audio_start_idx\n",
" user_data[user_idx][\"prev_audio_end_idx\"] = audio_end_idx\n",
" user_data[user_idx][\"prev_token_end_idx\"] = token_end_idx\n",
" # display results for user\n",
" output_text = \" \".join([\n",
" w \n",
" for w in user_data[user_idx][\"words\"] + get_text_from_tokens(tokens[token_end_idx:]).split() \n",
" if len(w) > 0\n",
" ])\n",
" display_text(user_idx, output_text)\n",
" n_step += 1"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b2c0b475",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "784da302",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "9af7bc0a",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "3fa8f8c2",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "fb348fce",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "d35edee1",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "8862c3a4",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "4279a35c",
"metadata": {},
"source": [
"## TODO:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cb1363ea",
"metadata": {},
"outputs": [],
"source": [
"# -- short term --\n",
"# TODO: reduce context to 2-3s?\n",
"# TODO: implement 'partial' and 'final' (probably final until one back and next one becomes partial)\n",
"# TODO: add punctuation for final\n",
"# TODO: add pyctcdecode for final (beam 5 and hotwords, no LM?)\n",
"# TODO: hotwords format (warn if not normalized?)\n",
"\n",
"# -- medium term --\n",
"# TODO: look at max gpu memory consumption\n",
"# https://stackoverflow.com/questions/58216000/get-total-amount-of-free-gpu-memory-and-available-using-pytorch\n",
"# TODO: add diarization for final\n",
"# TODO: how do we deal with 2-channel, esp diarize\n",
"# TODO: implement async batch transcript\n",
"\n",
"# -- long term --\n",
"# TODO: use vad to avoid partial word errors at end?\n",
"# TODO: consider what to do about normalization: model.cfg.preprocessor.normalize\n",
"# TODO: ask nvidia for causal conformer, squeezeformer and LM\n",
"# TODO: test transducer\n",
"# TODO: implement cache aware for speed"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "33e6051e",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "b4c88258",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "d6f6cc14",
"metadata": {},
"source": [
"## test accuracy"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "a3aeade8",
"metadata": {},
"outputs": [],
"source": [
"import re\n",
"import json\n",
"import funcy\n",
"\n",
"from suno_utils.audio import Tokens\n",
"from suno_utils.utils.metrics import get_wer_bulk, get_cer_bulk\n",
"\n",
"data_dir = \"/mnt/data-ssd-1/data/private/customer/sanas/2022-08-04-fili-callcenter/to_sanas/2022-08-12\"\n",
"# data_dir = \"/mnt/data-ssd-1/data/private/customer/sanas/2022-08-04-fili-callcenter/to_sanas/2022-08-15\"\n",
"\n",
"dev_meta = []\n",
"with open(os.path.join(data_dir, \"segment_metadata.jsonl\")) as f:\n",
" for line in f:\n",
" line = line.strip()\n",
" if len(line) == 0:\n",
" continue\n",
" m = json.loads(line)\n",
" # only allowed meta tag is [laughter] and hesitations need to be removed\n",
" text = m[\"transcript\"][\"text\"]\n",
" text = re.sub(r\"\\s*\\-\\-\\s*\", \" \", text)\n",
" text = re.sub(r\"\\s*\\[laughter\\]\\s*\", \" \", text)\n",
" text = normalize_whitespace(text)\n",
" if \"[\" in text:\n",
" continue\n",
" dev_meta.append({\n",
" \"duration_s\": m[\"duration_s\"],\n",
" \"filepath\": os.path.join(data_dir, m[\"uri\"]),\n",
" \"text\": text,\n",
" \"text_norm\": Tokens.from_dict(m[\"transcript_normalized\"][\"tokens\"]).plaintext,\n",
" })\n",
"print(len(dev_meta), \"files loaded\")\n",
"print(round(np.sum([m[\"duration_s\"] for m in dev_meta]) / 60 / 60, 1), \"hours\")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "ecb7eb53",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 942/942 [02:59<00:00, 5.24it/s]\n"
]
}
],
"source": [
"greedy_preds = []\n",
"for chunk_meta in tqdm.tqdm(funcy.chunks(16, dev_meta), total=int(np.ceil(len(dev_meta) / 16))):\n",
" audio_arr_list = [Audio.from_file(m[\"filepath\"]).array_float for m in chunk_meta]\n",
" logits_list = predict_logits(model, audio_arr_list)\n",
" greedy_preds.extend([decode_logits(model, logits) for logits in logits_list])"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "2912dfb4",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"20.6% WER\n"
]
}
],
"source": [
"wer_val = get_wer_bulk(\n",
" [m[\"text_norm\"] for m in dev_meta],\n",
" greedy_preds,\n",
")\n",
"print(\"{}% WER\".format(round(wer_val * 100, 1)))"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "0ab940a8",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 942/942 [02:07<00:00, 7.40it/s]\n"
]
}
],
"source": [
"decoder_preds = []\n",
"for chunk_meta in tqdm.tqdm(funcy.chunks(16, dev_meta), total=int(np.ceil(len(dev_meta) / 16))):\n",
" audio_arr_list = [Audio.from_file(m[\"filepath\"]).array_float for m in chunk_meta]\n",
" logits_list = predict_logits(model, audio_arr_list)\n",
" decoder_preds.extend([\n",
" decoder.decode(\n",
" logits,\n",
" beam_width=5,\n",
" beam_prune_logp=-10.0,\n",
" token_min_logp=-5.0,\n",
" ) for logits in logits_list\n",
" ])"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "1a5aeb0d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"20.5% WER for decoder preds\n"
]
}
],
"source": [
"wer_val = get_wer_bulk(\n",
" [m[\"text_norm\"] for m in dev_meta],\n",
" decoder_preds,\n",
")\n",
"print(\"{}% WER for decoder preds\".format(round(wer_val * 100, 1)))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "741486ef",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "a7699038",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "3c6f3495",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "6b0bc9e4",
"metadata": {},
"source": [
"## Decoders"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e6ce11fe",
"metadata": {},
"outputs": [],
"source": [
"# https://pytorch.org/audio/main/tutorials/asr_inference_with_ctc_decoder_tutorial.html\n",
"from torchaudio.models.decoder import ctc_decoder\n",
"torch_decoder = ctc_decoder(\n",
" lexicon=lexicon_file,\n",
" tokens=tokens_file,\n",
"# lm=kenlm_file,\n",
" beam_size=beam_width,\n",
" beam_threshold=10,\n",
")\n",
"\n",
"decoder = build_ctcdecoder(model.decoder.vocabulary)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "05513faf",
"metadata": {},
"outputs": [],
"source": [
"\" \".join(torch_decoder(logits)[0][0].words)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7636a875",
"metadata": {},
"outputs": [],
"source": [
"# pyctcdecoe\n",
"decoder.decode(\n",
" logits,\n",
" beam_width=5,\n",
" beam_prune_logp=-10.0,\n",
" token_min_logp=-5.0,\n",
") for logits in logits_list\n",
"\n",
"\n",
"# maxtasksperchild=10 in Pool if we instantiate outside\n",
"# with multiprocessing.get_context(\"fork\").Pool(10) as pool:\n",
"# _ = decoder.decode_beams_batch(pool, logits_list[:128], beam_width=50)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "00f1fff8",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "4a58a45b",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "c7c13780",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "6c29be9e",
"metadata": {},
"source": [
"## cap/punct"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "df08134c",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[NeMo W 2022-10-24 17:30:00 optimizers:67] Could not import distributed_fused_adam optimizer from Apex\n",
"[NeMo W 2022-10-24 17:30:02 experimental:27] Module is experimental, not ready for production and is not fully supported. Use at your own risk.\n",
"[NeMo W 2022-10-24 17:30:03 experimental:27] Module is experimental, not ready for production and is not fully supported. Use at your own risk.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"100% [......................................................................] 245117658 / 245117658"
]
}
],
"source": [
"import nemo\n",
"import nemo.collections.nlp as nemo_nlp\n",
"from suno_utils.utils.display import suppress_logging\n",
"from IPython.utils.io import capture_output\n",
"\n",
"with suppress_logging():\n",
" punct_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(\n",
" model_name=\"punctuation_en_bert\",\n",
" )\n",
" punct_2_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(\n",
" model_name=\"punctuation_en_distilbert\",\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": 35,
"id": "806e91f7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CPU times: user 733 ms, sys: 0 ns, total: 733 ms\n",
"Wall time: 732 ms\n"
]
}
],
"source": [
"%%time\n",
"with suppress_logging():\n",
" with capture_output():\n",
" for _ in range(10):\n",
" out = punct_model.add_punctuation_capitalization(\n",
" ['how are you i recently came across this interesting place oh really very cool '*2]*16\n",
" )\n",
"# ~12 ms per pred, ~3 with batch 16, ~1.5 with batch 64"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bd9129f4",
"metadata": {},
"outputs": [],
"source": [
"import nemo\n",
"import nemo.collections.nlp as nemo_nlp\n",
"from suno_utils.utils.display import suppress_logging\n",
"from IPython.utils.io import capture_output\n",
"\n",
"with suppress_logging():\n",
" punct_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(\n",
" model_name=\"punctuation_en_bert\",\n",
" )\n",
" punct_2_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(\n",
" model_name=\"punctuation_en_distilbert\",\n",
" )\n",
"\n",
"%%time\n",
"with suppress_logging():\n",
" with capture_output():\n",
" for _ in range(10):\n",
" out = punct_model.add_punctuation_capitalization(\n",
" ['how are you i recently came across this interesting place oh really very cool']*16\n",
" )\n",
"# ~12 ms per pred, ~3 with batch 16, ~1.5 with batch 64"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "29672670",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "fdae0519",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "2735f6d0",
"metadata": {},
"source": [
"## Denormalize"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "072ae551",
"metadata": {},
"outputs": [],
"source": [
"from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer\n",
"\n",
"with suppress_logging():\n",
" itn_model = InverseNormalizer(lang=\"en\")\n",
"# itn_2_model = nemo_nlp.models.ThutmoseTaggerModel.from_pretrained(model_name=\"itn_en_thutmose_bert\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e4f7138b",
"metadata": {},
"outputs": [],
"source": [
"spoken = \"we paid fifteen dollars for this desk from a t and t i think cause i work at the f b i\"\n",
"print(itn_model.inverse_normalize(spoken, verbose=False))\n",
"# print(itn_2_model._infer([spoken])[0].split(\"\\t\")[0])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1eae4639",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "d1ea137f",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "4be9e432",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "3bc27346",
"metadata": {},
"source": [
"## Playground"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b28a55a8",
"metadata": {},
"outputs": [],
"source": [
"docker run --gpus all \\\n",
" -it \\\n",
" -v /home/georg/code/NeMo:/NeMo \\\n",
" -v /home/georg:/home/georg \\\n",
" -p 8339:8339 \\\n",
" --shm-size=8g \\\n",
" --ulimit memlock=-1 \\\n",
" --ulimit stack=67108864 \\\n",
" nvcr.io/nvidia/pytorch:22.09-py3\n",
" \n",
"# -v /mnt/data-ssd-1:/mnt/data-ssd-1 \\"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "79dd68db",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "4f1fd993",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "f94eaa3b",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "04069c03",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}