{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "bd70534c",
"metadata": {},
"outputs": [],
"source": [
"from suno_utils.utils.notebook import _get_asr_model_ctc, capture_output, suppress_logging, Audio\n",
"from pyctcdecode import build_ctcdecoder\n",
"import portion\n",
"import re\n",
"import numpy as np\n",
"from nemo.collections.nlp.models import TokenClassificationModel\n",
"\n",
"def _logit_idx_to_time(idx, frame_dur_s=0.1):\n",
" return round(idx * frame_dur_s, 2)\n",
"\n",
"def _generate_audio_bleep(duration_s, freq=440.0, volume=0.4, sample_rate=16_000):\n",
" x = np.arange(int(duration_s * sample_rate))\n",
" arr = np.sin(2 * np.pi * freq * (x / sample_rate))\n",
" arr = (arr * np.iinfo(np.int16).max).clip(\n",
" np.iinfo(np.int16).min, np.iinfo(np.int16).max\n",
" ).astype(np.int16)\n",
" audio = Audio(arr.tobytes(), sample_rate, 2, 1)\n",
" return audio\n",
"\n",
"asr_model = _get_asr_model_ctc(\"cpu\")\n",
"decoder = build_ctcdecoder(asr_model.decoder.vocabulary)\n",
"with suppress_logging():\n",
" ner_model = TokenClassificationModel.from_pretrained(\"ner_en_bert\")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "96f0ee0c",
"metadata": {},
"outputs": [],
"source": [
"filepath = \"mikey.wav\"\n",
"\n",
"a = Audio.from_file(filepath)\n",
"with capture_output():\n",
" logits = asr_model.transcribe([filepath], logprobs=True)[0]\n",
"_, _, tokens, _, _ = decoder.decode_beams(logits)[0]\n",
"frame_dur_s = a.duration_s / logits.shape[0]\n",
"timed_tokens = []\n",
"for t, (st, et) in tokens:\n",
" timed_tokens.append((t, (_logit_idx_to_time(st, frame_dur_s=frame_dur_s), \n",
" _logit_idx_to_time(et, frame_dur_s=frame_dur_s))))\n",
"# find per\n",
"s = \" \".join([e[0] for e in timed_tokens])\n",
"with suppress_logging():\n",
" st = ner_model.add_predictions([s])[0]\n",
"per_tags = [\"per\" if \"PER]\" in e else \"\" for e in st.split()]\n",
"# find num\n",
"tmp_num_tags = [re.match(r\"^zero|one|two|three|four|five|six|seven|eight|nine|oh|o$\", t) for t in s.split()]\n",
"# keep only at least 2 consecutive\n",
"num_tags = [\"\"] * len(tmp_num_tags)\n",
"prev_m = False\n",
"for n, m in enumerate(tmp_num_tags):\n",
" if prev_m and m:\n",
" num_tags[n] = \"num\"\n",
" prev_m = m\n",
"prev_m = False\n",
"for n, m in enumerate(tmp_num_tags[::-1]):\n",
" if prev_m and m:\n",
" num_tags[len(tmp_num_tags) - n - 1] = \"num\"\n",
" prev_m = m\n",
"# merge\n",
"b_bleep_token = [len(a + b) > 0 for a, b in zip(per_tags, num_tags)]\n",
"# merge times of consecutive\n",
"bleep_timespans = portion.empty()\n",
"prev_bleep_end_s = -1\n",
"for (_, (start_s, end_s)), b_bleep in zip(timed_tokens, b_bleep_token):\n",
" if b_bleep:\n",
" if prev_bleep_end_s >= 0:\n",
" bleep_timespans |= portion.closedopen(prev_bleep_end_s, end_s)\n",
" else:\n",
" bleep_timespans |= portion.closedopen(start_s, end_s)\n",
" \n",
" if b_bleep:\n",
" prev_bleep_end_s = end_s\n",
" else:\n",
" prev_bleep_end_s = -1\n",
"# pad\n",
"pad_s = 0.02\n",
"for t in bleep_timespans:\n",
" bleep_timespans |= portion.closedopen(\n",
" round(np.max([t.lower - pad_s, 0]), 2), \n",
" round(np.min([t.upper + pad_s, a.duration_s]), 2)\n",
" )\n",
"# bleep\n",
"audios = []\n",
"offs_s = 0\n",
"for t in bleep_timespans:\n",
" if t.lower > 0:\n",
" audios.append(a.get_slice(offs_s, t.lower))\n",
" beep_duration_s = t.upper - t.lower\n",
" bleep = _generate_audio_bleep(beep_duration_s, sample_rate=a.sample_rate)\n",
" audios.append(bleep)\n",
" offs_s = t.upper\n",
"if offs_s < a.duration_s:\n",
" audios.append(a.get_slice(offs_s))\n",
"total_bytes = b\"\".join([e.bytes for e in audios])\n",
"bleeped_audio = Audio(total_bytes, audios[0].sample_rate, audios[0].byte_width, audios[0].n_channels)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "7c54db09",
"metadata": {},
"outputs": [],
"source": [
"bleeped_audio.write_wav(\"mikey_censored.wav\")"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "7fc0a2cc",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
" \n",
" "
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"bleeped_audio.play()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f8de4bbe",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "ea2d38a9",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "fb79c0a9",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "4db4cae8",
"metadata": {},
"source": [
"## Playground"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ed270cc0",
"metadata": {},
"outputs": [],
"source": [
"import nemo\n",
"import nemo.collections.nlp as nemo_nlp\n",
"from nemo.collections.nlp.models import TokenClassificationModel\n",
"from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer\n",
"from suno_utils.utils.notebook import Audio\n",
"\n",
"inverse_normalizer = InverseNormalizer(lang=\"en\", cache_dir=\"CACHE_DIR\")\n",
"punct_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(model_name=\"punctuation_en_bert\")\n",
"ner_model = TokenClassificationModel.from_pretrained(\"ner_en_bert\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1767925e",
"metadata": {},
"outputs": [],
"source": [
"s = a.transcribe()\n",
"s = inverse_normalizer.normalize(s)\n",
"s = punct_model.add_punctuation_capitalization([s])[0]\n",
"s = ner_model.add_predictions([s])[0]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e3c353de",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "a339b124",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "b7dca5e4",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "a92de179",
"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
}