Whisper_Transcription_+_NeMo_Diarization.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "eCmjcOc9yEtQ",
        "jbsUt3SwyhjD"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jarvisx17/ASR/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>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Installing Dependencies"
      ],
      "metadata": {
        "id": "eCmjcOc9yEtQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tn1c-CoDv2kw"
      },
      "outputs": [],
      "source": [
        "!pip install torch==1.13.1+cu116 torchaudio==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116\n",
        "!pip install torchtext==0.14.1\n",
        "!pip install nemo_toolkit[asr]==1.15.0\n",
        "!pip install git+https://github.com/openai/whisper.git\n",
        "!pip install git+https://github.com/m-bain/whisperX.git@4cb167a225c0ebaea127fd6049abfaa3af9f8bb4\n",
        "!pip install git+https://github.com/facebookresearch/demucs#egg=demucs\n",
        "!pip install deepmultilingualpunctuation\n",
        "!pip install pytorch-lightning==1.8.6\n",
        "!pip install transformers==4.26.1"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "5q0BFjAVzKuo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import wget\n",
        "from omegaconf import OmegaConf\n",
        "import json\n",
        "import shutil\n",
        "from whisper import load_model\n",
        "import whisperx\n",
        "import torch\n",
        "import librosa\n",
        "import soundfile\n",
        "from nemo.collections.asr.models.msdd_models import NeuralDiarizer\n",
        "from deepmultilingualpunctuation import PunctuationModel\n",
        "import re"
      ],
      "metadata": {
        "id": "YzhncHP0ytbQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Helper Functions"
      ],
      "metadata": {
        "id": "jbsUt3SwyhjD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def create_config():\n",
        "    data_dir = \"./\"\n",
        "    DOMAIN_TYPE = \"telephonic\"  # Can be meeting or telephonic 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(data_dir, CONFIG_FILE_NAME)\n",
        "    if not os.path.exists(MODEL_CONFIG):\n",
        "        MODEL_CONFIG = wget.download(CONFIG_URL, data_dir)\n",
        "\n",
        "    config = OmegaConf.load(MODEL_CONFIG)\n",
        "\n",
        "    ROOT = os.getcwd()\n",
        "    data_dir = os.path.join(ROOT, \"data\")\n",
        "    os.makedirs(data_dir, exist_ok=True)\n",
        "\n",
        "    meta = {\n",
        "        \"audio_filepath\": \"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(\"data/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",
        "\n",
        "    config.num_workers = 1  # Workaround for multiprocessing hanging with ipython issue\n",
        "\n",
        "    output_dir = \"nemo_outputs\"  # os.path.join(ROOT, 'outputs')\n",
        "    os.makedirs(output_dir, exist_ok=True)\n",
        "    config.diarizer.manifest_filepath = \"data/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)\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",
        "    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:\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",
        "    for sentence_dict in sentences_speaker_mapping:\n",
        "        sp = sentence_dict[\"speaker\"]\n",
        "        text = sentence_dict[\"text\"]\n",
        "        f.write(f\"\\n\\n{sp}: {text}\")\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",
        "    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 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",
        "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",
        "wav2vec2_langs = [\n",
        "    \"en\",\n",
        "    \"fr\",\n",
        "    \"de\",\n",
        "    \"es\",\n",
        "    \"it\",\n",
        "    \"nl\",\n",
        "    \"pt\",\n",
        "    \"ja\",\n",
        "    \"zh\",\n",
        "    \"uk\",\n",
        "    \"pt\",\n",
        "    \"ar\",\n",
        "    \"ru\",\n",
        "    \"pl\",\n",
        "    \"hu\",\n",
        "    \"fi\",\n",
        "    \"fa\",\n",
        "    \"el\",\n",
        "    \"tr\",\n",
        "]"
      ],
      "metadata": {
        "id": "Se6Hc7CZygxu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Options"
      ],
      "metadata": {
        "id": "B7qWQb--1Xcw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Name of the audio file\n",
        "audio_path = '9d54248d-60f5-4661-97d1-88b23568b2db.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')\n",
        "whisper_model_name = 'medium.en'"
      ],
      "metadata": {
        "id": "ONlFrSnD0FOp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Processing"
      ],
      "metadata": {
        "id": "h-cY1ZEy2KVI"
      }
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "7ZS4xXmE2NGP"
      }
    },
    {
      "cell_type": "code",
      "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_ft --two-stems=vocals \"{audio_path}\" -o \"temp_outputs\"'\n",
        "    )\n",
        "\n",
        "    if return_code != 0:\n",
        "        print(\n",
        "            \"Source splitting failed, using original audio file.\"\n",
        "        )\n",
        "        vocal_target = audio_path\n",
        "    else:\n",
        "        vocal_target = f\"temp_outputs/htdemucs_ft/{audio_path[:-4]}/vocals.wav\"\n",
        "else:\n",
        "    vocal_target = audio_path"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HKcgQUrAzsJZ",
        "outputId": "dc2a1d96-20da-4749-9d64-21edacfba1b1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Source splitting failed, using original audio file. Use --no-stem argument to disable it.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "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",
        "\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."
      ],
      "metadata": {
        "id": "UYg9VWb22Tz8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Large models result in considerably better and more aligned (words, timestamps) mapping.\n",
        "whisper_model = load_model(whisper_model_name)\n",
        "whisper_results = whisper_model.transcribe(vocal_target, beam_size=None, verbose=False)\n",
        "\n",
        "# clear gpu vram\n",
        "del whisper_model\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "device = \"cuda\"\n",
        "alignment_model, metadata = whisperx.load_align_model(\n",
        "    language_code=whisper_results[\"language\"], device=device\n",
        ")\n",
        "result_aligned = whisperx.align(\n",
        "    whisper_results[\"segments\"], alignment_model, metadata, vocal_target, device\n",
        ")\n",
        "\n",
        "# clear gpu vram\n",
        "del alignment_model\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "5-VKFn530oTl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Convert audio to mono for NeMo combatibility"
      ],
      "metadata": {
        "id": "7EEaJPsQ21Rx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "signal, sample_rate = librosa.load(vocal_target, sr=None)\n",
        "ROOT = os.getcwd()\n",
        "temp_path = os.path.join(ROOT, \"temp_outputs\")\n",
        "if not os.path.exists(temp_path):\n",
        "    os.mkdir(temp_path)\n",
        "os.chdir(temp_path)\n",
        "soundfile.write(\"mono_file.wav\", signal, sample_rate, \"PCM_24\")"
      ],
      "metadata": {
        "id": "rVPuY1VS0xN7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "D1gkViCf2-CV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Initialize NeMo MSDD diarization model\n",
        "msdd_model = NeuralDiarizer(cfg=create_config())\n",
        "msdd_model.diarize()\n",
        "\n",
        "del msdd_model\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "C7jIpBCH02RL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Mapping Spekers to Sentences According to Timestamps"
      ],
      "metadata": {
        "id": "NmkZYaDAEOAg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Reading timestamps <> Speaker Labels mapping\n",
        "\n",
        "output_dir = \"nemo_outputs\"\n",
        "\n",
        "speaker_ts = []\n",
        "with open(f\"{output_dir}/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(result_aligned[\"word_segments\"], speaker_ts, \"start\")"
      ],
      "metadata": {
        "id": "E65LUGQe02zw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "8Ruxc8S1EXtW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "if whisper_results[\"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)\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",
        "\n",
        "\n",
        "    wsm = get_realigned_ws_mapping_with_punctuation(wsm)\n",
        "else:\n",
        "    print(\n",
        "        f'Punctuation restoration is not available for {whisper_results[\"language\"]} language.'\n",
        "    )\n",
        "\n",
        "ssm = get_sentences_speaker_mapping(wsm, speaker_ts)"
      ],
      "metadata": {
        "id": "pgfC5hA41BXu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Cleanup and Exporing the results"
      ],
      "metadata": {
        "id": "vF2QAtLOFvwZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "os.chdir(ROOT)  # back to parent dir\n",
        "with open(f\"{audio_path[:-4]}.txt\", \"w\", encoding=\"utf-8-sig\") as f:\n",
        "    get_speaker_aware_transcript(ssm, f)\n",
        "\n",
        "with open(f\"{audio_path[:-4]}.srt\", \"w\", encoding=\"utf-8-sig\") as srt:\n",
        "    write_srt(ssm, srt)\n",
        "\n",
        "cleanup(temp_path)"
      ],
      "metadata": {
        "id": "kFTyKI6B1MI0"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}