Upload 13 files

whisper_diarization_main/.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.srt
+*.txt
+*.wav
+*.mp3
+*.m4a
+

whisper_diarization_main/LICENSE

@@ -0,0 +1,24 @@
+BSD 2-Clause License
+
+Copyright (c) 2023, Mahmoud Ashraf
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

whisper_diarization_main/README.md

@@ -0,0 +1,93 @@
+<h1 align="center">Speaker Diarization Using OpenAI Whisper</h1>
+<p align="center">
+  <a href="https://github.com/MahmoudAshraf97/whisper-diarization/stargazers">
+    <img src="https://img.shields.io/github/stars/MahmoudAshraf97/whisper-diarization.svg?colorA=orange&colorB=orange&logo=github"
+         alt="GitHub stars">
+  </a>
+  <a href="https://github.com/MahmoudAshraf97/whisper-diarization/issues">
+        <img src="https://img.shields.io/github/issues/MahmoudAshraf97/whisper-diarization.svg"
+             alt="GitHub issues">
+  </a>
+  <a href="https://github.com/MahmoudAshraf97/whisper-diarization/blob/master/LICENSE">
+        <img src="https://img.shields.io/github/license/MahmoudAshraf97/whisper-diarization.svg"
+             alt="GitHub license">
+  </a>
+  <a href="https://twitter.com/intent/tweet?text=&url=https%3A%2F%2Fgithub.com%2FMahmoudAshraf97%2Fwhisper-diarization">
+  <img src="https://img.shields.io/twitter/url/https/github.com/MahmoudAshraf97/whisper-diarization.svg?style=social" alt="Twitter">
+  </a> 
+  </a>
+  <a href="https://colab.research.google.com/github/MahmoudAshraf97/whisper-diarization/blob/main/Whisper_Transcription_%2B_NeMo_Diarization.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab">
+  </a>
+ 
+</p>
+
+
+# 
+Speaker Diarization pipeline based on OpenAI Whisper
+I'd like to thank [@m-bain](https://github.com/m-bain) for Batched Whisper Inference, [@mu4farooqi](https://github.com/mu4farooqi) for punctuation realignment algorithm
+
+<img src="https://github.blog/wp-content/uploads/2020/09/github-stars-logo_Color.png" alt="drawing" width="25"/> **Please, star the project on github (see top-right corner) if you appreciate my contribution to the community!**
+
+## What is it
+This repository combines Whisper ASR capabilities with Voice Activity Detection (VAD) and Speaker Embedding to identify the speaker for each sentence in the transcription generated by Whisper. First, the vocals are extracted from the audio to increase the speaker embedding accuracy, then the transcription is generated using Whisper, then the timestamps are corrected and aligned using WhisperX to help minimize diarization error due to time shift. The audio is then passed into MarbleNet for VAD and segmentation to exclude silences, TitaNet is then used to extract speaker embeddings to identify the speaker for each segment, the result is then associated with the timestamps generated by WhisperX to detect the speaker for each word based on timestamps and then realigned using punctuation models to compensate for minor time shifts.
+
+
+WhisperX and NeMo parameters are coded into diarize.py and helpers.py, I will add the CLI arguments to change them later
+## Installation
+`PyTorch`, `FFMPEG` and `Cython` are needed as prerequisites to install the requirements
+```
+pip install cython torch
+```
+or
+```
+pip install torch
+sudo apt update && sudo apt install cython3
+```
+```
+# on Ubuntu or Debian
+sudo apt update && sudo apt install ffmpeg
+
+# on Arch Linux
+sudo pacman -S ffmpeg
+
+# on MacOS using Homebrew (https://brew.sh/)
+brew install ffmpeg
+
+# on Windows using Chocolatey (https://chocolatey.org/)
+choco install ffmpeg
+
+# on Windows using Scoop (https://scoop.sh/)
+scoop install ffmpeg
+```
+```
+pip install -r requirements.txt
+```
+## Usage 
+
+```
+python diarize.py -a AUDIO_FILE_NAME
+```
+
+If your system has enough VRAM (>=10GB), you can use `diarize_parallel.py` instead, the difference is that it runs NeMo in parallel with Whisper, this can be beneficial in some cases and the result is the same since the two models are nondependent on each other. This is still experimental, so expect errors and sharp edges. Your feedback is welcome.
+
+## Command Line Options
+
+- `-a AUDIO_FILE_NAME`: The name of the audio file to be processed
+- `--no-stem`: Disables source separation
+- `--whisper-model`: The model to be used for ASR, default is `medium.en`
+- `--suppress_numerals`: Transcribes numbers in their pronounced letters instead of digits, improves alignment accuracy
+- `--device`: Choose which device to use, defaults to "cuda" if available
+- `--language`: Manually select language, useful if language detection failed
+- `--batch-size`: Batch size for batched inference, reduce if you run out of memory, set to 0 for non-batched inference
+
+## Known Limitations
+- Overlapping speakers are yet to be addressed, a possible approach would be to separate the audio file and isolate only one speaker, then feed it into the pipeline but this will need much more computation
+- There might be some errors, please raise an issue if you encounter any.
+
+## Future Improvements
+- Implement a maximum length per sentence for SRT
+
+## Acknowledgements
+Special Thanks for [@adamjonas](https://github.com/adamjonas) for supporting this project
+This work is based on [OpenAI's Whisper](https://github.com/openai/whisper) , [Faster Whisper](https://github.com/guillaumekln/faster-whisper) , [Nvidia NeMo](https://github.com/NVIDIA/NeMo) , and [Facebook's Demucs](https://github.com/facebookresearch/demucs)

whisper_diarization_main/Whisper_Transcription_+_NeMo_Diarization.ipynb

@@ -0,0 +1,1000 @@
+{
+  "cells": [
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/MahmoudAshraf97/whisper-diarization/blob/main/Whisper_Transcription_%2B_NeMo_Diarization.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eCmjcOc9yEtQ"
+      },
+      "source": [
+        "# Installing Dependencies"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Tn1c-CoDv2kw"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install git+https://github.com/m-bain/whisperX.git@78dcfaab51005aa703ee21375f81ed31bc248560\n",
+        "!pip install --no-build-isolation nemo_toolkit[asr]==1.23.0\n",
+        "!pip install --no-deps git+https://github.com/facebookresearch/demucs#egg=demucs\n",
+        "!pip install git+https://github.com/oliverguhr/deepmultilingualpunctuation.git\n",
+        "!pip install git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YzhncHP0ytbQ"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "import wget\n",
+        "from omegaconf import OmegaConf\n",
+        "import json\n",
+        "import shutil\n",
+        "import torch\n",
+        "import torchaudio\n",
+        "from nemo.collections.asr.models.msdd_models import NeuralDiarizer\n",
+        "from deepmultilingualpunctuation import PunctuationModel\n",
+        "import re\n",
+        "import logging\n",
+        "import nltk\n",
+        "from whisperx.utils import LANGUAGES, TO_LANGUAGE_CODE\n",
+        "from ctc_forced_aligner import (\n",
+        "    load_alignment_model,\n",
+        "    generate_emissions,\n",
+        "    preprocess_text,\n",
+        "    get_alignments,\n",
+        "    get_spans,\n",
+        "    postprocess_results,\n",
+        ")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jbsUt3SwyhjD"
+      },
+      "source": [
+        "# Helper Functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Se6Hc7CZygxu"
+      },
+      "outputs": [],
+      "source": [
+        "punct_model_langs = [\n",
+        "    \"en\",\n",
+        "    \"fr\",\n",
+        "    \"de\",\n",
+        "    \"es\",\n",
+        "    \"it\",\n",
+        "    \"nl\",\n",
+        "    \"pt\",\n",
+        "    \"bg\",\n",
+        "    \"pl\",\n",
+        "    \"cs\",\n",
+        "    \"sk\",\n",
+        "    \"sl\",\n",
+        "]\n",
+        "langs_to_iso = {\n",
+        "    \"af\": \"afr\",\n",
+        "    \"am\": \"amh\",\n",
+        "    \"ar\": \"ara\",\n",
+        "    \"as\": \"asm\",\n",
+        "    \"az\": \"aze\",\n",
+        "    \"ba\": \"bak\",\n",
+        "    \"be\": \"bel\",\n",
+        "    \"bg\": \"bul\",\n",
+        "    \"bn\": \"ben\",\n",
+        "    \"bo\": \"tib\",\n",
+        "    \"br\": \"bre\",\n",
+        "    \"bs\": \"bos\",\n",
+        "    \"ca\": \"cat\",\n",
+        "    \"cs\": \"cze\",\n",
+        "    \"cy\": \"wel\",\n",
+        "    \"da\": \"dan\",\n",
+        "    \"de\": \"ger\",\n",
+        "    \"el\": \"gre\",\n",
+        "    \"en\": \"eng\",\n",
+        "    \"es\": \"spa\",\n",
+        "    \"et\": \"est\",\n",
+        "    \"eu\": \"baq\",\n",
+        "    \"fa\": \"per\",\n",
+        "    \"fi\": \"fin\",\n",
+        "    \"fo\": \"fao\",\n",
+        "    \"fr\": \"fre\",\n",
+        "    \"gl\": \"glg\",\n",
+        "    \"gu\": \"guj\",\n",
+        "    \"ha\": \"hau\",\n",
+        "    \"haw\": \"haw\",\n",
+        "    \"he\": \"heb\",\n",
+        "    \"hi\": \"hin\",\n",
+        "    \"hr\": \"hrv\",\n",
+        "    \"ht\": \"hat\",\n",
+        "    \"hu\": \"hun\",\n",
+        "    \"hy\": \"arm\",\n",
+        "    \"id\": \"ind\",\n",
+        "    \"is\": \"ice\",\n",
+        "    \"it\": \"ita\",\n",
+        "    \"ja\": \"jpn\",\n",
+        "    \"jw\": \"jav\",\n",
+        "    \"ka\": \"geo\",\n",
+        "    \"kk\": \"kaz\",\n",
+        "    \"km\": \"khm\",\n",
+        "    \"kn\": \"kan\",\n",
+        "    \"ko\": \"kor\",\n",
+        "    \"la\": \"lat\",\n",
+        "    \"lb\": \"ltz\",\n",
+        "    \"ln\": \"lin\",\n",
+        "    \"lo\": \"lao\",\n",
+        "    \"lt\": \"lit\",\n",
+        "    \"lv\": \"lav\",\n",
+        "    \"mg\": \"mlg\",\n",
+        "    \"mi\": \"mao\",\n",
+        "    \"mk\": \"mac\",\n",
+        "    \"ml\": \"mal\",\n",
+        "    \"mn\": \"mon\",\n",
+        "    \"mr\": \"mar\",\n",
+        "    \"ms\": \"may\",\n",
+        "    \"mt\": \"mlt\",\n",
+        "    \"my\": \"bur\",\n",
+        "    \"ne\": \"nep\",\n",
+        "    \"nl\": \"dut\",\n",
+        "    \"nn\": \"nno\",\n",
+        "    \"no\": \"nor\",\n",
+        "    \"oc\": \"oci\",\n",
+        "    \"pa\": \"pan\",\n",
+        "    \"pl\": \"pol\",\n",
+        "    \"ps\": \"pus\",\n",
+        "    \"pt\": \"por\",\n",
+        "    \"ro\": \"rum\",\n",
+        "    \"ru\": \"rus\",\n",
+        "    \"sa\": \"san\",\n",
+        "    \"sd\": \"snd\",\n",
+        "    \"si\": \"sin\",\n",
+        "    \"sk\": \"slo\",\n",
+        "    \"sl\": \"slv\",\n",
+        "    \"sn\": \"sna\",\n",
+        "    \"so\": \"som\",\n",
+        "    \"sq\": \"alb\",\n",
+        "    \"sr\": \"srp\",\n",
+        "    \"su\": \"sun\",\n",
+        "    \"sv\": \"swe\",\n",
+        "    \"sw\": \"swa\",\n",
+        "    \"ta\": \"tam\",\n",
+        "    \"te\": \"tel\",\n",
+        "    \"tg\": \"tgk\",\n",
+        "    \"th\": \"tha\",\n",
+        "    \"tk\": \"tuk\",\n",
+        "    \"tl\": \"tgl\",\n",
+        "    \"tr\": \"tur\",\n",
+        "    \"tt\": \"tat\",\n",
+        "    \"uk\": \"ukr\",\n",
+        "    \"ur\": \"urd\",\n",
+        "    \"uz\": \"uzb\",\n",
+        "    \"vi\": \"vie\",\n",
+        "    \"yi\": \"yid\",\n",
+        "    \"yo\": \"yor\",\n",
+        "    \"yue\": \"yue\",\n",
+        "    \"zh\": \"chi\",\n",
+        "}\n",
+        "\n",
+        "\n",
+        "whisper_langs = sorted(LANGUAGES.keys()) + sorted(\n",
+        "    [k.title() for k in TO_LANGUAGE_CODE.keys()]\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def create_config(output_dir):\n",
+        "    DOMAIN_TYPE = \"telephonic\"  # Can be meeting, telephonic, or general based on domain type of the audio file\n",
+        "    CONFIG_FILE_NAME = f\"diar_infer_{DOMAIN_TYPE}.yaml\"\n",
+        "    CONFIG_URL = f\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/{CONFIG_FILE_NAME}\"\n",
+        "    MODEL_CONFIG = os.path.join(output_dir, CONFIG_FILE_NAME)\n",
+        "    if not os.path.exists(MODEL_CONFIG):\n",
+        "        MODEL_CONFIG = wget.download(CONFIG_URL, output_dir)\n",
+        "\n",
+        "    config = OmegaConf.load(MODEL_CONFIG)\n",
+        "\n",
+        "    data_dir = os.path.join(output_dir, \"data\")\n",
+        "    os.makedirs(data_dir, exist_ok=True)\n",
+        "\n",
+        "    meta = {\n",
+        "        \"audio_filepath\": os.path.join(output_dir, \"mono_file.wav\"),\n",
+        "        \"offset\": 0,\n",
+        "        \"duration\": None,\n",
+        "        \"label\": \"infer\",\n",
+        "        \"text\": \"-\",\n",
+        "        \"rttm_filepath\": None,\n",
+        "        \"uem_filepath\": None,\n",
+        "    }\n",
+        "    with open(os.path.join(data_dir, \"input_manifest.json\"), \"w\") as fp:\n",
+        "        json.dump(meta, fp)\n",
+        "        fp.write(\"\\n\")\n",
+        "\n",
+        "    pretrained_vad = \"vad_multilingual_marblenet\"\n",
+        "    pretrained_speaker_model = \"titanet_large\"\n",
+        "    config.num_workers = 0  # Workaround for multiprocessing hanging with ipython issue\n",
+        "    config.diarizer.manifest_filepath = os.path.join(data_dir, \"input_manifest.json\")\n",
+        "    config.diarizer.out_dir = (\n",
+        "        output_dir  # Directory to store intermediate files and prediction outputs\n",
+        "    )\n",
+        "\n",
+        "    config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model\n",
+        "    config.diarizer.oracle_vad = (\n",
+        "        False  # compute VAD provided with model_path to vad config\n",
+        "    )\n",
+        "    config.diarizer.clustering.parameters.oracle_num_speakers = False\n",
+        "\n",
+        "    # Here, we use our in-house pretrained NeMo VAD model\n",
+        "    config.diarizer.vad.model_path = pretrained_vad\n",
+        "    config.diarizer.vad.parameters.onset = 0.8\n",
+        "    config.diarizer.vad.parameters.offset = 0.6\n",
+        "    config.diarizer.vad.parameters.pad_offset = -0.05\n",
+        "    config.diarizer.msdd_model.model_path = (\n",
+        "        \"diar_msdd_telephonic\"  # Telephonic speaker diarization model\n",
+        "    )\n",
+        "\n",
+        "    return config\n",
+        "\n",
+        "\n",
+        "def get_word_ts_anchor(s, e, option=\"start\"):\n",
+        "    if option == \"end\":\n",
+        "        return e\n",
+        "    elif option == \"mid\":\n",
+        "        return (s + e) / 2\n",
+        "    return s\n",
+        "\n",
+        "\n",
+        "def get_words_speaker_mapping(wrd_ts, spk_ts, word_anchor_option=\"start\"):\n",
+        "    s, e, sp = spk_ts[0]\n",
+        "    wrd_pos, turn_idx = 0, 0\n",
+        "    wrd_spk_mapping = []\n",
+        "    for wrd_dict in wrd_ts:\n",
+        "        ws, we, wrd = (\n",
+        "            int(wrd_dict[\"start\"] * 1000),\n",
+        "            int(wrd_dict[\"end\"] * 1000),\n",
+        "            wrd_dict[\"text\"],\n",
+        "        )\n",
+        "        wrd_pos = get_word_ts_anchor(ws, we, word_anchor_option)\n",
+        "        while wrd_pos > float(e):\n",
+        "            turn_idx += 1\n",
+        "            turn_idx = min(turn_idx, len(spk_ts) - 1)\n",
+        "            s, e, sp = spk_ts[turn_idx]\n",
+        "            if turn_idx == len(spk_ts) - 1:\n",
+        "                e = get_word_ts_anchor(ws, we, option=\"end\")\n",
+        "        wrd_spk_mapping.append(\n",
+        "            {\"word\": wrd, \"start_time\": ws, \"end_time\": we, \"speaker\": sp}\n",
+        "        )\n",
+        "    return wrd_spk_mapping\n",
+        "\n",
+        "\n",
+        "sentence_ending_punctuations = \".?!\"\n",
+        "\n",
+        "\n",
+        "def get_first_word_idx_of_sentence(word_idx, word_list, speaker_list, max_words):\n",
+        "    is_word_sentence_end = (\n",
+        "        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations\n",
+        "    )\n",
+        "    left_idx = word_idx\n",
+        "    while (\n",
+        "        left_idx > 0\n",
+        "        and word_idx - left_idx < max_words\n",
+        "        and speaker_list[left_idx - 1] == speaker_list[left_idx]\n",
+        "        and not is_word_sentence_end(left_idx - 1)\n",
+        "    ):\n",
+        "        left_idx -= 1\n",
+        "\n",
+        "    return left_idx if left_idx == 0 or is_word_sentence_end(left_idx - 1) else -1\n",
+        "\n",
+        "\n",
+        "def get_last_word_idx_of_sentence(word_idx, word_list, max_words):\n",
+        "    is_word_sentence_end = (\n",
+        "        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations\n",
+        "    )\n",
+        "    right_idx = word_idx\n",
+        "    while (\n",
+        "        right_idx < len(word_list) - 1\n",
+        "        and right_idx - word_idx < max_words\n",
+        "        and not is_word_sentence_end(right_idx)\n",
+        "    ):\n",
+        "        right_idx += 1\n",
+        "\n",
+        "    return (\n",
+        "        right_idx\n",
+        "        if right_idx == len(word_list) - 1 or is_word_sentence_end(right_idx)\n",
+        "        else -1\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "def get_realigned_ws_mapping_with_punctuation(\n",
+        "    word_speaker_mapping, max_words_in_sentence=50\n",
+        "):\n",
+        "    is_word_sentence_end = (\n",
+        "        lambda x: x >= 0\n",
+        "        and word_speaker_mapping[x][\"word\"][-1] in sentence_ending_punctuations\n",
+        "    )\n",
+        "    wsp_len = len(word_speaker_mapping)\n",
+        "\n",
+        "    words_list, speaker_list = [], []\n",
+        "    for k, line_dict in enumerate(word_speaker_mapping):\n",
+        "        word, speaker = line_dict[\"word\"], line_dict[\"speaker\"]\n",
+        "        words_list.append(word)\n",
+        "        speaker_list.append(speaker)\n",
+        "\n",
+        "    k = 0\n",
+        "    while k < len(word_speaker_mapping):\n",
+        "        line_dict = word_speaker_mapping[k]\n",
+        "        if (\n",
+        "            k < wsp_len - 1\n",
+        "            and speaker_list[k] != speaker_list[k + 1]\n",
+        "            and not is_word_sentence_end(k)\n",
+        "        ):\n",
+        "            left_idx = get_first_word_idx_of_sentence(\n",
+        "                k, words_list, speaker_list, max_words_in_sentence\n",
+        "            )\n",
+        "            right_idx = (\n",
+        "                get_last_word_idx_of_sentence(\n",
+        "                    k, words_list, max_words_in_sentence - k + left_idx - 1\n",
+        "                )\n",
+        "                if left_idx > -1\n",
+        "                else -1\n",
+        "            )\n",
+        "            if min(left_idx, right_idx) == -1:\n",
+        "                k += 1\n",
+        "                continue\n",
+        "\n",
+        "            spk_labels = speaker_list[left_idx : right_idx + 1]\n",
+        "            mod_speaker = max(set(spk_labels), key=spk_labels.count)\n",
+        "            if spk_labels.count(mod_speaker) < len(spk_labels) // 2:\n",
+        "                k += 1\n",
+        "                continue\n",
+        "\n",
+        "            speaker_list[left_idx : right_idx + 1] = [mod_speaker] * (\n",
+        "                right_idx - left_idx + 1\n",
+        "            )\n",
+        "            k = right_idx\n",
+        "\n",
+        "        k += 1\n",
+        "\n",
+        "    k, realigned_list = 0, []\n",
+        "    while k < len(word_speaker_mapping):\n",
+        "        line_dict = word_speaker_mapping[k].copy()\n",
+        "        line_dict[\"speaker\"] = speaker_list[k]\n",
+        "        realigned_list.append(line_dict)\n",
+        "        k += 1\n",
+        "\n",
+        "    return realigned_list\n",
+        "\n",
+        "\n",
+        "def get_sentences_speaker_mapping(word_speaker_mapping, spk_ts):\n",
+        "    sentence_checker = nltk.tokenize.PunktSentenceTokenizer().text_contains_sentbreak\n",
+        "    s, e, spk = spk_ts[0]\n",
+        "    prev_spk = spk\n",
+        "\n",
+        "    snts = []\n",
+        "    snt = {\"speaker\": f\"Speaker {spk}\", \"start_time\": s, \"end_time\": e, \"text\": \"\"}\n",
+        "\n",
+        "    for wrd_dict in word_speaker_mapping:\n",
+        "        wrd, spk = wrd_dict[\"word\"], wrd_dict[\"speaker\"]\n",
+        "        s, e = wrd_dict[\"start_time\"], wrd_dict[\"end_time\"]\n",
+        "        if spk != prev_spk or sentence_checker(snt[\"text\"] + \" \" + wrd):\n",
+        "            snts.append(snt)\n",
+        "            snt = {\n",
+        "                \"speaker\": f\"Speaker {spk}\",\n",
+        "                \"start_time\": s,\n",
+        "                \"end_time\": e,\n",
+        "                \"text\": \"\",\n",
+        "            }\n",
+        "        else:\n",
+        "            snt[\"end_time\"] = e\n",
+        "        snt[\"text\"] += wrd + \" \"\n",
+        "        prev_spk = spk\n",
+        "\n",
+        "    snts.append(snt)\n",
+        "    return snts\n",
+        "\n",
+        "\n",
+        "def get_speaker_aware_transcript(sentences_speaker_mapping, f):\n",
+        "    previous_speaker = sentences_speaker_mapping[0][\"speaker\"]\n",
+        "    f.write(f\"{previous_speaker}: \")\n",
+        "\n",
+        "    for sentence_dict in sentences_speaker_mapping:\n",
+        "        speaker = sentence_dict[\"speaker\"]\n",
+        "        sentence = sentence_dict[\"text\"]\n",
+        "\n",
+        "        # If this speaker doesn't match the previous one, start a new paragraph\n",
+        "        if speaker != previous_speaker:\n",
+        "            f.write(f\"\\n\\n{speaker}: \")\n",
+        "            previous_speaker = speaker\n",
+        "\n",
+        "        # No matter what, write the current sentence\n",
+        "        f.write(sentence + \" \")\n",
+        "\n",
+        "\n",
+        "def format_timestamp(\n",
+        "    milliseconds: float, always_include_hours: bool = False, decimal_marker: str = \".\"\n",
+        "):\n",
+        "    assert milliseconds >= 0, \"non-negative timestamp expected\"\n",
+        "\n",
+        "    hours = milliseconds // 3_600_000\n",
+        "    milliseconds -= hours * 3_600_000\n",
+        "\n",
+        "    minutes = milliseconds // 60_000\n",
+        "    milliseconds -= minutes * 60_000\n",
+        "\n",
+        "    seconds = milliseconds // 1_000\n",
+        "    milliseconds -= seconds * 1_000\n",
+        "\n",
+        "    hours_marker = f\"{hours:02d}:\" if always_include_hours or hours > 0 else \"\"\n",
+        "    return (\n",
+        "        f\"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}\"\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "def write_srt(transcript, file):\n",
+        "    \"\"\"\n",
+        "    Write a transcript to a file in SRT format.\n",
+        "\n",
+        "    \"\"\"\n",
+        "    for i, segment in enumerate(transcript, start=1):\n",
+        "        # write srt lines\n",
+        "        print(\n",
+        "            f\"{i}\\n\"\n",
+        "            f\"{format_timestamp(segment['start_time'], always_include_hours=True, decimal_marker=',')} --> \"\n",
+        "            f\"{format_timestamp(segment['end_time'], always_include_hours=True, decimal_marker=',')}\\n\"\n",
+        "            f\"{segment['speaker']}: {segment['text'].strip().replace('-->', '->')}\\n\",\n",
+        "            file=file,\n",
+        "            flush=True,\n",
+        "        )\n",
+        "\n",
+        "\n",
+        "def find_numeral_symbol_tokens(tokenizer):\n",
+        "    numeral_symbol_tokens = [\n",
+        "        -1,\n",
+        "    ]\n",
+        "    for token, token_id in tokenizer.get_vocab().items():\n",
+        "        has_numeral_symbol = any(c in \"0123456789%$£\" for c in token)\n",
+        "        if has_numeral_symbol:\n",
+        "            numeral_symbol_tokens.append(token_id)\n",
+        "    return numeral_symbol_tokens\n",
+        "\n",
+        "\n",
+        "def _get_next_start_timestamp(word_timestamps, current_word_index, final_timestamp):\n",
+        "    # if current word is the last word\n",
+        "    if current_word_index == len(word_timestamps) - 1:\n",
+        "        return word_timestamps[current_word_index][\"start\"]\n",
+        "\n",
+        "    next_word_index = current_word_index + 1\n",
+        "    while current_word_index < len(word_timestamps) - 1:\n",
+        "        if word_timestamps[next_word_index].get(\"start\") is None:\n",
+        "            # if next word doesn't have a start timestamp\n",
+        "            # merge it with the current word and delete it\n",
+        "            word_timestamps[current_word_index][\"word\"] += (\n",
+        "                \" \" + word_timestamps[next_word_index][\"word\"]\n",
+        "            )\n",
+        "\n",
+        "            word_timestamps[next_word_index][\"word\"] = None\n",
+        "            next_word_index += 1\n",
+        "            if next_word_index == len(word_timestamps):\n",
+        "                return final_timestamp\n",
+        "\n",
+        "        else:\n",
+        "            return word_timestamps[next_word_index][\"start\"]\n",
+        "\n",
+        "\n",
+        "def filter_missing_timestamps(\n",
+        "    word_timestamps, initial_timestamp=0, final_timestamp=None\n",
+        "):\n",
+        "    # handle the first and last word\n",
+        "    if word_timestamps[0].get(\"start\") is None:\n",
+        "        word_timestamps[0][\"start\"] = (\n",
+        "            initial_timestamp if initial_timestamp is not None else 0\n",
+        "        )\n",
+        "        word_timestamps[0][\"end\"] = _get_next_start_timestamp(\n",
+        "            word_timestamps, 0, final_timestamp\n",
+        "        )\n",
+        "\n",
+        "    result = [\n",
+        "        word_timestamps[0],\n",
+        "    ]\n",
+        "\n",
+        "    for i, ws in enumerate(word_timestamps[1:], start=1):\n",
+        "        # if ws doesn't have a start and end\n",
+        "        # use the previous end as start and next start as end\n",
+        "        if ws.get(\"start\") is None and ws.get(\"word\") is not None:\n",
+        "            ws[\"start\"] = word_timestamps[i - 1][\"end\"]\n",
+        "            ws[\"end\"] = _get_next_start_timestamp(word_timestamps, i, final_timestamp)\n",
+        "\n",
+        "        if ws[\"word\"] is not None:\n",
+        "            result.append(ws)\n",
+        "    return result\n",
+        "\n",
+        "\n",
+        "def cleanup(path: str):\n",
+        "    \"\"\"path could either be relative or absolute.\"\"\"\n",
+        "    # check if file or directory exists\n",
+        "    if os.path.isfile(path) or os.path.islink(path):\n",
+        "        # remove file\n",
+        "        os.remove(path)\n",
+        "    elif os.path.isdir(path):\n",
+        "        # remove directory and all its content\n",
+        "        shutil.rmtree(path)\n",
+        "    else:\n",
+        "        raise ValueError(\"Path {} is not a file or dir.\".format(path))\n",
+        "\n",
+        "\n",
+        "def process_language_arg(language: str, model_name: str):\n",
+        "    \"\"\"\n",
+        "    Process the language argument to make sure it's valid and convert language names to language codes.\n",
+        "    \"\"\"\n",
+        "    if language is not None:\n",
+        "        language = language.lower()\n",
+        "    if language not in LANGUAGES:\n",
+        "        if language in TO_LANGUAGE_CODE:\n",
+        "            language = TO_LANGUAGE_CODE[language]\n",
+        "        else:\n",
+        "            raise ValueError(f\"Unsupported language: {language}\")\n",
+        "\n",
+        "    if model_name.endswith(\".en\") and language != \"en\":\n",
+        "        if language is not None:\n",
+        "            logging.warning(\n",
+        "                f\"{model_name} is an English-only model but received '{language}'; using English instead.\"\n",
+        "            )\n",
+        "        language = \"en\"\n",
+        "    return language\n",
+        "\n",
+        "\n",
+        "def transcribe_batched(\n",
+        "    audio_file: str,\n",
+        "    language: str,\n",
+        "    batch_size: int,\n",
+        "    model_name: str,\n",
+        "    compute_dtype: str,\n",
+        "    suppress_numerals: bool,\n",
+        "    device: str,\n",
+        "):\n",
+        "    import whisperx\n",
+        "\n",
+        "    # Faster Whisper batched\n",
+        "    whisper_model = whisperx.load_model(\n",
+        "        model_name,\n",
+        "        device,\n",
+        "        compute_type=compute_dtype,\n",
+        "        asr_options={\"suppress_numerals\": suppress_numerals},\n",
+        "    )\n",
+        "    audio = whisperx.load_audio(audio_file)\n",
+        "    result = whisper_model.transcribe(audio, language=language, batch_size=batch_size)\n",
+        "    del whisper_model\n",
+        "    torch.cuda.empty_cache()\n",
+        "    return result[\"segments\"], result[\"language\"], audio"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "B7qWQb--1Xcw"
+      },
+      "source": [
+        "# Options"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ONlFrSnD0FOp"
+      },
+      "outputs": [],
+      "source": [
+        "# Name of the audio file\n",
+        "audio_path = \"20200128-Pieter Wuille (part 1 of 2) - Episode 1.mp3\"\n",
+        "\n",
+        "# Whether to enable music removal from speech, helps increase diarization quality but uses alot of ram\n",
+        "enable_stemming = True\n",
+        "\n",
+        "# (choose from 'tiny.en', 'tiny', 'base.en', 'base', 'small.en', 'small', 'medium.en', 'medium', 'large-v1', 'large-v2', 'large-v3', 'large')\n",
+        "whisper_model_name = \"large-v2\"\n",
+        "\n",
+        "# replaces numerical digits with their pronounciation, increases diarization accuracy\n",
+        "suppress_numerals = True\n",
+        "\n",
+        "batch_size = 8\n",
+        "\n",
+        "language = None  # autodetect language\n",
+        "\n",
+        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "h-cY1ZEy2KVI"
+      },
+      "source": [
+        "# Processing"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7ZS4xXmE2NGP"
+      },
+      "source": [
+        "## Separating music from speech using Demucs\n",
+        "\n",
+        "---\n",
+        "\n",
+        "By isolating the vocals from the rest of the audio, it becomes easier to identify and track individual speakers based on the spectral and temporal characteristics of their speech signals. Source separation is just one of many techniques that can be used as a preprocessing step to help improve the accuracy and reliability of the overall diarization process."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "HKcgQUrAzsJZ",
+        "outputId": "dc2a1d96-20da-4749-9d64-21edacfba1b1"
+      },
+      "outputs": [],
+      "source": [
+        "if enable_stemming:\n",
+        "    # Isolate vocals from the rest of the audio\n",
+        "\n",
+        "    return_code = os.system(\n",
+        "        f'python3 -m demucs.separate -n htdemucs --two-stems=vocals \"{audio_path}\" -o \"temp_outputs\"'\n",
+        "    )\n",
+        "\n",
+        "    if return_code != 0:\n",
+        "        logging.warning(\"Source splitting failed, using original audio file.\")\n",
+        "        vocal_target = audio_path\n",
+        "    else:\n",
+        "        vocal_target = os.path.join(\n",
+        "            \"temp_outputs\",\n",
+        "            \"htdemucs\",\n",
+        "            os.path.splitext(os.path.basename(audio_path))[0],\n",
+        "            \"vocals.wav\",\n",
+        "        )\n",
+        "else:\n",
+        "    vocal_target = audio_path"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UYg9VWb22Tz8"
+      },
+      "source": [
+        "## Transcriping audio using Whisper and realligning timestamps using Wav2Vec2\n",
+        "---\n",
+        "This code uses two different open-source models to transcribe speech and perform forced alignment on the resulting transcription.\n",
+        "\n",
+        "The first model is called OpenAI Whisper, which is a speech recognition model that can transcribe speech with high accuracy. The code loads the whisper model and uses it to transcribe the vocal_target file.\n",
+        "\n",
+        "The output of the transcription process is a set of text segments with corresponding timestamps indicating when each segment was spoken.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5-VKFn530oTl"
+      },
+      "outputs": [],
+      "source": [
+        "compute_type = \"float16\"\n",
+        "# or run on GPU with INT8\n",
+        "# compute_type = \"int8_float16\"\n",
+        "# or run on CPU with INT8\n",
+        "# compute_type = \"int8\"\n",
+        "\n",
+        "whisper_results, language, audio_waveform = transcribe_batched(\n",
+        "    vocal_target,\n",
+        "    language,\n",
+        "    batch_size,\n",
+        "    whisper_model_name,\n",
+        "    compute_type,\n",
+        "    suppress_numerals,\n",
+        "    device,\n",
+        ")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Aligning the transcription with the original audio using Wav2Vec2\n",
+        "---\n",
+        "The second model used is called wav2vec2, which is a large-scale neural network that is designed to learn representations of speech that are useful for a variety of speech processing tasks, including speech recognition and alignment.\n",
+        "\n",
+        "The code loads the wav2vec2 alignment model and uses it to align the transcription segments with the original audio signal contained in the vocal_target file. This process involves finding the exact timestamps in the audio signal where each segment was spoken and aligning the text accordingly.\n",
+        "\n",
+        "By combining the outputs of the two models, the code produces a fully aligned transcription of the speech contained in the vocal_target file. This aligned transcription can be useful for a variety of speech processing tasks, such as speaker diarization, sentiment analysis, and language identification.\n",
+        "\n",
+        "If there's no Wav2Vec2 model available for your language, word timestamps generated by whisper will be used instead."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "alignment_model, alignment_tokenizer, alignment_dictionary = load_alignment_model(\n",
+        "    device,\n",
+        "    dtype=torch.float16 if device == \"cuda\" else torch.float32,\n",
+        ")\n",
+        "\n",
+        "audio_waveform = (\n",
+        "    torch.from_numpy(audio_waveform)\n",
+        "    .to(alignment_model.dtype)\n",
+        "    .to(alignment_model.device)\n",
+        ")\n",
+        "\n",
+        "emissions, stride = generate_emissions(\n",
+        "    alignment_model, audio_waveform, batch_size=batch_size\n",
+        ")\n",
+        "\n",
+        "del alignment_model\n",
+        "torch.cuda.empty_cache()\n",
+        "\n",
+        "full_transcript = \"\".join(segment[\"text\"] for segment in whisper_results)\n",
+        "\n",
+        "tokens_starred, text_starred = preprocess_text(\n",
+        "    full_transcript,\n",
+        "    romanize=True,\n",
+        "    language=langs_to_iso[language],\n",
+        ")\n",
+        "\n",
+        "segments, scores, blank_id = get_alignments(\n",
+        "    emissions,\n",
+        "    tokens_starred,\n",
+        "    alignment_dictionary,\n",
+        ")\n",
+        "\n",
+        "spans = get_spans(tokens_starred, segments, alignment_tokenizer.decode(blank_id))\n",
+        "\n",
+        "word_timestamps = postprocess_results(text_starred, spans, stride, scores)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7EEaJPsQ21Rx"
+      },
+      "source": [
+        "## Convert audio to mono for NeMo combatibility"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "ROOT = os.getcwd()\n",
+        "temp_path = os.path.join(ROOT, \"temp_outputs\")\n",
+        "os.makedirs(temp_path, exist_ok=True)\n",
+        "torchaudio.save(\n",
+        "    os.path.join(temp_path, \"mono_file.wav\"),\n",
+        "    audio_waveform.cpu().unsqueeze(0).float(),\n",
+        "    16000,\n",
+        "    channels_first=True,\n",
+        ")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "D1gkViCf2-CV"
+      },
+      "source": [
+        "## Speaker Diarization using NeMo MSDD Model\n",
+        "---\n",
+        "This code uses a model called Nvidia NeMo MSDD (Multi-scale Diarization Decoder) to perform speaker diarization on an audio signal. Speaker diarization is the process of separating an audio signal into different segments based on who is speaking at any given time."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "C7jIpBCH02RL"
+      },
+      "outputs": [],
+      "source": [
+        "# Initialize NeMo MSDD diarization model\n",
+        "msdd_model = NeuralDiarizer(cfg=create_config(temp_path)).to(\"cuda\")\n",
+        "msdd_model.diarize()\n",
+        "\n",
+        "del msdd_model\n",
+        "torch.cuda.empty_cache()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NmkZYaDAEOAg"
+      },
+      "source": [
+        "## Mapping Spekers to Sentences According to Timestamps"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E65LUGQe02zw"
+      },
+      "outputs": [],
+      "source": [
+        "# Reading timestamps <> Speaker Labels mapping\n",
+        "\n",
+        "speaker_ts = []\n",
+        "with open(os.path.join(temp_path, \"pred_rttms\", \"mono_file.rttm\"), \"r\") as f:\n",
+        "    lines = f.readlines()\n",
+        "    for line in lines:\n",
+        "        line_list = line.split(\" \")\n",
+        "        s = int(float(line_list[5]) * 1000)\n",
+        "        e = s + int(float(line_list[8]) * 1000)\n",
+        "        speaker_ts.append([s, e, int(line_list[11].split(\"_\")[-1])])\n",
+        "\n",
+        "wsm = get_words_speaker_mapping(word_timestamps, speaker_ts, \"start\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8Ruxc8S1EXtW"
+      },
+      "source": [
+        "## Realligning Speech segments using Punctuation\n",
+        "---\n",
+        "\n",
+        "This code provides a method for disambiguating speaker labels in cases where a sentence is split between two different speakers. It uses punctuation markings to determine the dominant speaker for each sentence in the transcription.\n",
+        "\n",
+        "```\n",
+        "Speaker A: It's got to come from somewhere else. Yeah, that one's also fun because you know the lows are\n",
+        "Speaker B: going to suck, right? So it's actually it hits you on both sides.\n",
+        "```\n",
+        "\n",
+        "For example, if a sentence is split between two speakers, the code takes the mode of speaker labels for each word in the sentence, and uses that speaker label for the whole sentence. This can help to improve the accuracy of speaker diarization, especially in cases where the Whisper model may not take fine utterances like \"hmm\" and \"yeah\" into account, but the Diarization Model (Nemo) may include them, leading to inconsistent results.\n",
+        "\n",
+        "The code also handles cases where one speaker is giving a monologue while other speakers are making occasional comments in the background. It ignores the comments and assigns the entire monologue to the speaker who is speaking the majority of the time. This provides a robust and reliable method for realigning speech segments to their respective speakers based on punctuation in the transcription."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pgfC5hA41BXu"
+      },
+      "outputs": [],
+      "source": [
+        "if language in punct_model_langs:\n",
+        "    # restoring punctuation in the transcript to help realign the sentences\n",
+        "    punct_model = PunctuationModel(model=\"kredor/punctuate-all\")\n",
+        "\n",
+        "    words_list = list(map(lambda x: x[\"word\"], wsm))\n",
+        "\n",
+        "    labled_words = punct_model.predict(words_list,chunk_size=230)\n",
+        "\n",
+        "    ending_puncts = \".?!\"\n",
+        "    model_puncts = \".,;:!?\"\n",
+        "\n",
+        "    # We don't want to punctuate U.S.A. with a period. Right?\n",
+        "    is_acronym = lambda x: re.fullmatch(r\"\\b(?:[a-zA-Z]\\.){2,}\", x)\n",
+        "\n",
+        "    for word_dict, labeled_tuple in zip(wsm, labled_words):\n",
+        "        word = word_dict[\"word\"]\n",
+        "        if (\n",
+        "            word\n",
+        "            and labeled_tuple[1] in ending_puncts\n",
+        "            and (word[-1] not in model_puncts or is_acronym(word))\n",
+        "        ):\n",
+        "            word += labeled_tuple[1]\n",
+        "            if word.endswith(\"..\"):\n",
+        "                word = word.rstrip(\".\")\n",
+        "            word_dict[\"word\"] = word\n",
+        "\n",
+        "else:\n",
+        "    logging.warning(\n",
+        "        f\"Punctuation restoration is not available for {language} language. Using the original punctuation.\"\n",
+        "    )\n",
+        "\n",
+        "wsm = get_realigned_ws_mapping_with_punctuation(wsm)\n",
+        "ssm = get_sentences_speaker_mapping(wsm, speaker_ts)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vF2QAtLOFvwZ"
+      },
+      "source": [
+        "## Cleanup and Exporing the results"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "kFTyKI6B1MI0"
+      },
+      "outputs": [],
+      "source": [
+        "with open(f\"{os.path.splitext(audio_path)[0]}.txt\", \"w\", encoding=\"utf-8-sig\") as f:\n",
+        "    get_speaker_aware_transcript(ssm, f)\n",
+        "\n",
+        "with open(f\"{os.path.splitext(audio_path)[0]}.srt\", \"w\", encoding=\"utf-8-sig\") as srt:\n",
+        "    write_srt(ssm, srt)\n",
+        "\n",
+        "cleanup(temp_path)"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "authorship_tag": "ABX9TyOyiQNkD+ROzss634BOsrSh",
+      "collapsed_sections": [
+        "eCmjcOc9yEtQ",
+        "jbsUt3SwyhjD"
+      ],
+      "include_colab_link": true,
+      "provenance": []
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "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.10.12"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

whisper_diarization_main/diarize.py

@@ -0,0 +1,235 @@
+import argparse
+import logging
+import os
+import re
+
+import torch
+import torchaudio
+from ctc_forced_aligner import (
+    generate_emissions,
+    get_alignments,
+    get_spans,
+    load_alignment_model,
+    postprocess_results,
+    preprocess_text,
+)
+from deepmultilingualpunctuation import PunctuationModel
+from nemo.collections.asr.models.msdd_models import NeuralDiarizer
+
+from helpers import (
+    cleanup,
+    create_config,
+    get_realigned_ws_mapping_with_punctuation,
+    get_sentences_speaker_mapping,
+    get_speaker_aware_transcript,
+    get_words_speaker_mapping,
+    langs_to_iso,
+    punct_model_langs,
+    whisper_langs,
+    write_srt,
+)
+from transcription_helpers import transcribe_batched
+
+mtypes = {"cpu": "int8", "cuda": "float16"}
+
+# Initialize parser
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-a", "--audio", help="name of the target audio file", required=True
+)
+parser.add_argument(
+    "--no-stem",
+    action="store_false",
+    dest="stemming",
+    default=True,
+    help="Disables source separation."
+    "This helps with long files that don't contain a lot of music.",
+)
+
+parser.add_argument(
+    "--suppress_numerals",
+    action="store_true",
+    dest="suppress_numerals",
+    default=False,
+    help="Suppresses Numerical Digits."
+    "This helps the diarization accuracy but converts all digits into written text.",
+)
+
+parser.add_argument(
+    "--whisper-model",
+    dest="model_name",
+    default="medium.en",
+    help="name of the Whisper model to use",
+)
+
+parser.add_argument(
+    "--batch-size",
+    type=int,
+    dest="batch_size",
+    default=8,
+    help="Batch size for batched inference, reduce if you run out of memory, set to 0 for non-batched inference",
+)
+
+parser.add_argument(
+    "--language",
+    type=str,
+    default=None,
+    choices=whisper_langs,
+    help="Language spoken in the audio, specify None to perform language detection",
+)
+
+parser.add_argument(
+    "--device",
+    dest="device",
+    default="cuda" if torch.cuda.is_available() else "cpu",
+    help="if you have a GPU use 'cuda', otherwise 'cpu'",
+)
+
+args = parser.parse_args()
+
+if args.stemming:
+    # Isolate vocals from the rest of the audio
+
+    return_code = os.system(
+        f'python3 -m demucs.separate -n htdemucs --two-stems=vocals "{args.audio}" -o "temp_outputs"'
+    )
+
+    if return_code != 0:
+        logging.warning(
+            "Source splitting failed, using original audio file. Use --no-stem argument to disable it."
+        )
+        vocal_target = args.audio
+    else:
+        vocal_target = os.path.join(
+            "temp_outputs",
+            "htdemucs",
+            os.path.splitext(os.path.basename(args.audio))[0],
+            "vocals.wav",
+        )
+else:
+    vocal_target = args.audio
+
+
+# Transcribe the audio file
+
+whisper_results, language, audio_waveform = transcribe_batched(
+    vocal_target,
+    args.language,
+    args.batch_size,
+    args.model_name,
+    mtypes[args.device],
+    args.suppress_numerals,
+    args.device,
+)
+
+# Forced Alignment
+alignment_model, alignment_tokenizer, alignment_dictionary = load_alignment_model(
+    args.device,
+    dtype=torch.float16 if args.device == "cuda" else torch.float32,
+)
+
+audio_waveform = (
+    torch.from_numpy(audio_waveform)
+    .to(alignment_model.dtype)
+    .to(alignment_model.device)
+)
+emissions, stride = generate_emissions(
+    alignment_model, audio_waveform, batch_size=args.batch_size
+)
+
+del alignment_model
+torch.cuda.empty_cache()
+
+full_transcript = "".join(segment["text"] for segment in whisper_results)
+
+tokens_starred, text_starred = preprocess_text(
+    full_transcript,
+    romanize=True,
+    language=langs_to_iso[language],
+)
+
+segments, scores, blank_id = get_alignments(
+    emissions,
+    tokens_starred,
+    alignment_dictionary,
+)
+
+spans = get_spans(tokens_starred, segments, alignment_tokenizer.decode(blank_id))
+
+word_timestamps = postprocess_results(text_starred, spans, stride, scores)
+
+
+# convert audio to mono for NeMo combatibility
+ROOT = os.getcwd()
+temp_path = os.path.join(ROOT, "temp_outputs")
+os.makedirs(temp_path, exist_ok=True)
+torchaudio.save(
+    os.path.join(temp_path, "mono_file.wav"),
+    audio_waveform.cpu().unsqueeze(0).float(),
+    16000,
+    channels_first=True,
+)
+
+
+# Initialize NeMo MSDD diarization model
+msdd_model = NeuralDiarizer(cfg=create_config(temp_path)).to(args.device)
+msdd_model.diarize()
+
+del msdd_model
+torch.cuda.empty_cache()
+
+# Reading timestamps <> Speaker Labels mapping
+
+
+speaker_ts = []
+with open(os.path.join(temp_path, "pred_rttms", "mono_file.rttm"), "r") as f:
+    lines = f.readlines()
+    for line in lines:
+        line_list = line.split(" ")
+        s = int(float(line_list[5]) * 1000)
+        e = s + int(float(line_list[8]) * 1000)
+        speaker_ts.append([s, e, int(line_list[11].split("_")[-1])])
+
+wsm = get_words_speaker_mapping(word_timestamps, speaker_ts, "start")
+
+if language in punct_model_langs:
+    # restoring punctuation in the transcript to help realign the sentences
+    punct_model = PunctuationModel(model="kredor/punctuate-all")
+
+    words_list = list(map(lambda x: x["word"], wsm))
+
+    labled_words = punct_model.predict(words_list, chunk_size=230)
+
+    ending_puncts = ".?!"
+    model_puncts = ".,;:!?"
+
+    # We don't want to punctuate U.S.A. with a period. Right?
+    is_acronym = lambda x: re.fullmatch(r"\b(?:[a-zA-Z]\.){2,}", x)
+
+    for word_dict, labeled_tuple in zip(wsm, labled_words):
+        word = word_dict["word"]
+        if (
+            word
+            and labeled_tuple[1] in ending_puncts
+            and (word[-1] not in model_puncts or is_acronym(word))
+        ):
+            word += labeled_tuple[1]
+            if word.endswith(".."):
+                word = word.rstrip(".")
+            word_dict["word"] = word
+
+else:
+    logging.warning(
+        f"Punctuation restoration is not available for {language} language. Using the original punctuation."
+    )
+
+wsm = get_realigned_ws_mapping_with_punctuation(wsm)
+ssm = get_sentences_speaker_mapping(wsm, speaker_ts)
+
+with open(f"{os.path.splitext(args.audio)[0]}.txt", "w", encoding="utf-8-sig") as f:
+    get_speaker_aware_transcript(ssm, f)
+
+with open(f"{os.path.splitext(args.audio)[0]}.srt", "w", encoding="utf-8-sig") as srt:
+    write_srt(ssm, srt)
+
+cleanup(temp_path)

whisper_diarization_main/diarize_parallel.py

@@ -0,0 +1,217 @@
+import argparse
+import logging
+import os
+import re
+import subprocess
+
+import torch
+from ctc_forced_aligner import (
+    generate_emissions,
+    get_alignments,
+    get_spans,
+    load_alignment_model,
+    postprocess_results,
+    preprocess_text,
+)
+from deepmultilingualpunctuation import PunctuationModel
+
+from helpers import (
+    cleanup,
+    get_realigned_ws_mapping_with_punctuation,
+    get_sentences_speaker_mapping,
+    get_speaker_aware_transcript,
+    get_words_speaker_mapping,
+    langs_to_iso,
+    punct_model_langs,
+    whisper_langs,
+    write_srt,
+)
+from transcription_helpers import transcribe_batched
+
+mtypes = {"cpu": "int8", "cuda": "float16"}
+
+# Initialize parser
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-a", "--audio", help="name of the target audio file", required=True
+)
+parser.add_argument(
+    "--no-stem",
+    action="store_false",
+    dest="stemming",
+    default=True,
+    help="Disables source separation."
+    "This helps with long files that don't contain a lot of music.",
+)
+
+parser.add_argument(
+    "--suppress_numerals",
+    action="store_true",
+    dest="suppress_numerals",
+    default=False,
+    help="Suppresses Numerical Digits."
+    "This helps the diarization accuracy but converts all digits into written text.",
+)
+
+parser.add_argument(
+    "--whisper-model",
+    dest="model_name",
+    default="medium.en",
+    help="name of the Whisper model to use",
+)
+
+parser.add_argument(
+    "--batch-size",
+    type=int,
+    dest="batch_size",
+    default=8,
+    help="Batch size for batched inference, reduce if you run out of memory, set to 0 for non-batched inference",
+)
+
+parser.add_argument(
+    "--language",
+    type=str,
+    default=None,
+    choices=whisper_langs,
+    help="Language spoken in the audio, specify None to perform language detection",
+)
+
+parser.add_argument(
+    "--device",
+    dest="device",
+    default="cuda" if torch.cuda.is_available() else "cpu",
+    help="if you have a GPU use 'cuda', otherwise 'cpu'",
+)
+
+args = parser.parse_args()
+
+if args.stemming:
+    # Isolate vocals from the rest of the audio
+
+    return_code = os.system(
+        f'python3 -m demucs.separate -n htdemucs --two-stems=vocals "{args.audio}" -o "temp_outputs"'
+    )
+
+    if return_code != 0:
+        logging.warning(
+            "Source splitting failed, using original audio file. Use --no-stem argument to disable it."
+        )
+        vocal_target = args.audio
+    else:
+        vocal_target = os.path.join(
+            "temp_outputs",
+            "htdemucs",
+            os.path.splitext(os.path.basename(args.audio))[0],
+            "vocals.wav",
+        )
+else:
+    vocal_target = args.audio
+
+logging.info("Starting Nemo process with vocal_target: ", vocal_target)
+nemo_process = subprocess.Popen(
+    ["python3", "nemo_process.py", "-a", vocal_target, "--device", args.device],
+)
+# Transcribe the audio file
+whisper_results, language, audio_waveform = transcribe_batched(
+    vocal_target,
+    args.language,
+    args.batch_size,
+    args.model_name,
+    mtypes[args.device],
+    args.suppress_numerals,
+    args.device,
+)
+
+# Forced Alignment
+alignment_model, alignment_tokenizer, alignment_dictionary = load_alignment_model(
+    args.device,
+    dtype=torch.float16 if args.device == "cuda" else torch.float32,
+)
+
+audio_waveform = (
+    torch.from_numpy(audio_waveform)
+    .to(alignment_model.dtype)
+    .to(alignment_model.device)
+)
+emissions, stride = generate_emissions(
+    alignment_model, audio_waveform, batch_size=args.batch_size
+)
+
+del alignment_model
+torch.cuda.empty_cache()
+
+full_transcript = "".join(segment["text"] for segment in whisper_results)
+
+tokens_starred, text_starred = preprocess_text(
+    full_transcript,
+    romanize=True,
+    language=langs_to_iso[language],
+)
+
+segments, scores, blank_id = get_alignments(
+    emissions,
+    tokens_starred,
+    alignment_dictionary,
+)
+
+spans = get_spans(tokens_starred, segments, alignment_tokenizer.decode(blank_id))
+
+word_timestamps = postprocess_results(text_starred, spans, stride, scores)
+
+# Reading timestamps <> Speaker Labels mapping
+nemo_process.communicate()
+ROOT = os.getcwd()
+temp_path = os.path.join(ROOT, "temp_outputs")
+
+speaker_ts = []
+with open(os.path.join(temp_path, "pred_rttms", "mono_file.rttm"), "r") as f:
+    lines = f.readlines()
+    for line in lines:
+        line_list = line.split(" ")
+        s = int(float(line_list[5]) * 1000)
+        e = s + int(float(line_list[8]) * 1000)
+        speaker_ts.append([s, e, int(line_list[11].split("_")[-1])])
+
+wsm = get_words_speaker_mapping(word_timestamps, speaker_ts, "start")
+
+if language in punct_model_langs:
+    # restoring punctuation in the transcript to help realign the sentences
+    punct_model = PunctuationModel(model="kredor/punctuate-all")
+
+    words_list = list(map(lambda x: x["word"], wsm))
+
+    labled_words = punct_model.predict(words_list, chunk_size=230)
+
+    ending_puncts = ".?!"
+    model_puncts = ".,;:!?"
+
+    # We don't want to punctuate U.S.A. with a period. Right?
+    is_acronym = lambda x: re.fullmatch(r"\b(?:[a-zA-Z]\.){2,}", x)
+
+    for word_dict, labeled_tuple in zip(wsm, labled_words):
+        word = word_dict["word"]
+        if (
+            word
+            and labeled_tuple[1] in ending_puncts
+            and (word[-1] not in model_puncts or is_acronym(word))
+        ):
+            word += labeled_tuple[1]
+            if word.endswith(".."):
+                word = word.rstrip(".")
+            word_dict["word"] = word
+
+else:
+    logging.warning(
+        f"Punctuation restoration is not available for {language} language. Using the original punctuation."
+    )
+
+wsm = get_realigned_ws_mapping_with_punctuation(wsm)
+ssm = get_sentences_speaker_mapping(wsm, speaker_ts)
+
+with open(f"{os.path.splitext(args.audio)[0]}.txt", "w", encoding="utf-8-sig") as f:
+    get_speaker_aware_transcript(ssm, f)
+
+with open(f"{os.path.splitext(args.audio)[0]}.srt", "w", encoding="utf-8-sig") as srt:
+    write_srt(ssm, srt)
+
+cleanup(temp_path)

whisper_diarization_main/helpers.py

@@ -0,0 +1,580 @@
+import json
+import logging
+import os
+import shutil
+
+import nltk
+import wget
+from omegaconf import OmegaConf
+from whisperx.alignment import DEFAULT_ALIGN_MODELS_HF, DEFAULT_ALIGN_MODELS_TORCH
+from whisperx.utils import LANGUAGES, TO_LANGUAGE_CODE
+
+punct_model_langs = [
+    "en",
+    "fr",
+    "de",
+    "es",
+    "it",
+    "nl",
+    "pt",
+    "bg",
+    "pl",
+    "cs",
+    "sk",
+    "sl",
+]
+wav2vec2_langs = list(DEFAULT_ALIGN_MODELS_TORCH.keys()) + list(
+    DEFAULT_ALIGN_MODELS_HF.keys()
+)
+
+whisper_langs = sorted(LANGUAGES.keys()) + sorted(
+    [k.title() for k in TO_LANGUAGE_CODE.keys()]
+)
+
+langs_to_iso = {
+    "aa": "aar",
+    "ab": "abk",
+    "ae": "ave",
+    "af": "afr",
+    "ak": "aka",
+    "am": "amh",
+    "an": "arg",
+    "ar": "ara",
+    "as": "asm",
+    "av": "ava",
+    "ay": "aym",
+    "az": "aze",
+    "ba": "bak",
+    "be": "bel",
+    "bg": "bul",
+    "bh": "bih",
+    "bi": "bis",
+    "bm": "bam",
+    "bn": "ben",
+    "bo": "tib",
+    "br": "bre",
+    "bs": "bos",
+    "ca": "cat",
+    "ce": "che",
+    "ch": "cha",
+    "co": "cos",
+    "cr": "cre",
+    "cs": "cze",
+    "cu": "chu",
+    "cv": "chv",
+    "cy": "wel",
+    "da": "dan",
+    "de": "ger",
+    "dv": "div",
+    "dz": "dzo",
+    "ee": "ewe",
+    "el": "gre",
+    "en": "eng",
+    "eo": "epo",
+    "es": "spa",
+    "et": "est",
+    "eu": "baq",
+    "fa": "per",
+    "ff": "ful",
+    "fi": "fin",
+    "fj": "fij",
+    "fo": "fao",
+    "fr": "fre",
+    "fy": "fry",
+    "ga": "gle",
+    "gd": "gla",
+    "gl": "glg",
+    "gn": "grn",
+    "gu": "guj",
+    "gv": "glv",
+    "ha": "hau",
+    "he": "heb",
+    "hi": "hin",
+    "ho": "hmo",
+    "hr": "hrv",
+    "ht": "hat",
+    "hu": "hun",
+    "hy": "arm",
+    "hz": "her",
+    "ia": "ina",
+    "id": "ind",
+    "ie": "ile",
+    "ig": "ibo",
+    "ii": "iii",
+    "ik": "ipk",
+    "io": "ido",
+    "is": "ice",
+    "it": "ita",
+    "iu": "iku",
+    "ja": "jpn",
+    "jv": "jav",
+    "ka": "geo",
+    "kg": "kon",
+    "ki": "kik",
+    "kj": "kua",
+    "kk": "kaz",
+    "kl": "kal",
+    "km": "khm",
+    "kn": "kan",
+    "ko": "kor",
+    "kr": "kau",
+    "ks": "kas",
+    "ku": "kur",
+    "kv": "kom",
+    "kw": "cor",
+    "ky": "kir",
+    "la": "lat",
+    "lb": "ltz",
+    "lg": "lug",
+    "li": "lim",
+    "ln": "lin",
+    "lo": "lao",
+    "lt": "lit",
+    "lu": "lub",
+    "lv": "lav",
+    "mg": "mlg",
+    "mh": "mah",
+    "mi": "mao",
+    "mk": "mac",
+    "ml": "mal",
+    "mn": "mon",
+    "mr": "mar",
+    "ms": "may",
+    "mt": "mlt",
+    "my": "bur",
+    "na": "nau",
+    "nb": "nob",
+    "nd": "nde",
+    "ne": "nep",
+    "ng": "ndo",
+    "nl": "dut",
+    "nn": "nno",
+    "no": "nor",
+    "nr": "nbl",
+    "nv": "nav",
+    "ny": "nya",
+    "oc": "oci",
+    "oj": "oji",
+    "om": "orm",
+    "or": "ori",
+    "os": "oss",
+    "pa": "pan",
+    "pi": "pli",
+    "pl": "pol",
+    "ps": "pus",
+    "pt": "por",
+    "qu": "que",
+    "rm": "roh",
+    "rn": "run",
+    "ro": "rum",
+    "ru": "rus",
+    "rw": "kin",
+    "sa": "san",
+    "sc": "srd",
+    "sd": "snd",
+    "se": "sme",
+    "sg": "sag",
+    "si": "sin",
+    "sk": "slo",
+    "sl": "slv",
+    "sm": "smo",
+    "sn": "sna",
+    "so": "som",
+    "sq": "alb",
+    "sr": "srp",
+    "ss": "ssw",
+    "st": "sot",
+    "su": "sun",
+    "sv": "swe",
+    "sw": "swa",
+    "ta": "tam",
+    "te": "tel",
+    "tg": "tgk",
+    "th": "tha",
+    "ti": "tir",
+    "tk": "tuk",
+    "tl": "tgl",
+    "tn": "tsn",
+    "to": "ton",
+    "tr": "tur",
+    "ts": "tso",
+    "tt": "tat",
+    "tw": "twi",
+    "ty": "tah",
+    "ug": "uig",
+    "uk": "ukr",
+    "ur": "urd",
+    "uz": "uzb",
+    "ve": "ven",
+    "vi": "vie",
+    "vo": "vol",
+    "wa": "wln",
+    "wo": "wol",
+    "xh": "xho",
+    "yi": "yid",
+    "yo": "yor",
+    "za": "zha",
+    "zh": "chi",
+    "zu": "zul",
+}
+
+
+def create_config(output_dir):
+    DOMAIN_TYPE = "telephonic"  # Can be meeting, telephonic, or general based on domain type of the audio file
+    CONFIG_LOCAL_DIRECTORY = "nemo_msdd_configs"
+    CONFIG_FILE_NAME = f"diar_infer_{DOMAIN_TYPE}.yaml"
+    MODEL_CONFIG_PATH = os.path.join(CONFIG_LOCAL_DIRECTORY, CONFIG_FILE_NAME)
+    if not os.path.exists(MODEL_CONFIG_PATH):
+        os.makedirs(CONFIG_LOCAL_DIRECTORY, exist_ok=True)
+        CONFIG_URL = f"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/{CONFIG_FILE_NAME}"
+        MODEL_CONFIG_PATH = wget.download(CONFIG_URL, MODEL_CONFIG_PATH)
+
+    config = OmegaConf.load(MODEL_CONFIG_PATH)
+
+    data_dir = os.path.join(output_dir, "data")
+    os.makedirs(data_dir, exist_ok=True)
+
+    meta = {
+        "audio_filepath": os.path.join(output_dir, "mono_file.wav"),
+        "offset": 0,
+        "duration": None,
+        "label": "infer",
+        "text": "-",
+        "rttm_filepath": None,
+        "uem_filepath": None,
+    }
+    with open(os.path.join(data_dir, "input_manifest.json"), "w") as fp:
+        json.dump(meta, fp)
+        fp.write("\n")
+
+    pretrained_vad = "vad_multilingual_marblenet"
+    pretrained_speaker_model = "titanet_large"
+    config.num_workers = 0
+    config.diarizer.manifest_filepath = os.path.join(data_dir, "input_manifest.json")
+    config.diarizer.out_dir = (
+        output_dir  # Directory to store intermediate files and prediction outputs
+    )
+
+    config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model
+    config.diarizer.oracle_vad = (
+        False  # compute VAD provided with model_path to vad config
+    )
+    config.diarizer.clustering.parameters.oracle_num_speakers = False
+
+    # Here, we use our in-house pretrained NeMo VAD model
+    config.diarizer.vad.model_path = pretrained_vad
+    config.diarizer.vad.parameters.onset = 0.8
+    config.diarizer.vad.parameters.offset = 0.6
+    config.diarizer.vad.parameters.pad_offset = -0.05
+    config.diarizer.msdd_model.model_path = (
+        "diar_msdd_telephonic"  # Telephonic speaker diarization model
+    )
+
+    return config
+
+
+def get_word_ts_anchor(s, e, option="start"):
+    if option == "end":
+        return e
+    elif option == "mid":
+        return (s + e) / 2
+    return s
+
+
+def get_words_speaker_mapping(wrd_ts, spk_ts, word_anchor_option="start"):
+    s, e, sp = spk_ts[0]
+    wrd_pos, turn_idx = 0, 0
+    wrd_spk_mapping = []
+    for wrd_dict in wrd_ts:
+        ws, we, wrd = (
+            int(wrd_dict["start"] * 1000),
+            int(wrd_dict["end"] * 1000),
+            wrd_dict["text"],
+        )
+        wrd_pos = get_word_ts_anchor(ws, we, word_anchor_option)
+        while wrd_pos > float(e):
+            turn_idx += 1
+            turn_idx = min(turn_idx, len(spk_ts) - 1)
+            s, e, sp = spk_ts[turn_idx]
+            if turn_idx == len(spk_ts) - 1:
+                e = get_word_ts_anchor(ws, we, option="end")
+        wrd_spk_mapping.append(
+            {"word": wrd, "start_time": ws, "end_time": we, "speaker": sp}
+        )
+    return wrd_spk_mapping
+
+
+sentence_ending_punctuations = ".?!"
+
+
+def get_first_word_idx_of_sentence(word_idx, word_list, speaker_list, max_words):
+    is_word_sentence_end = (
+        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
+    )
+    left_idx = word_idx
+    while (
+        left_idx > 0
+        and word_idx - left_idx < max_words
+        and speaker_list[left_idx - 1] == speaker_list[left_idx]
+        and not is_word_sentence_end(left_idx - 1)
+    ):
+        left_idx -= 1
+
+    return left_idx if left_idx == 0 or is_word_sentence_end(left_idx - 1) else -1
+
+
+def get_last_word_idx_of_sentence(word_idx, word_list, max_words):
+    is_word_sentence_end = (
+        lambda x: x >= 0 and word_list[x][-1] in sentence_ending_punctuations
+    )
+    right_idx = word_idx
+    while (
+        right_idx < len(word_list) - 1
+        and right_idx - word_idx < max_words
+        and not is_word_sentence_end(right_idx)
+    ):
+        right_idx += 1
+
+    return (
+        right_idx
+        if right_idx == len(word_list) - 1 or is_word_sentence_end(right_idx)
+        else -1
+    )
+
+
+def get_realigned_ws_mapping_with_punctuation(
+    word_speaker_mapping, max_words_in_sentence=50
+):
+    is_word_sentence_end = (
+        lambda x: x >= 0
+        and word_speaker_mapping[x]["word"][-1] in sentence_ending_punctuations
+    )
+    wsp_len = len(word_speaker_mapping)
+
+    words_list, speaker_list = [], []
+    for k, line_dict in enumerate(word_speaker_mapping):
+        word, speaker = line_dict["word"], line_dict["speaker"]
+        words_list.append(word)
+        speaker_list.append(speaker)
+
+    k = 0
+    while k < len(word_speaker_mapping):
+        line_dict = word_speaker_mapping[k]
+        if (
+            k < wsp_len - 1
+            and speaker_list[k] != speaker_list[k + 1]
+            and not is_word_sentence_end(k)
+        ):
+            left_idx = get_first_word_idx_of_sentence(
+                k, words_list, speaker_list, max_words_in_sentence
+            )
+            right_idx = (
+                get_last_word_idx_of_sentence(
+                    k, words_list, max_words_in_sentence - k + left_idx - 1
+                )
+                if left_idx > -1
+                else -1
+            )
+            if min(left_idx, right_idx) == -1:
+                k += 1
+                continue
+
+            spk_labels = speaker_list[left_idx : right_idx + 1]
+            mod_speaker = max(set(spk_labels), key=spk_labels.count)
+            if spk_labels.count(mod_speaker) < len(spk_labels) // 2:
+                k += 1
+                continue
+
+            speaker_list[left_idx : right_idx + 1] = [mod_speaker] * (
+                right_idx - left_idx + 1
+            )
+            k = right_idx
+
+        k += 1
+
+    k, realigned_list = 0, []
+    while k < len(word_speaker_mapping):
+        line_dict = word_speaker_mapping[k].copy()
+        line_dict["speaker"] = speaker_list[k]
+        realigned_list.append(line_dict)
+        k += 1
+
+    return realigned_list
+
+
+def get_sentences_speaker_mapping(word_speaker_mapping, spk_ts):
+    sentence_checker = nltk.tokenize.PunktSentenceTokenizer().text_contains_sentbreak
+    s, e, spk = spk_ts[0]
+    prev_spk = spk
+
+    snts = []
+    snt = {"speaker": f"Speaker {spk}", "start_time": s, "end_time": e, "text": ""}
+
+    for wrd_dict in word_speaker_mapping:
+        wrd, spk = wrd_dict["word"], wrd_dict["speaker"]
+        s, e = wrd_dict["start_time"], wrd_dict["end_time"]
+        if spk != prev_spk or sentence_checker(snt["text"] + " " + wrd):
+            snts.append(snt)
+            snt = {
+                "speaker": f"Speaker {spk}",
+                "start_time": s,
+                "end_time": e,
+                "text": "",
+            }
+        else:
+            snt["end_time"] = e
+        snt["text"] += wrd + " "
+        prev_spk = spk
+
+    snts.append(snt)
+    return snts
+
+
+def get_speaker_aware_transcript(sentences_speaker_mapping, f):
+    previous_speaker = sentences_speaker_mapping[0]["speaker"]
+    f.write(f"{previous_speaker}: ")
+
+    for sentence_dict in sentences_speaker_mapping:
+        speaker = sentence_dict["speaker"]
+        sentence = sentence_dict["text"]
+
+        # If this speaker doesn't match the previous one, start a new paragraph
+        if speaker != previous_speaker:
+            f.write(f"\n\n{speaker}: ")
+            previous_speaker = speaker
+
+        # No matter what, write the current sentence
+        f.write(sentence + " ")
+
+
+def format_timestamp(
+    milliseconds: float, always_include_hours: bool = False, decimal_marker: str = "."
+):
+    assert milliseconds >= 0, "non-negative timestamp expected"
+
+    hours = milliseconds // 3_600_000
+    milliseconds -= hours * 3_600_000
+
+    minutes = milliseconds // 60_000
+    milliseconds -= minutes * 60_000
+
+    seconds = milliseconds // 1_000
+    milliseconds -= seconds * 1_000
+
+    hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
+    return (
+        f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
+    )
+
+
+def write_srt(transcript, file):
+    """
+    Write a transcript to a file in SRT format.
+
+    """
+    for i, segment in enumerate(transcript, start=1):
+        # write srt lines
+        print(
+            f"{i}\n"
+            f"{format_timestamp(segment['start_time'], always_include_hours=True, decimal_marker=',')} --> "
+            f"{format_timestamp(segment['end_time'], always_include_hours=True, decimal_marker=',')}\n"
+            f"{segment['speaker']}: {segment['text'].strip().replace('-->', '->')}\n",
+            file=file,
+            flush=True,
+        )
+
+
+def find_numeral_symbol_tokens(tokenizer):
+    numeral_symbol_tokens = [
+        -1,
+    ]
+    for token, token_id in tokenizer.get_vocab().items():
+        has_numeral_symbol = any(c in "0123456789%$£" for c in token)
+        if has_numeral_symbol:
+            numeral_symbol_tokens.append(token_id)
+    return numeral_symbol_tokens
+
+
+def _get_next_start_timestamp(word_timestamps, current_word_index, final_timestamp):
+    # if current word is the last word
+    if current_word_index == len(word_timestamps) - 1:
+        return word_timestamps[current_word_index]["start"]
+
+    next_word_index = current_word_index + 1
+    while current_word_index < len(word_timestamps) - 1:
+        if word_timestamps[next_word_index].get("start") is None:
+            # if next word doesn't have a start timestamp
+            # merge it with the current word and delete it
+            word_timestamps[current_word_index]["word"] += (
+                " " + word_timestamps[next_word_index]["word"]
+            )
+
+            word_timestamps[next_word_index]["word"] = None
+            next_word_index += 1
+            if next_word_index == len(word_timestamps):
+                return final_timestamp
+
+        else:
+            return word_timestamps[next_word_index]["start"]
+
+
+def filter_missing_timestamps(
+    word_timestamps, initial_timestamp=0, final_timestamp=None
+):
+    # handle the first and last word
+    if word_timestamps[0].get("start") is None:
+        word_timestamps[0]["start"] = (
+            initial_timestamp if initial_timestamp is not None else 0
+        )
+        word_timestamps[0]["end"] = _get_next_start_timestamp(
+            word_timestamps, 0, final_timestamp
+        )
+
+    result = [
+        word_timestamps[0],
+    ]
+
+    for i, ws in enumerate(word_timestamps[1:], start=1):
+        # if ws doesn't have a start and end
+        # use the previous end as start and next start as end
+        if ws.get("start") is None and ws.get("word") is not None:
+            ws["start"] = word_timestamps[i - 1]["end"]
+            ws["end"] = _get_next_start_timestamp(word_timestamps, i, final_timestamp)
+
+        if ws["word"] is not None:
+            result.append(ws)
+    return result
+
+
+def cleanup(path: str):
+    """path could either be relative or absolute."""
+    # check if file or directory exists
+    if os.path.isfile(path) or os.path.islink(path):
+        # remove file
+        os.remove(path)
+    elif os.path.isdir(path):
+        # remove directory and all its content
+        shutil.rmtree(path)
+    else:
+        raise ValueError("Path {} is not a file or dir.".format(path))
+
+
+def process_language_arg(language: str, model_name: str):
+    """
+    Process the language argument to make sure it's valid and convert language names to language codes.
+    """
+    if language is not None:
+        language = language.lower()
+    if language not in LANGUAGES:
+        if language in TO_LANGUAGE_CODE:
+            language = TO_LANGUAGE_CODE[language]
+        else:
+            raise ValueError(f"Unsupported language: {language}")
+
+    if model_name.endswith(".en") and language != "en":
+        if language is not None:
+            logging.warning(
+                f"{model_name} is an English-only model but received '{language}'; using English instead."
+            )
+        language = "en"
+    return language

whisper_diarization_main/nemo_msdd_configs/diar_infer_general.yaml

@@ -0,0 +1,95 @@
+# This YAML file is created for all types of offline speaker diarization inference tasks in `<NeMo git root>/example/speaker_tasks/diarization` folder.
+# The inference parameters for VAD, speaker embedding extractor, clustering module, MSDD module, ASR decoder are all included in this YAML file. 
+# All the keys under `diarizer` key (`vad`, `speaker_embeddings`, `clustering`, `msdd_model`, `asr`) can be selectively used for its own purpose and also can be ignored if the module is not used.
+# The configurations in this YAML file is optimized to show balanced performances on various types of domain. VAD is optimized on multilingual ASR datasets and diarizer is optimized on DIHARD3 development set.
+# An example line in an input manifest file (`.json` format):
+# {"audio_filepath": "/path/to/audio_file", "offset": 0, "duration": null, "label": "infer", "text": "-", "num_speakers": null, "rttm_filepath": "/path/to/rttm/file", "uem_filepath": "/path/to/uem/file"}
+name: &name "ClusterDiarizer"
+
+num_workers: 1
+sample_rate: 16000
+batch_size: 64
+device: null # can specify a specific device, i.e: cuda:1 (default cuda if cuda available, else cpu)
+verbose: True # enable additional logging
+
+diarizer:
+  manifest_filepath: ???
+  out_dir: ???
+  oracle_vad: False # If True, uses RTTM files provided in the manifest file to get speech activity (VAD) timestamps
+  collar: 0.25 # Collar value for scoring
+  ignore_overlap: True # Consider or ignore overlap segments while scoring
+
+  vad:
+    model_path: vad_multilingual_marblenet # .nemo local model path or pretrained VAD model name 
+    external_vad_manifest: null # This option is provided to use external vad and provide its speech activity labels for speaker embeddings extraction. Only one of model_path or external_vad_manifest should be set
+
+    parameters: # Tuned by detection error rate (false alarm + miss) on multilingual ASR evaluation datasets
+      window_length_in_sec: 0.63  # Window length in sec for VAD context input 
+      shift_length_in_sec: 0.08 # Shift length in sec for generate frame level VAD prediction
+      smoothing: False # False or type of smoothing method (eg: median)
+      overlap: 0.5 # Overlap ratio for overlapped mean/median smoothing filter
+      onset: 0.5 # Onset threshold for detecting the beginning and end of a speech 
+      offset: 0.3 # Offset threshold for detecting the end of a speech
+      pad_onset: 0.2 # Adding durations before each speech segment 
+      pad_offset: 0.2 # Adding durations after each speech segment 
+      min_duration_on: 0.5 # Threshold for small non_speech deletion
+      min_duration_off: 0.5 # Threshold for short speech segment deletion
+      filter_speech_first: True 
+
+  speaker_embeddings:
+    model_path: titanet_large # .nemo local model path or pretrained model name (titanet_large, ecapa_tdnn or speakerverification_speakernet)
+    parameters:
+      window_length_in_sec: [1.9,1.2,0.5] # Window length(s) in sec (floating-point number). either a number or a list. ex) 1.5 or [1.5,1.0,0.5]
+      shift_length_in_sec: [0.95,0.6,0.25] # Shift length(s) in sec (floating-point number). either a number or a list. ex) 0.75 or [0.75,0.5,0.25]
+      multiscale_weights: [1,1,1] # Weight for each scale. should be null (for single scale) or a list matched with window/shift scale count. ex) [0.33,0.33,0.33]
+      save_embeddings: True # If True, save speaker embeddings in pickle format. This should be True if clustering result is used for other models, such as `msdd_model`.
+  
+  clustering:
+    parameters:
+      oracle_num_speakers: False # If True, use num of speakers value provided in manifest file.
+      max_num_speakers: 8 # Max number of speakers for each recording. If an oracle number of speakers is passed, this value is ignored.
+      enhanced_count_thres: 80 # If the number of segments is lower than this number, enhanced speaker counting is activated.
+      max_rp_threshold: 0.25 # Determines the range of p-value search: 0 < p <= max_rp_threshold. 
+      sparse_search_volume: 10 # The higher the number, the more values will be examined with more time. 
+      maj_vote_spk_count: False  # If True, take a majority vote on multiple p-values to estimate the number of speakers.
+      chunk_cluster_count: 50 # Number of forced clusters (overclustering) per unit chunk in long-form audio clustering.
+      embeddings_per_chunk: 10000 # Number of embeddings in each chunk for long-form audio clustering. Adjust based on GPU memory capacity. (default: 10000, approximately 40 mins of audio) 
+
+  
+  msdd_model:
+    model_path: null  # .nemo local model path or pretrained model name for multiscale diarization decoder (MSDD)
+    parameters:
+      use_speaker_model_from_ckpt: True # If True, use speaker embedding model in checkpoint. If False, the provided speaker embedding model in config will be used.
+      infer_batch_size: 25 # Batch size for MSDD inference. 
+      sigmoid_threshold: [0.7] # Sigmoid threshold for generating binarized speaker labels. The smaller the more generous on detecting overlaps.
+      seq_eval_mode: False # If True, use oracle number of speaker and evaluate F1 score for the given speaker sequences. Default is False.
+      split_infer: True # If True, break the input audio clip to short sequences and calculate cluster average embeddings for inference.
+      diar_window_length: 50 # The length of split short sequence when split_infer is True.
+      overlap_infer_spk_limit: 5 # If the estimated number of speakers are larger than this number, overlap speech is not estimated.
+  
+  asr:
+    model_path: null # Provide NGC cloud ASR model name. stt_en_conformer_ctc_* models are recommended for diarization purposes.
+    parameters:
+      asr_based_vad: False # if True, speech segmentation for diarization is based on word-timestamps from ASR inference.
+      asr_based_vad_threshold: 1.0 # Threshold (in sec) that caps the gap between two words when generating VAD timestamps using ASR based VAD.
+      asr_batch_size: null # Batch size can be dependent on each ASR model. Default batch sizes are applied if set to null.
+      decoder_delay_in_sec: null # Native decoder delay. null is recommended to use the default values for each ASR model.
+      word_ts_anchor_offset: null # Offset to set a reference point from the start of the word. Recommended range of values is [-0.05  0.2]. 
+      word_ts_anchor_pos: "start" # Select which part of the word timestamp we want to use. The options are: 'start', 'end', 'mid'.
+      fix_word_ts_with_VAD: False # Fix the word timestamp using VAD output. You must provide a VAD model to use this feature.
+      colored_text: False # If True, use colored text to distinguish speakers in the output transcript.
+      print_time: True # If True, the start and end time of each speaker turn is printed in the output transcript.
+      break_lines: False # If True, the output transcript breaks the line to fix the line width (default is 90 chars)
+    
+    ctc_decoder_parameters: # Optional beam search decoder (pyctcdecode)
+      pretrained_language_model: null # KenLM model file: .arpa model file or .bin binary file.
+      beam_width: 32
+      alpha: 0.5
+      beta: 2.5
+
+    realigning_lm_parameters: # Experimental feature
+      arpa_language_model: null # Provide a KenLM language model in .arpa format.
+      min_number_of_words: 3 # Min number of words for the left context.
+      max_number_of_words: 10 # Max number of words for the right context.
+      logprob_diff_threshold: 1.2  # The threshold for the difference between two log probability values from two hypotheses.
+

whisper_diarization_main/nemo_msdd_configs/diar_infer_meeting.yaml

@@ -0,0 +1,94 @@
+# This YAML file is created for all types of offline speaker diarization inference tasks in `<NeMo git root>/example/speaker_tasks/diarization` folder.
+# The inference parameters for VAD, speaker embedding extractor, clustering module, MSDD module, ASR decoder are all included in this YAML file. 
+# All the keys under `diarizer` key (`vad`, `speaker_embeddings`, `clustering`, `msdd_model`, `asr`) can be selectively used for its own purpose and also can be ignored if the module is not used.
+# The configurations in this YAML file is suitable for 3~5 speakers participating in a meeting and may not show the best performance on other types of dialogues.
+# An example line in an input manifest file (`.json` format):
+# {"audio_filepath": "/path/to/audio_file", "offset": 0, "duration": null, "label": "infer", "text": "-", "num_speakers": null, "rttm_filepath": "/path/to/rttm/file", "uem_filepath": "/path/to/uem/file"}
+name: &name "ClusterDiarizer"
+
+num_workers: 1
+sample_rate: 16000
+batch_size: 64
+device: null # can specify a specific device, i.e: cuda:1 (default cuda if cuda available, else cpu)
+verbose: True # enable additional logging
+
+diarizer:
+  manifest_filepath: ???
+  out_dir: ???
+  oracle_vad: False # If True, uses RTTM files provided in the manifest file to get speech activity (VAD) timestamps
+  collar: 0.25 # Collar value for scoring
+  ignore_overlap: True # Consider or ignore overlap segments while scoring
+
+  vad:
+    model_path:  vad_multilingual_marblenet # .nemo local model path or pretrained VAD model name 
+    external_vad_manifest: null # This option is provided to use external vad and provide its speech activity labels for speaker embeddings extraction. Only one of model_path or external_vad_manifest should be set
+
+    parameters: # Tuned parameters for CH109 (using the 11 multi-speaker sessions as dev set) 
+      window_length_in_sec: 0.63  # Window length in sec for VAD context input 
+      shift_length_in_sec: 0.01 # Shift length in sec for generate frame level VAD prediction
+      smoothing: False # False or type of smoothing method (eg: median)
+      overlap: 0.5 # Overlap ratio for overlapped mean/median smoothing filter
+      onset: 0.9 # Onset threshold for detecting the beginning and end of a speech 
+      offset: 0.5 # Offset threshold for detecting the end of a speech
+      pad_onset: 0 # Adding durations before each speech segment 
+      pad_offset: 0 # Adding durations after each speech segment 
+      min_duration_on: 0 # Threshold for small non_speech deletion
+      min_duration_off: 0.6 # Threshold for short speech segment deletion
+      filter_speech_first: True 
+
+  speaker_embeddings:
+    model_path: titanet_large # .nemo local model path or pretrained model name (titanet_large, ecapa_tdnn or speakerverification_speakernet)
+    parameters:
+      window_length_in_sec: [3.0,2.5,2.0,1.5,1.0,0.5] # Window length(s) in sec (floating-point number). either a number or a list. ex) 1.5 or [1.5,1.0,0.5]
+      shift_length_in_sec: [1.5,1.25,1.0,0.75,0.5,0.25] # Shift length(s) in sec (floating-point number). either a number or a list. ex) 0.75 or [0.75,0.5,0.25]
+      multiscale_weights: [1,1,1,1,1,1] # Weight for each scale. should be null (for single scale) or a list matched with window/shift scale count. ex) [0.33,0.33,0.33]
+      save_embeddings: True # If True, save speaker embeddings in pickle format. This should be True if clustering result is used for other models, such as `msdd_model`.
+  
+  clustering:
+    parameters:
+      oracle_num_speakers: False # If True, use num of speakers value provided in manifest file.
+      max_num_speakers: 8 # Max number of speakers for each recording. If an oracle number of speakers is passed, this value is ignored.
+      enhanced_count_thres: 80 # If the number of segments is lower than this number, enhanced speaker counting is activated.
+      max_rp_threshold: 0.25 # Determines the range of p-value search: 0 < p <= max_rp_threshold. 
+      sparse_search_volume: 30 # The higher the number, the more values will be examined with more time. 
+      maj_vote_spk_count: False  # If True, take a majority vote on multiple p-values to estimate the number of speakers.
+      chunk_cluster_count: 50 # Number of forced clusters (overclustering) per unit chunk in long-form audio clustering.
+      embeddings_per_chunk: 10000 # Number of embeddings in each chunk for long-form audio clustering. Adjust based on GPU memory capacity. (default: 10000, approximately 40 mins of audio) 
+  
+  msdd_model:
+    model_path: null # .nemo local model path or pretrained model name for multiscale diarization decoder (MSDD)
+    parameters:
+      use_speaker_model_from_ckpt: True # If True, use speaker embedding model in checkpoint. If False, the provided speaker embedding model in config will be used.
+      infer_batch_size: 25 # Batch size for MSDD inference. 
+      sigmoid_threshold: [0.7] # Sigmoid threshold for generating binarized speaker labels. The smaller the more generous on detecting overlaps.
+      seq_eval_mode: False # If True, use oracle number of speaker and evaluate F1 score for the given speaker sequences. Default is False.
+      split_infer: True # If True, break the input audio clip to short sequences and calculate cluster average embeddings for inference.
+      diar_window_length: 50 # The length of split short sequence when split_infer is True.
+      overlap_infer_spk_limit: 5 # If the estimated number of speakers are larger than this number, overlap speech is not estimated.
+  
+  asr:
+    model_path: stt_en_conformer_ctc_large # Provide NGC cloud ASR model name. stt_en_conformer_ctc_* models are recommended for diarization purposes.
+    parameters:
+      asr_based_vad: False # if True, speech segmentation for diarization is based on word-timestamps from ASR inference.
+      asr_based_vad_threshold: 1.0 # Threshold (in sec) that caps the gap between two words when generating VAD timestamps using ASR based VAD.
+      asr_batch_size: null # Batch size can be dependent on each ASR model. Default batch sizes are applied if set to null.
+      decoder_delay_in_sec: null # Native decoder delay. null is recommended to use the default values for each ASR model.
+      word_ts_anchor_offset: null # Offset to set a reference point from the start of the word. Recommended range of values is [-0.05  0.2]. 
+      word_ts_anchor_pos: "start" # Select which part of the word timestamp we want to use. The options are: 'start', 'end', 'mid'.
+      fix_word_ts_with_VAD: False # Fix the word timestamp using VAD output. You must provide a VAD model to use this feature.
+      colored_text: False # If True, use colored text to distinguish speakers in the output transcript.
+      print_time: True # If True, the start and end time of each speaker turn is printed in the output transcript.
+      break_lines: False # If True, the output transcript breaks the line to fix the line width (default is 90 chars)
+    
+    ctc_decoder_parameters: # Optional beam search decoder (pyctcdecode)
+      pretrained_language_model: null # KenLM model file: .arpa model file or .bin binary file.
+      beam_width: 32
+      alpha: 0.5
+      beta: 2.5
+
+    realigning_lm_parameters: # Experimental feature
+      arpa_language_model: null # Provide a KenLM language model in .arpa format.
+      min_number_of_words: 3 # Min number of words for the left context.
+      max_number_of_words: 10 # Max number of words for the right context.
+      logprob_diff_threshold: 1.2  # The threshold for the difference between two log probability values from two hypotheses.
+

whisper_diarization_main/nemo_msdd_configs/diar_infer_telephonic.yaml

@@ -0,0 +1,94 @@
+# This YAML file is created for all types of offline speaker diarization inference tasks in `<NeMo git root>/example/speaker_tasks/diarization` folder.
+# The inference parameters for VAD, speaker embedding extractor, clustering module, MSDD module, ASR decoder are all included in this YAML file. 
+# All the keys under `diarizer` key (`vad`, `speaker_embeddings`, `clustering`, `msdd_model`, `asr`) can be selectively used for its own purpose and also can be ignored if the module is not used.
+# The configurations in this YAML file is suitable for telephone recordings involving 2~8 speakers in a session and may not show the best performance on the other types of acoustic conditions or dialogues.
+# An example line in an input manifest file (`.json` format):
+# {"audio_filepath": "/path/to/audio_file", "offset": 0, "duration": null, "label": "infer", "text": "-", "num_speakers": null, "rttm_filepath": "/path/to/rttm/file", "uem_filepath": "/path/to/uem/file"}
+name: &name "ClusterDiarizer"
+
+num_workers: 1
+sample_rate: 16000
+batch_size: 64
+device: null # can specify a specific device, i.e: cuda:1 (default cuda if cuda available, else cpu)
+verbose: True # enable additional logging
+
+diarizer:
+  manifest_filepath: ???
+  out_dir: ???
+  oracle_vad: False # If True, uses RTTM files provided in the manifest file to get speech activity (VAD) timestamps
+  collar: 0.25 # Collar value for scoring
+  ignore_overlap: True # Consider or ignore overlap segments while scoring
+
+  vad:
+    model_path: vad_multilingual_marblenet # .nemo local model path or pretrained VAD model name 
+    external_vad_manifest: null # This option is provided to use external vad and provide its speech activity labels for speaker embeddings extraction. Only one of model_path or external_vad_manifest should be set
+
+    parameters: # Tuned parameters for CH109 (using the 11 multi-speaker sessions as dev set) 
+      window_length_in_sec: 0.15  # Window length in sec for VAD context input 
+      shift_length_in_sec: 0.01 # Shift length in sec for generate frame level VAD prediction
+      smoothing: "median" # False or type of smoothing method (eg: median)
+      overlap: 0.5 # Overlap ratio for overlapped mean/median smoothing filter
+      onset: 0.1 # Onset threshold for detecting the beginning and end of a speech 
+      offset: 0.1 # Offset threshold for detecting the end of a speech
+      pad_onset: 0.1 # Adding durations before each speech segment 
+      pad_offset: 0 # Adding durations after each speech segment 
+      min_duration_on: 0 # Threshold for small non_speech deletion
+      min_duration_off: 0.2 # Threshold for short speech segment deletion
+      filter_speech_first: True 
+
+  speaker_embeddings:
+    model_path: titanet_large # .nemo local model path or pretrained model name (titanet_large, ecapa_tdnn or speakerverification_speakernet)
+    parameters:
+      window_length_in_sec: [1.5,1.25,1.0,0.75,0.5] # Window length(s) in sec (floating-point number). either a number or a list. ex) 1.5 or [1.5,1.0,0.5]
+      shift_length_in_sec: [0.75,0.625,0.5,0.375,0.25] # Shift length(s) in sec (floating-point number). either a number or a list. ex) 0.75 or [0.75,0.5,0.25]
+      multiscale_weights: [1,1,1,1,1] # Weight for each scale. should be null (for single scale) or a list matched with window/shift scale count. ex) [0.33,0.33,0.33]
+      save_embeddings: True # If True, save speaker embeddings in pickle format. This should be True if clustering result is used for other models, such as `msdd_model`.
+  
+  clustering:
+    parameters:
+      oracle_num_speakers: False # If True, use num of speakers value provided in manifest file.
+      max_num_speakers: 8 # Max number of speakers for each recording. If an oracle number of speakers is passed, this value is ignored.
+      enhanced_count_thres: 80 # If the number of segments is lower than this number, enhanced speaker counting is activated.
+      max_rp_threshold: 0.25 # Determines the range of p-value search: 0 < p <= max_rp_threshold. 
+      sparse_search_volume: 30 # The higher the number, the more values will be examined with more time. 
+      maj_vote_spk_count: False  # If True, take a majority vote on multiple p-values to estimate the number of speakers.
+      chunk_cluster_count: 50 # Number of forced clusters (overclustering) per unit chunk in long-form audio clustering.
+      embeddings_per_chunk: 10000 # Number of embeddings in each chunk for long-form audio clustering. Adjust based on GPU memory capacity. (default: 10000, approximately 40 mins of audio) 
+  
+  msdd_model:
+    model_path: diar_msdd_telephonic # .nemo local model path or pretrained model name for multiscale diarization decoder (MSDD)
+    parameters:
+      use_speaker_model_from_ckpt: True # If True, use speaker embedding model in checkpoint. If False, the provided speaker embedding model in config will be used.
+      infer_batch_size: 25 # Batch size for MSDD inference.
+      sigmoid_threshold: [0.7] # Sigmoid threshold for generating binarized speaker labels. The smaller the more generous on detecting overlaps.
+      seq_eval_mode: False # If True, use oracle number of speaker and evaluate F1 score for the given speaker sequences. Default is False.
+      split_infer: True # If True, break the input audio clip to short sequences and calculate cluster average embeddings for inference.
+      diar_window_length: 50 # The length of split short sequence when split_infer is True.
+      overlap_infer_spk_limit: 5 # If the estimated number of speakers are larger than this number, overlap speech is not estimated.
+  
+  asr:
+    model_path: stt_en_conformer_ctc_large # Provide NGC cloud ASR model name. stt_en_conformer_ctc_* models are recommended for diarization purposes.
+    parameters:
+      asr_based_vad: False # if True, speech segmentation for diarization is based on word-timestamps from ASR inference.
+      asr_based_vad_threshold: 1.0 # Threshold (in sec) that caps the gap between two words when generating VAD timestamps using ASR based VAD.
+      asr_batch_size: null # Batch size can be dependent on each ASR model. Default batch sizes are applied if set to null.
+      decoder_delay_in_sec: null # Native decoder delay. null is recommended to use the default values for each ASR model.
+      word_ts_anchor_offset: null # Offset to set a reference point from the start of the word. Recommended range of values is [-0.05  0.2]. 
+      word_ts_anchor_pos: "start" # Select which part of the word timestamp we want to use. The options are: 'start', 'end', 'mid'.
+      fix_word_ts_with_VAD: False # Fix the word timestamp using VAD output. You must provide a VAD model to use this feature.
+      colored_text: False # If True, use colored text to distinguish speakers in the output transcript.
+      print_time: True # If True, the start and end time of each speaker turn is printed in the output transcript.
+      break_lines: False # If True, the output transcript breaks the line to fix the line width (default is 90 chars)
+    
+    ctc_decoder_parameters: # Optional beam search decoder (pyctcdecode)
+      pretrained_language_model: null # KenLM model file: .arpa model file or .bin binary file.
+      beam_width: 32
+      alpha: 0.5
+      beta: 2.5
+
+    realigning_lm_parameters: # Experimental feature
+      arpa_language_model: null # Provide a KenLM language model in .arpa format.
+      min_number_of_words: 3 # Min number of words for the left context.
+      max_number_of_words: 10 # Max number of words for the right context.
+      logprob_diff_threshold: 1.2  # The threshold for the difference between two log probability values from two hypotheses.
+

whisper_diarization_main/nemo_process.py

@@ -0,0 +1,31 @@
+import argparse
+import os
+
+import torch
+from nemo.collections.asr.models.msdd_models import NeuralDiarizer
+from pydub import AudioSegment
+
+from helpers import create_config
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-a", "--audio", help="name of the target audio file", required=True
+)
+parser.add_argument(
+    "--device",
+    dest="device",
+    default="cuda" if torch.cuda.is_available() else "cpu",
+    help="if you have a GPU use 'cuda', otherwise 'cpu'",
+)
+args = parser.parse_args()
+
+# convert audio to mono for NeMo combatibility
+sound = AudioSegment.from_file(args.audio).set_channels(1)
+ROOT = os.getcwd()
+temp_path = os.path.join(ROOT, "temp_outputs")
+os.makedirs(temp_path, exist_ok=True)
+sound.export(os.path.join(temp_path, "mono_file.wav"), format="wav")
+
+# Initialize NeMo MSDD diarization model
+msdd_model = NeuralDiarizer(cfg=create_config(temp_path)).to(args.device)
+msdd_model.diarize()

whisper_diarization_main/requirements.txt

@@ -0,0 +1,7 @@
+wget
+nemo_toolkit[asr]==1.23.0
+git+https://github.com/m-bain/whisperX.git@78dcfaab51005aa703ee21375f81ed31bc248560
+git+https://github.com/adefossez/demucs.git
+git+https://github.com/oliverguhr/deepmultilingualpunctuation.git
+git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git
+

whisper_diarization_main/transcription_helpers.py

@@ -0,0 +1,74 @@
+import torch
+
+
+def transcribe(
+    audio_file: str,
+    language: str,
+    model_name: str,
+    compute_dtype: str,
+    suppress_numerals: bool,
+    device: str,
+):
+    from faster_whisper import WhisperModel
+
+    from helpers import find_numeral_symbol_tokens, wav2vec2_langs
+
+    # Faster Whisper non-batched
+    # Run on GPU with FP16
+    whisper_model = WhisperModel(model_name, device=device, compute_type=compute_dtype)
+
+    # or run on GPU with INT8
+    # model = WhisperModel(model_size, device="cuda", compute_type="int8_float16")
+    # or run on CPU with INT8
+    # model = WhisperModel(model_size, device="cpu", compute_type="int8")
+
+    if suppress_numerals:
+        numeral_symbol_tokens = find_numeral_symbol_tokens(whisper_model.hf_tokenizer)
+    else:
+        numeral_symbol_tokens = None
+
+    if language is not None and language in wav2vec2_langs:
+        word_timestamps = False
+    else:
+        word_timestamps = True
+
+    segments, info = whisper_model.transcribe(
+        audio_file,
+        language=language,
+        beam_size=5,
+        word_timestamps=word_timestamps,  # TODO: disable this if the language is supported by wav2vec2
+        suppress_tokens=numeral_symbol_tokens,
+        vad_filter=True,
+    )
+    whisper_results = []
+    for segment in segments:
+        whisper_results.append(segment._asdict())
+    # clear gpu vram
+    del whisper_model
+    torch.cuda.empty_cache()
+    return whisper_results, info.language
+
+
+def transcribe_batched(
+    audio_file: str,
+    language: str,
+    batch_size: int,
+    model_name: str,
+    compute_dtype: str,
+    suppress_numerals: bool,
+    device: str,
+):
+    import whisperx
+
+    # Faster Whisper batched
+    whisper_model = whisperx.load_model(
+        model_name,
+        device,
+        compute_type=compute_dtype,
+        asr_options={"suppress_numerals": suppress_numerals},
+    )
+    audio = whisperx.load_audio(audio_file)
+    result = whisper_model.transcribe(audio, language=language, batch_size=batch_size)
+    del whisper_model
+    torch.cuda.empty_cache()
+    return result["segments"], result["language"], audio