{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "cc5fc7b1", "metadata": {}, "outputs": [], "source": [ "# https://github.com/suno-ai/glocke?nspiel/blob/main/suno_api/suno_api/asr/nemo_model.py\n", "# https://github.com/suno-ai/glockenspiel/blob/main/suno_api/suno_api/asr/nemo_worker.py" ] }, { "cell_type": "code", "execution_count": 2, "id": "a0863295", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"" ] }, { "cell_type": "code", "execution_count": 3, "id": "cccf4b4c", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[NeMo W 2022-10-21 14:49:04 optimizers:77] Could not import distributed_fused_adam optimizer from Apex\n" ] } ], "source": [ "import contextlib\n", "import time\n", "\n", "import tqdm\n", "import numpy as np\n", "import torch\n", "import nemo.collections.asr as nemo_asr\n", "from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer\n", "\n", "from suno_utils.audio import Audio\n", "from suno_utils.utils.display import suppress_logging\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", "# torch.backends.cudnn.benchmark = False\n", "# torch.backends.cudnn.enabled = False\n", "# torch.cuda.empty_cache()\n", "# gc.collect()\n", " \n", "audio = Audio.from_file(\"/home/georg/data/sample_audio/russia.wav\", sample_rate=16_000, byte_width=2)" ] }, { "cell_type": "code", "execution_count": 4, "id": "f33bb6f9", "metadata": {}, "outputs": [], "source": [ "# load nemo\n", "def load_model():\n", " with suppress_logging():\n", " model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(model_name=\"stt_en_conformer_ctc_large\")\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", " return model\n", "\n", "# https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/*\n", "# model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(model_name=\"stt_en_conformer_ctc_large\")\n", "# model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(model_name=\"stt_en_conformer_ctc_xlarge\")\n", "# model = nemo_asr.models.EncDecRNNTBPEModel.from_pretrained(model_name=\"stt_en_conformer_transducer_large\")\n", "# model = nemo_asr.models.EncDecRNNTBPEModel.from_pretrained(model_name=\"stt_en_conformer_transducer_xlarge\")\n", "\n", "model = load_model()\n", "# out = model.transcribe([\"/home/georg/data/sample_audio/russia.wav\"])[0]" ] }, { "cell_type": "code", "execution_count": 5, "id": "daa91d2a", "metadata": {}, "outputs": [], "source": [ "# # load riva\n", "# import yaml\n", "# import torch\n", "# from omegaconf.dictconfig import DictConfig\n", "# from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE\n", "# with suppress_logging():\n", "# model_dir = \"/home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0\"\n", "# weights_fp = os.path.join(model_dir, \"model_weights.ckpt\")\n", "# state_dict = torch.load(weights_fp, map_location=\"cpu\")\n", "# with open(os.path.join(model_dir, \"model_config.yaml\")) as f:\n", "# d = yaml.safe_load(f)\n", "# d[\"tokenizer\"] = {\n", "# \"dir\": model_dir, \n", "# \"type\": \"bpe\",\n", "# \"model_path\": os.path.join(model_dir, d[\"tokenizer\"][\"model_path\"]),\n", "# \"vocab_path\": os.path.join(model_dir, d[\"tokenizer\"][\"vocab_path\"]),\n", "# \"spe_tokenizer_vocab\": os.path.join(model_dir, d[\"tokenizer\"][\"spe_tokenizer_vocab\"]),\n", "# }\n", "# del d[\"train_ds\"]\n", "# d[\"validation_ds\"][\"manifest_filepath\"] = \"/home/georg/dummy.json\"\n", "# del d[\"test_ds\"]\n", "# riva_model = EncDecCTCModelBPE(cfg=DictConfig(d))\n", "# riva_model.load_state_dict(state_dict);\n", "# riva_model.preprocessor.featurizer.dither = 0.0\n", "# riva_model.preprocessor.featurizer.pad_to = 0\n", "# riva_model.eval()\n", "# riva_model.encoder.freeze()\n", "# riva_model.decoder.freeze()\n", "# # out = riva_model.transcribe([\"/home/georg/data/sample_audio/russia.wav\"])[0]" ] }, { "cell_type": "markdown", "id": "7dc053c1", "metadata": {}, "source": [ "### test speed" ] }, { "cell_type": "code", "execution_count": 80, "id": "bb466141", "metadata": {}, "outputs": [], "source": [ "# n_repl = 1\n", "n_repl = 16 * 4\n", "\n", "arr_10s = audio.array_float.reshape(1, -1)\n", "arr_10s_16x = np.array([audio.get_slice(to_s=10).array_float] * n_repl).reshape(n_repl, -1)\n", "a = torch.from_numpy(arr_10s_16x).to(model.device)\n", "b = torch.from_numpy(np.array([arr_10s.shape[1]] * n_repl)).to(model.device)\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", "execution_count": 81, "id": "120a9634", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 11.8 ms, sys: 0 ns, total: 11.8 ms\n", "Wall time: 11.1 ms\n" ] } ], "source": [ "%%time\n", "# with torch.backends.cudnn.flags(benchmark=False, deterministic=False):\n", "with torch.inference_mode():\n", " with autocast(): # this slows it down (at least when amp is not installed)\n", " with torch.no_grad():\n", " aa, bb = model.preprocessor(\n", " input_signal=a, length=b,\n", " )\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", "execution_count": 82, "id": "6b6da578", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 207 ms, sys: 38.7 ms, total: 246 ms\n", "Wall time: 245 ms\n" ] } ], "source": [ "%%time\n", "# with torch.backends.cudnn.flags(benchmark=False, deterministic=False):\n", "with torch.inference_mode():\n", " with autocast(): # this slows it down (at least when amp is not installed)\n", " with torch.no_grad():\n", " out = model.forward(processed_signal=aa, processed_signal_length=bb)\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", "execution_count": 83, "id": "bbf08dcf", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 210 ms, sys: 42.8 ms, total: 253 ms\n", "Wall time: 252 ms\n" ] } ], "source": [ "%%time\n", "# with torch.backends.cudnn.flags(benchmark=False, deterministic=False):\n", "with torch.inference_mode():\n", " with autocast(): # this slows it down (at least when amp is not installed)\n", " with torch.no_grad():\n", " out = model.forward(input_signal=a, input_signal_length=b)\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", "execution_count": 84, "id": "314e656e", "metadata": {}, "outputs": [], "source": [ "del a, b, aa, bb, out\n", "torch.cuda.synchronize()" ] }, { "cell_type": "code", "execution_count": null, "id": "2351ca99", "metadata": {}, "outputs": [], "source": [ "# lambda\n", "# - 10s clips\n", "# 500x realtime\n", "# 3200x realtime when batched\n", "# - 2s clips\n", "# 100x realtime\n", "# 3200x realtime when batched\n", "\n", "# note: doing 2s preds with 8s context - equivalent of 12000x realtime for 10s\n", "\n", "# lambda cpu\n", "# - 10s clips\n", "# 50x realtime\n", "# 60x realtime when batched\n", "# - 2s clips\n", "# 30x realtime\n", "# 100x realtime when batched\n", "\n", "# cloud machine (g4dn.2xlarge)\n", "# - 10s clips, max batch size 256/128 (with/without autocast)\n", "# 300x realtime\n", "# 1600x realtime when batched\n", "# - 2s clips\n", "# 65x realtime\n", "# 900x realtime when batched" ] }, { "cell_type": "markdown", "id": "3740b607", "metadata": {}, "source": [ "### correct\n", "President Joe Biden, backed by the full symbolic power of the Western Alliance, is locked in a showdown with Russian President Vladimir Putin, who is using Ukraine as a hostage to try to force the US to renegotiate the settled outcome of the Cold War.\n", "\n", "### nemo final\n", "president joe biden backed by the full symbolic power of the western alliance is locked in a showdown with russian president vladimir putin who is using ukraine as a hostage to try to force the uus to renegotiate the settled outcome of the cold war\n", "\n", "### riva final\n", "president joe biden backed by the full symbolic power of the western alliance is locked in a showdown with russian president of vladimir putin who is using ukraine as a hostage to try to force the us to renegotiate the settled outcome of the cold war\n", "\n", "### rev final\n", "President Joe Biden backed by the full symbolic power of the Western Alliance is locked in a showdown with Russian president of Vladimir Putin, who is using Ukraine as a hostage to try to force the us to renegotiate the settled outcome of the cult war.\n", "\n", "### rev partial\n", "president joe biden backed by the full symbolic power of the western alliance is locked in a showdown with russian president of vladimir putin who is using ukraine as a hostage to try to force the us to renegotiate the settled outcome of the cult\n", "\n", "### deepgram final\n", "President Joe Biden backed by the full symbolic power of the West an alliance is locked in a showdown with Russian president of Vladimir Putin who is using Ukraine as a hostage to try to force the US to renegotiate the settled outcome of the Cold War.\n", "\n", "### deepgram partial\n", "President Joe Biden backed by the full symbolic power of the West an alliance (because he's) is locked in a showdown with Russian president of Vladimir Putin who's using (new) Ukraine as a hostage to try to force the US to (read) renegotiate the settled outcome of the Cold War.\n", "\n", "### whisper large\n", "President Joe Biden, backed by the full symbolic power of the Western Alliance, is locked in a showdown with Russian President Vladimir Putin, who is using Ukraine as a hostage to try to force the US to renegotiate the settled outcome of the Cold War.\n", "\n", "### whisper medium.en\n", "President Joe Biden, backed by the full symbolic power of the Western Alliance, is locked in a showdown with Russian President Vladimir Putin, who is using Ukraine as a hostage to try to force the US to renegotiate the settled outcome of the Cold War." ] }, { "cell_type": "code", "execution_count": null, "id": "9e7e3837", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "e6f05786", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "929bcabc", "metadata": {}, "source": [ "## Manual sliding window" ] }, { "cell_type": "code", "execution_count": 255, "id": "e3540ab0", "metadata": {}, "outputs": [], "source": [ "vocab = {n: c for n, c in enumerate(model._cfg[\"decoder\"][\"vocabulary\"] + [\"\"])}\n", "\n", "def _decode_tokens(probs):\n", " return [vocab.get(idx) for idx in probs.argmax(axis=1)]\n", " \n", "# def _find_pause_streaks(token_list, min_duration=2):\n", "# # find space/blank stretches or blank 3+ stretches at end\n", "# # get all words before and remember where we stopped (backoff 2 from right boundary)\n", "# streaks = []\n", "# n_start = None\n", "# trailing_bound = False\n", "# for n, token in enumerate(token_list):\n", "# if token == \"\":\n", "# if n_start is None:\n", "# n_start = n\n", "# continue\n", "# elif token == \"▁\":\n", "# if n_start is None:\n", "# n_start = n\n", "# trailing_bound = True\n", "# elif token.startswith(\"▁\"):\n", "# if n_start is not None and n - n_start > 0:\n", "# streaks.append((n_start, n))\n", "# else:\n", "# streaks.append((n, n))\n", "# trailing_bound = False\n", "# else:\n", "# if trailing_bound and n_start is not None and n - n_start > 0:\n", "# streaks.append((n_start, n))\n", "# trailing_bound = False\n", "# n_start = None\n", "# if n_start is not None and len(token_list) - n_start > 0:\n", "# streaks.append((n_start, len(token_list)))\n", "# # filter to reliable pauses if possible\n", "# filtered_streaks = [(s, e) for s, e in streaks if s < len(token_list) - 15]\n", "# if len(filtered_streaks) == 0:\n", "# filtered_streaks = streaks[:]\n", "# streaks = filtered_streaks[:]\n", "# filtered_streaks = [(s, e) for s, e in streaks if e - s >= min_duration]\n", "# if len(filtered_streaks) == 0:\n", "# filtered_streaks = streaks[:]\n", "# streaks = filtered_streaks[:]\n", "# # find good breakpoint\n", "# if len(streaks) == 0:\n", "# best_break_idx = len(token_list)\n", "# else:\n", "# best_start, best_end = streaks[-1]\n", "# best_break_idx = int(np.max([best_start, best_end - 1]))\n", "# return best_break_idx\n", " \n", "def _find_pause_streaks_2(token_list):\n", " best_idx = None\n", " for n, t in enumerate(token_list[::-1]):\n", " if t.startswith(\"▁\"):\n", " best_idx = len(token_list) - n - 1\n", " if best_idx is None:\n", " best_idx = len(token_list) - 5\n", " return best_idx\n", " \n", "def _decode(token_list):\n", " tokens = []\n", " p = None\n", " for t in token_list:\n", " if t != p:\n", " tokens.append(t)\n", " p = t\n", " return \"\".join(tokens).replace(\"▁\", \" \").strip()\n", " \n", "# def greedy_decode(arr):\n", "# deduped = []\n", "# p = None\n", "# for idx in arr:\n", "# if idx != p:\n", "# deduped.append(idx)\n", "# p = idx\n", "# return \"\".join([vocab.get(idx) for idx in deduped]).replace(\"▁\", \" \").strip()" ] }, { "cell_type": "code", "execution_count": 251, "id": "75a84027", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:00<00:00, 40.64it/s]\n" ] } ], "source": [ "# every 500 ms\n", "jump_duration_ms = 1_000\n", "context_duration_s = 5\n", "\n", "jump_dist = int(jump_duration_ms / 1_000 * 16_000)\n", "context_dist = int(context_duration_s * 16_000)\n", "min_pred_ms = 10\n", "min_pred_dist = int(min_pred_ms / 1_000 * 16_000)\n", "arr = audio.array_float.reshape(1, -1)\n", "probs_container = []\n", "for n in tqdm.tqdm(range(int(np.ceil(arr.shape[-1] // jump_dist)))):\n", " end_idx = (n + 1) * jump_dist\n", " start_idx = np.max([0, end_idx - context_dist])\n", " use_arr = arr[:,start_idx:end_idx]\n", " with torch.inference_mode():\n", " with autocast(): # this slows it down (at least when amp is not installed)\n", " with torch.no_grad():\n", " a = torch.from_numpy(use_arr).to(model.device)\n", " b = torch.from_numpy(np.array([use_arr.shape[1]])).to(model.device)\n", " raw_logits, _, _ = model.forward(input_signal=a, input_signal_length=b)\n", " logits = raw_logits.detach().cpu().numpy().squeeze()\n", " del a, b, raw_logits\n", " probs = np.exp(logits)\n", " probs_container.append(probs)\n", " if arr.shape[-1] - end_idx < min_pred_dist:\n", " break" ] }, { "cell_type": "code", "execution_count": 262, "id": "c8973470", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "president joe\n", "biden backed\n", "by the full symbolic\n", "power of the western\n", "alliance is locked\n", "in a showdown with\n", "the russian president of vladimir\n", "putin who\n", "is using ukraine\n", "as a hostage\n", "to try to force\n", "the uus to renegotiate\n", "the settled\n", "outcome of the cold war\n" ] } ], "source": [ "# find right-most gap of at least 3 frames with p_blank >= 0.8\n", "n_jump_logits = int(round(jump_duration_ms * 0.026))\n", "prev_idx = 0\n", "last_token = None\n", "wiggle_room = 2\n", "for n_probs, probs in enumerate(probs_container):\n", " token_list = _decode_tokens(probs)\n", " if n_probs == len(probs_container) - 1:\n", " text = _decode(token_list[prev_idx:])\n", " print(text)\n", " break\n", "# idx = _find_pause_streaks(token_list[-n_jump_logits:]) - n_jump_logits\n", " idx = _find_pause_streaks_2(token_list[-n_jump_logits:]) - n_jump_logits\n", " # wiggle token if position doesn't fit\n", " if token_list[prev_idx-1] != last_token:\n", " for n in range(wiggle_room):\n", " if token_list[prev_idx-1+n] == last_token:\n", " prev_idx = prev_idx + n\n", " break\n", " if token_list[prev_idx-1-n] == last_token:\n", " prev_idx = prev_idx - n\n", " break\n", " last_token = token_list[idx-1]\n", " text = _decode(token_list[prev_idx:idx])\n", " print(text)\n", " prev_idx = idx - n_jump_logits" ] }, { "cell_type": "code", "execution_count": null, "id": "374fb4c5", "metadata": {}, "outputs": [], "source": [ "# TODO: go down to every 500ms\n", "# TODO: improve break function \n", "# either a token back like now, or at least out 10 frames back if there is blankish there?\n", "# TODO: implement realtime loop image version in notebook\n", "# TODO: use vad to avoid partial word errors at end?\n", "# TODO: batching using _speech_collate_fn\n", "# TODO: consider what to do about normalization: model.cfg.preprocessor.normalize\n", "# TODO: test causal model\n", "# TODO: test transducer\n", "# TODO: implement cache aware for speed" ] }, { "cell_type": "code", "execution_count": 111, "id": "bc5c3cf9", "metadata": {}, "outputs": [], "source": [ "def _speech_collate_fn(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, 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" ] }, { "cell_type": "code", "execution_count": null, "id": "75f6b379", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "33d3cb92", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "64641c16", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "4be9e432", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "a99752c3", "metadata": {}, "source": [ "## Try streaming" ] }, { "cell_type": "code", "execution_count": 46, "id": "63f430ec", "metadata": {}, "outputs": [], "source": [ "vocab = {n: c for n, c in enumerate(model._cfg[\"decoder\"][\"vocabulary\"] + [\"\"])}\n", "\n", "def _decode_arr(idx_list):\n", " tokens = []\n", " p = None\n", " for idx in idx_list:\n", " if idx != p:\n", " tokens.append(vocab.get(idx, \"\"))\n", " p = idx\n", " return \"\".join(tokens).replace(\"▁\", \" \").strip()" ] }, { "cell_type": "code", "execution_count": null, "id": "51b4c1c0", "metadata": {}, "outputs": [], "source": [ "# reload model\n", "# model = load_model()" ] }, { "cell_type": "code", "execution_count": null, "id": "fe782c67", "metadata": {}, "outputs": [], "source": [ "# returned 10 instead of 13 frames, no cache" ] }, { "cell_type": "code", "execution_count": 55, "id": "b1cdc5ae", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "CacheAwareStreamingConfig(chunk_size=[797, 800], shift_size=[197, 200], cache_drop_size=150, last_channel_cache_size=400, valid_out_len=50, pre_encode_cache_size=[0, 5], drop_extra_pre_encoded=2, last_channel_num=18, last_time_num=18)" ] }, "execution_count": 55, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# 1 input frame is 10ms\n", "# 1 output logit is 38.46ms\n", "model.encoder.setup_streaming_params(\n", " chunk_size=200, left_chunks=2, shift_size=50, max_context=10_000,\n", ")\n", "# model.encoder.streaming_cfg.valid_out_len = 11\n", "model.encoder.streaming_cfg" ] }, { "cell_type": "code", "execution_count": 107, "id": "797b458a", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 42.87it/s]\n" ] } ], "source": [ "# every ~500 ms\n", "jump_duration_ms = 500\n", "window_duration_ms = jump_duration_ms\n", "extra_context_duration_s = 4\n", "\n", "def _calc_drop_extra_pre_encoded(asr_model, step_num):\n", " return 0 if step_num == 0 else model.encoder.streaming_cfg.drop_extra_pre_encoded\n", "\n", "jump_dist = int(jump_duration_ms / 1_000 * 16_000)\n", "window_dist = int(window_duration_ms / 1_000 * 16_000)\n", "min_pred_ms = 10\n", "min_pred_dist = int(min_pred_ms / 1_000 * 16_000)\n", "arr = audio.array_float.reshape(1, -1)\n", "out_container = []\n", "n_tot = int(np.ceil(arr.shape[-1] // jump_dist))\n", "pred_out_stream = None\n", "transcribed_texts = None\n", "cache_last_channel = None\n", "cache_last_time = None\n", "for n in tqdm.tqdm(range(n_tot)):\n", " is_last_step = (n == n_tot - 1)\n", " end_idx = (n + 1) * jump_dist\n", " start_idx = np.max([0, end_idx - window_dist])\n", " use_arr = arr[:,start_idx:end_idx]\n", " with torch.inference_mode():\n", " with autocast(): # this slows it down (at least when amp is not installed)\n", " with torch.no_grad():\n", " a = torch.from_numpy(use_arr).to(model.device)\n", " b = torch.from_numpy(np.array([use_arr.shape[1]])).to(model.device)\n", " aa, bb = model.preprocessor(\n", " input_signal=a, length=b,\n", " )\n", "# raw_logits, _, _ = model.forward(input_signal=a, input_signal_length=b)\n", " (\n", " pred_out_stream, # argmax label array\n", " _, # transcribed text\n", " cache_last_channel,\n", " cache_last_time,\n", " _,\n", " ) = model.conformer_stream_step(\n", " processed_signal=aa,\n", " processed_signal_length=bb,\n", " cache_last_channel=cache_last_channel,\n", " cache_last_time=cache_last_time,\n", " keep_all_outputs=is_last_step, # should be true in last step\n", " previous_pred_out=pred_out_stream,\n", " drop_extra_pre_encoded=_calc_drop_extra_pre_encoded(model, n),\n", " return_transcription=False,\n", " )\n", "# if n == 1:\n", "# break\n", " del a, b, aa, bb\n", " # TODO: use self.encoder.cache_aware_stream_step directly to get log probs\n", " # see: nemo_asr.parts.mixin.mixins.py\n", " out = pred_out_stream[0].detach().cpu().numpy().squeeze().tolist()\n", " out_container.append(out)\n", " if arr.shape[-1] - end_idx < min_pred_dist:\n", " break" ] }, { "cell_type": "code", "execution_count": 108, "id": "7d0ac752", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "'president go bye back by the full symbolic power the western alliance his lot walck in the show down with russian pres it in the blood mir put mm was using ukraine as a whole hostage to try out of force the us read and go no sh the sttleld outcome of the cold war'" ] }, "execution_count": 108, "metadata": {}, "output_type": "execute_result" } ], "source": [ "_decode_arr(out_container[-1])" ] }, { "cell_type": "code", "execution_count": null, "id": "2ff43f43", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "5fabba89", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "8adc9fb6", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "c73fcd96", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "3bc27346", "metadata": {}, "source": [ "## Playground" ] }, { "cell_type": "code", "execution_count": 5, "id": "6695329f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "CacheAwareStreamingConfig(chunk_size=[77, 80], shift_size=[37, 40], cache_drop_size=10, last_channel_cache_size=1000, valid_out_len=10, pre_encode_cache_size=[0, 5], drop_extra_pre_encoded=2, last_channel_num=18, last_time_num=18)" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# 1 frame is 10ms\n", "# 1 size in chunk is 40ms\n", "model.encoder.setup_streaming_params(\n", " chunk_size=20, left_chunks=50, shift_size=10, max_context=10_000,\n", ")\n", "# model.encoder.streaming_cfg.valid_out_len = 11\n", "model.encoder.streaming_cfg" ] }, { "cell_type": "code", "execution_count": 10, "id": "005ed3d2", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[NeMo I 2022-09-30 18:17:00 features:225] PADDING: 0\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "[NeMo W 2022-09-30 18:17:00 nemo_logging:349] /home/georg/code/NeMo/nemo/collections/asr/parts/utils/streaming_utils.py:1384: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:172.)\n", " audio_signal = torch.from_numpy(audio).unsqueeze_(0).to(device)\n", " \n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "\n", "presi\n", "president jo\n", "president joe bid\n", "president joe biden bac\n", "president joe biden backed by the\n", "president joe biden backed by the full sy\n", "president joe biden backed by the full symbolic\n", "president joe biden backed by the full symbolic power of\n", "president joe biden backed by the full symbolic power of the western\n", "president joe biden backed by the full symbolic power of the westernn all\n", "president joe biden backed by the full symbolic power of the westernn alliance\n", "president joe biden backed by the full symbolic power of the westernn alliance is lock\n", "done\n" ] } ], "source": [ "streaming_buffer = CacheAwareStreamingAudioBuffer(\n", " model=model, \n", " online_normalization=False\n", ")\n", "_ = streaming_buffer.append_audio(audio.array)\n", "# processed_signal, processed_signal_length = streaming_buffer.get_all_audios()\n", "\n", "batch_size = len(streaming_buffer.streams_length)\n", "cache_last_channel, cache_last_time = model.encoder.get_initial_cache_state(batch_size=batch_size)\n", "\n", "def calc_drop_extra_pre_encoded(asr_model, step_num):\n", " # for the first step there is no need to drop any tokens after the downsampling as no caching is being used\n", " if step_num == 0:\n", " return 0\n", " else:\n", " return model.encoder.streaming_cfg.drop_extra_pre_encoded\n", "\n", "previous_hypotheses = None\n", "streaming_buffer_iter = iter(streaming_buffer)\n", "pred_out_stream = None\n", "for step_num, (chunk_audio, chunk_lengths) in enumerate(streaming_buffer_iter):\n", " with torch.inference_mode():\n", " with autocast():\n", " # keep_all_outputs needs to be True for the last step of streaming when model is trained with \n", " # att_context_style=regular otherwise the last outputs would get dropped\n", " with torch.no_grad():\n", "# print(chunk_audio.size(-1))\n", " t0 = time.time()\n", " (\n", " pred_out_stream,\n", " transcribed_texts,\n", " cache_last_channel,\n", " cache_last_time,\n", " _,\n", " ) = model.conformer_stream_step(\n", " processed_signal=chunk_audio,\n", " processed_signal_length=chunk_lengths,\n", " cache_last_channel=cache_last_channel,\n", " cache_last_time=cache_last_time,\n", " keep_all_outputs=streaming_buffer.is_buffer_empty(),\n", " previous_pred_out=pred_out_stream,\n", " drop_extra_pre_encoded=calc_drop_extra_pre_encoded(model, step_num),\n", " return_transcription=True,\n", " )\n", " t1 = time.time()\n", "# print(\"pred took {}ms\".format(int(round((t1 - t0) * 1000))))\n", " print(transcribed_texts[0])\n", " if len(transcribed_texts[0]) >= 80:\n", " print(\"done\")\n", " break\n", "# print(pred_out_stream[0].shape)" ] }, { "cell_type": "code", "execution_count": null, "id": "42322fc6", "metadata": {}, "outputs": [], "source": [ "# chunk_size=100, shift_size=100\n", "# torch.Size([1, 80, 397])\n", "# torch.Size([1, 80, 405])\n", "# torch.Size([1, 80, 405])\n", "# torch.Size([1, 80, 405])\n", "# torch.Size([1, 80, 15])\n", "\n", "# chunk_size=200, shift_size=100\n", "# torch.Size([1, 80, 797])\n", "# torch.Size([1, 80, 805])\n", "# torch.Size([1, 80, 805])\n", "# torch.Size([1, 80, 415])\n", "# torch.Size([1, 80, 15])\n", "\n", "# chunk_size=200, shift_size=200\n", "# torch.Size([1, 80, 797])\n", "# torch.Size([1, 80, 805])\n", "# torch.Size([1, 80, 15])\n", "\n", "# chunk_size=1000, shift_size=1000\n", "# torch.Size([1, 80, 1607])" ] }, { "cell_type": "code", "execution_count": 8, "id": "b82b59d7", "metadata": {}, "outputs": [], "source": [ "online_normalization = True\n", "streaming_buffer = CacheAwareStreamingAudioBuffer(model=model, online_normalization=online_normalization)\n", "processed_signal, processed_signal_length, stream_id = streaming_buffer.append_audio(audio.array)\n", "processed_signal, processed_signal_length = streaming_buffer.get_all_audios()\n", "streaming_buffer.reset_buffer()" ] }, { "cell_type": "code", "execution_count": 121, "id": "ef65e3b3", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "['president joe biden back 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 s to renegotiate the settled outcome of the cold war']" ] }, "execution_count": 121, "metadata": {}, "output_type": "execute_result" } ], "source": [ "batch_size = 1\n", "cache_last_channel, cache_last_time = model.encoder.get_initial_cache_state(batch_size=batch_size)\n", "with torch.inference_mode():\n", " with autocast():\n", " # keep_all_outputs needs to be True for the last step of streaming when model is trained with \n", " # att_context_style=regular otherwise the last outputs would get dropped\n", " with torch.no_grad():\n", " (\n", " pred_out_stream,\n", " transcribed_texts,\n", " cache_last_channel,\n", " cache_last_time,\n", " previous_hypotheses,\n", " ) = model.conformer_stream_step(\n", " processed_signal=processed_signal,\n", " processed_signal_length=processed_signal_length,\n", " cache_last_channel=cache_last_channel,\n", " cache_last_time=cache_last_time,\n", " keep_all_outputs=True,\n", " previous_hypotheses=None,\n", " previous_pred_out=None,\n", " drop_extra_pre_encoded=0,\n", " return_transcription=True,\n", " )\n", "transcribed_texts" ] }, { "cell_type": "code", "execution_count": 10, "id": "f05683c8", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[NeMo I 2022-09-30 18:17:00 features:225] PADDING: 0\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "[NeMo W 2022-09-30 18:17:00 nemo_logging:349] /home/georg/code/NeMo/nemo/collections/asr/parts/utils/streaming_utils.py:1384: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:172.)\n", " audio_signal = torch.from_numpy(audio).unsqueeze_(0).to(device)\n", " \n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "\n", "\n", "presi\n", "president jo\n", "president joe bid\n", "president joe biden bac\n", "president joe biden backed by the\n", "president joe biden backed by the full sy\n", "president joe biden backed by the full symbolic\n", "president joe biden backed by the full symbolic power of\n", "president joe biden backed by the full symbolic power of the western\n", "president joe biden backed by the full symbolic power of the westernn all\n", "president joe biden backed by the full symbolic power of the westernn alliance\n", "president joe biden backed by the full symbolic power of the westernn alliance is lock\n", "done\n" ] } ], "source": [ "streaming_buffer = CacheAwareStreamingAudioBuffer(\n", " model=model, \n", " online_normalization=False\n", ")\n", "_ = streaming_buffer.append_audio(audio.array)\n", "# processed_signal, processed_signal_length = streaming_buffer.get_all_audios()\n", "\n", "batch_size = len(streaming_buffer.streams_length)\n", "cache_last_channel, cache_last_time = model.encoder.get_initial_cache_state(batch_size=batch_size)\n", "\n", "def calc_drop_extra_pre_encoded(asr_model, step_num):\n", " # for the first step there is no need to drop any tokens after the downsampling as no caching is being used\n", " if step_num == 0:\n", " return 0\n", " else:\n", " return model.encoder.streaming_cfg.drop_extra_pre_encoded\n", "\n", "previous_hypotheses = None\n", "streaming_buffer_iter = iter(streaming_buffer)\n", "pred_out_stream = None\n", "for step_num, (chunk_audio, chunk_lengths) in enumerate(streaming_buffer_iter):\n", " with torch.inference_mode():\n", " with autocast():\n", " # keep_all_outputs needs to be True for the last step of streaming when model is trained with \n", " # att_context_style=regular otherwise the last outputs would get dropped\n", " with torch.no_grad():\n", "# print(chunk_audio.size(-1))\n", " t0 = time.time()\n", " (\n", " pred_out_stream,\n", " transcribed_texts,\n", " cache_last_channel,\n", " cache_last_time,\n", " _,\n", " ) = model.conformer_stream_step(\n", " processed_signal=chunk_audio,\n", " processed_signal_length=chunk_lengths,\n", " cache_last_channel=cache_last_channel,\n", " cache_last_time=cache_last_time,\n", " keep_all_outputs=streaming_buffer.is_buffer_empty(),\n", " previous_pred_out=pred_out_stream,\n", " drop_extra_pre_encoded=calc_drop_extra_pre_encoded(model, step_num),\n", " return_transcription=True,\n", " )\n", " t1 = time.time()\n", "# print(\"pred took {}ms\".format(int(round((t1 - t0) * 1000))))\n", " print(transcribed_texts[0])\n", " if len(transcribed_texts[0]) >= 80:\n", " print(\"done\")\n", " break\n", "# print(pred_out_stream[0].shape)" ] }, { "cell_type": "code", "execution_count": 51, "id": "c9f67904", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[NeMo I 2022-10-20 16:37:51 features:225] PADDING: 0\n", "[NeMo I 2022-10-20 16:37:51 features:225] PADDING: 0\n" ] } ], "source": [ "n_repl = 16*4\n", "\n", "arr_1 = np.array(audio.get_slice(to_s=8).array_float)\n", "arr_2 = np.array(audio.get_slice(from_s=8, to_s=10).array_float)\n", "\n", "streaming_buffer = CacheAwareStreamingAudioBuffer(\n", " model=model, \n", " online_normalization=False\n", ")\n", "for _ in range(n_repl):\n", " _ = streaming_buffer.append_audio(arr_1)\n", "processed_signal, processed_signal_length = streaming_buffer.get_all_audios()\n", "\n", "streaming_buffer = CacheAwareStreamingAudioBuffer(\n", " model=model, \n", " online_normalization=False\n", ")\n", "for _ in range(n_repl):\n", " _ = streaming_buffer.append_audio(arr_2)\n", "processed_signal_2, processed_signal_length_2 = streaming_buffer.get_all_audios()" ] }, { "cell_type": "code", "execution_count": 52, "id": "7ae2fd22", "metadata": {}, "outputs": [], "source": [ "with torch.inference_mode():\n", " with autocast(): # this slows it down (at least when amp is not installed)\n", " with torch.no_grad():\n", " (\n", " pred_out_stream,\n", " transcribed_texts,\n", " cache_last_channel,\n", " cache_last_time,\n", " _,\n", " ) = model.conformer_stream_step(\n", " processed_signal=processed_signal,\n", " processed_signal_length=processed_signal_length,\n", " cache_last_channel=None,\n", " cache_last_time=None,\n", " keep_all_outputs=False,\n", " previous_pred_out=None,\n", " drop_extra_pre_encoded=None,\n", " return_transcription=False,\n", " )" ] }, { "cell_type": "code", "execution_count": null, "id": "26583fce", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "784584a1", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "6b2a3679", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "0b31131c", "metadata": {}, "outputs": [], "source": [] }, { "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": "052ebb1f", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "5c92290a", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "f2f10880", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "3c55ac8f", "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 }