asr-speaker-diarization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2bf1191b",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/rivaasrasr-wordboosting/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How do I Use Speaker Diarization with Riva ASR?\n",
    "\n",
    "This tutorial walks you through the speaker diarization feature available with Riva ASR."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcfba84d",
   "metadata": {},
   "source": [
    "## NVIDIA Riva Overview\n",
    "\n",
    "NVIDIA Riva is a GPU-accelerated SDK for building speech AI applications that are customized for your use case and deliver real-time performance. <br/>\n",
    "Riva offers a rich set of speech and natural language understanding services such as:\n",
    "\n",
    "- Automated speech recognition (ASR)\n",
    "- Text-to-Speech synthesis (TTS)\n",
    "- A collection of natural language processing (NLP) services, such as named entity recognition (NER), punctuation, and intent classification.\n",
    "\n",
    "In this tutorial, we will show how to use the speaker diarization feature with Riva ASR to get the transcript with each word tagged with the id of the speaker who has spoken that word. <br> \n",
    "To understand the basics of Riva ASR APIs, refer to [Getting started with Riva ASR in Python](https://github.com/nvidia-riva/tutorials/tree/stable/asr-python-basics.ipynb). <br>\n",
    "\n",
    "For more information about Riva, refer to the [Riva developer documentation](https://developer.nvidia.com/riva)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27f37136",
   "metadata": {},
   "source": [
    "## Speaker Diarization with Riva ASR APIs\n",
    "\n",
    "[Speaker Diarization](https://en.wikipedia.org/wiki/Speaker_diarisation) is the process of segmenting audio recordings by speaker labels and aims to answer the question “who spoke when?” Riva ASR supports speaker diarization, which can be enabled by passing [`SpeakerDiarizationConfig`](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/protos.html#_CPPv424SpeakerDiarizationConfig) with `enable_speaker_diarization` set to `True`. Riva speaker diarization segments input audio, extracts speaker embeddings of the segments, counts the number of speakers, and then assigns each audio segment a corresponding speaker tag. When speaker diarization is enabled, Riva ASR returns the ASR transcript to the client, along with a speaker tag for each word in the transcript. Speaker diarization is language agnostic and can work with any language supported by Riva ASR.\n",
    "\n",
    "### Note:\n",
    "Speaker diarization is only supported with Riva ASR offline API. Speaker diarization is an alpha release and will increase ASR latency once the feature is enabled. Refer to [Performance](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-performance.html) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3304f0f0",
   "metadata": {},
   "source": [
    "# Requirements and Setup\n",
    "\n",
    "1. Enable the speaker diarization model.  \n",
    "Speaker Diarization is an optional model, so it needs to be enabled in `config.sh` by unncommenting the line containing the `rmir_diarizer_offline` word. Since speaker diarization works only with Riva ASR offline API, make sure that the offline ASR model is also enabled in `config.sh`.\n",
    "\n",
    "2. Deploy the models and start the Riva Speech Skills server.  \n",
    "Deploy the models enabled in the previous step by running `bash riva_init.sh` and then start the Riva server by running `bash riva_start.sh`. Refer to the [Riva Skills Quick Start Guide](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/quick-start-guide.html) for more information.\n",
    "\n",
    "3. Install the Riva client library.  \n",
    "Perform the steps in the [requirements and setup for the Riva client](https://github.com/nvidia-riva/tutorials#requirements) section to install the Riva client library."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41f7589e",
   "metadata": {},
   "source": [
    "#### Import the Riva Client Libraries\n",
    "\n",
    "Let's import some of the required libraries, including the Riva client libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73626be6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import IPython.display as ipd\n",
    "import grpc\n",
    "\n",
    "import riva.client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0ab5f9f",
   "metadata": {},
   "source": [
    "#### Create a Riva Client and Connect to the Riva Speech API Server\n",
    "\n",
    "The following URI assumes a local deployment of the Riva Speech API server is on the default port. In case the server deployment is on a different host or via a Helm chart on Kubernetes, use an appropriate URI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "244caad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate client\n",
    "auth = riva.client.Auth(uri='localhost:50051')\n",
    "riva_asr = riva.client.ASRService(auth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02154940",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Load a sample audio file from local disk\n",
    "# This example uses a .wav file with LINEAR_PCM encoding.\n",
    "# Sample file taken from https://freesound.org/people/SamKolber/sounds/203020/\n",
    "path = \"audio_samples/interview-with-bill.wav\"\n",
    "with io.open(path, 'rb') as fh:\n",
    "    content = fh.read()\n",
    "ipd.Audio(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0515ac0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating RecognitionConfig\n",
    "config = riva.client.RecognitionConfig(\n",
    "  language_code=\"en-US\",\n",
    "  max_alternatives=1,\n",
    "  enable_automatic_punctuation=True,\n",
    "  enable_word_time_offsets=True,\n",
    ")\n",
    "\n",
    "# Use utility function to add SpeakerDiarizationConfig with enable_speaker_diarization=True\n",
    "# Value of max_speaker_count in SpeakerDiarizationConfig has no effect as of now. It will be honoured in future.\n",
    "riva.client.asr.add_speaker_diarization_to_config(config, diarization_enable=True)\n",
    "\n",
    "# ASR inference call with Recognize\n",
    "response = riva_asr.offline_recognize(content, config)\n",
    "print(\"ASR Transcript with Speaker Diarization:\\n\", response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df24b66d",
   "metadata": {},
   "source": [
    "The ASR transcript is split into multiple results based on speech pauses. For each result, every word in the transcript is assigned a speaker tag indicating which speaker spoke that word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda5c652",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pretty print transcript with color coded speaker tags. Black color text indicates no speaker tag was assigned.\n",
    "for result in response.results:\n",
    "    for word in result.alternatives[0].words:\n",
    "        color = '\\033['+ str(30 + word.speaker_tag) + 'm'\n",
    "        print(color, word.word, end=\"\")\n",
    "      "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb21c265",
   "metadata": {},
   "source": [
    "This completes the tutorial for using speaker diarization with Riva ASR."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64c36fe7",
   "metadata": {},
   "source": [
    "## Go Deeper into Riva Capabilities\n",
    "\n",
    "Now that you have a basic introduction to the Riva ASR APIs, you can try:\n",
    "\n",
    "### Additional Riva Tutorials\n",
    "\n",
    "Checkout more Riva tutorials [here](https://github.com/nvidia-riva/tutorials) to understand how to use some of the advanced features of Riva ASR, including customizing ASR for your specific needs.\n",
    "\n",
    "\n",
    "### Sample Applications\n",
    "\n",
    "Riva comes with various sample applications. They demonstrate how to use the APIs to build various applications. Refer to [Riva Sampple Apps](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/samples/index.html) for more information.  \n",
    "\n",
    "\n",
    "### Additional Resources\n",
    "\n",
    "For more information about each of the Riva APIs and their functionalities, refer to the [documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/protos.html)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}