{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/rivaasrasr-basics/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
        "\n",
        "# NVIDIA ASR NIM Tutorial\n",
        "\n",
        "This tutorial walks you through the various features of NVIDIA ASR NIM and how to use the APIs in a Python application. NVIDIA ASR NIM uses the gRPC API to serve offline and online use cases."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Prerequisites\n",
        "\n",
        "1. Deploy NVIDIA ASR NIM with the Parakeet 1.1b en-US model by following the [NVIDIA ASR NIM](https://docs.nvidia.com/nim/speech/latest/asr/index.html) documentation.\n",
        "2. Install the Riva Python Client library:\n",
        "    ```bash\n",
        "    sudo apt-get install python3-pip\n",
        "    pip install -U nvidia-riva-client\n",
        "    ```\n",
        "3. Clone the Git repository at https://github.com/nvidia-riva/tutorials for audio samples. The repository is assumed to be cloned in the `$HOME` directory."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Offline Recognition\n",
        "\n",
        "In offline transcription, the entire input speech is submitted to the service, and the final transcript is received in one response.\n",
        "\n",
        "Try generating the transcripts using the Riva ASR APIs for some sample audio clips in English."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Begin by importing some of the necessary libraries, including the Riva Client libraries."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Import required libraries\n",
        "import io\n",
        "from pathlib import Path\n",
        "import grpc\n",
        "import riva.client\n",
        "import IPython.display as ipd\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The following URI assumes a local deployment of the NVIDIA ASR NIM 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",
      "metadata": {},
      "source": [
        "# Create a Riva client and connect to the Riva ASR NIM\n",
        "auth = riva.client.Auth(uri='0.0.0.0:50051')\n",
        "asr_service = riva.client.ASRService(auth)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Get list of available ASR models for offline use case\n",
        "print(\"Available ASR models for offline use case\")\n",
        "config_response = asr_service.stub.GetRivaSpeechRecognitionConfig(riva.client.proto.riva_asr_pb2.RivaSpeechRecognitionConfigRequest())\n",
        "for model_config in config_response.model_config:\n",
        "    if model_config.parameters[\"type\"] == \"offline\":\n",
        "        print(f\"{model_config.parameters['language_code']} : {model_config.model_name}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "NVIDIA ASR NIM supports 16-bit, single-channel audio in `LPCM`, `alaw`, and `ulaw` encoding in `.raw` (headerless) and `.wav` format. It also supports `.opus` and `.flac` formats. File format is auto-detected from the provided input audio."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": true
      },
      "source": [
        "# This example uses a .wav file with LINEAR_PCM encoding.\n",
        "# read in an audio file from local disk\n",
        "path = Path(\"~/tutorials/audio_samples/en-US_sample.wav\").expanduser()\n",
        "with io.open(path, 'rb') as fh:\n",
        "    content = fh.read()\n",
        "ipd.Audio(path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Setup configuration parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Set up an recognition config\n",
        "config = riva.client.RecognitionConfig()\n",
        "config.language_code = \"en-US\"                    # Language code of the audio clip\n",
        "config.max_alternatives = 1                       # How many top-N hypotheses to return. Only value of 1 is supported.\n",
        "config.enable_automatic_punctuation = True        # Enable punctuation and capitalization\n",
        "config.audio_channel_count = 1                    # Mono - Default\n",
        "config.verbatim_transcripts = False               # Set to True to return verbatim transcripts\n",
        "config.profanity_filter = False                   # Set to True to filter and replace profane words with first letter followed by asterisks (e.g. \"f***\")\n",
        "\n",
        "# In cases where audio samples are submitted in `.raw` format, you need to set the following parameters appropriately:\n",
        "# config.encoding = riva.client.AudioEncoding.LINEAR_PCM\n",
        "# config.sample_rate_hertz = 16000\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Submit the request to the server and print response."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": false
      },
      "source": [
        "# Make a gRPC request and invoke the ASR service\n",
        "request = riva.client.proto.riva_asr_pb2.RecognizeRequest(config=config, audio=content)\n",
        "response = asr_service.stub.Recognize(request)\n",
        "\n",
        "# Full response shows additional information like word time offsets and\n",
        "# word confidence along with the transcript.\n",
        "print(f\"Full Response: {response}\")\n",
        "\n",
        "# Print the final transcript by combining transcripts from all results\n",
        "final_transcript = \"\"\n",
        "for res in response.results:\n",
        "    final_transcript += res.alternatives[0].transcript\n",
        "print(\"Final transcript:\", final_transcript)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Speaker Diarization\n",
        "\n",
        "Profiles with `diarizer=sortformer` use the [Sortformer model](https://huggingface.co/nvidia/diar_sortformer_4spk-v1) for speaker diarization."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import riva.client.proto.riva_asr_pb2 as riva_asr_pb2\n",
        "\n",
        "# This example uses a .wav file with LINEAR_PCM encoding with 2 speakers\n",
        "# read in an audio file from local disk\n",
        "path = Path(\"~/tutorials/audio_samples/2_speaker_en_US.wav\").expanduser()\n",
        "with io.open(path, 'rb') as fh:\n",
        "    content = fh.read()\n",
        "\n",
        "# Set up an recognition config\n",
        "config.enable_word_time_offsets = True            # Enable word timestamps since they're needed for diarization\n",
        "\n",
        "# Add speaker diarization configuration\n",
        "riva.client.asr.add_speaker_diarization_to_config(config, diarization_enable=True, diarization_max_speakers=8)\n",
        "\n",
        "# Make a gRPC request and invoke the ASR service\n",
        "request = riva.client.proto.riva_asr_pb2.RecognizeRequest(config=config, audio=content)\n",
        "response = asr_service.stub.Recognize(request)\n",
        "\n",
        "# Full response shows additional information like word time offsets and\n",
        "# word confidence along with the transcript.\n",
        "print(f\"Response with Speaker Diarization: {response}\")\n",
        "\n",
        "# Print the final transcript per speaker\n",
        "speaker = \"\"\n",
        "for result in response.results:\n",
        "    for word in result.alternatives[0].words:\n",
        "        if speaker != word.speaker_tag:\n",
        "            speaker = word.speaker_tag\n",
        "            print(f\"\\nSpeaker {speaker}: \", end=\"\")\n",
        "        print(word.word, end=\" \")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Offline recognition for non-English languages\n",
        "\n",
        "You can run NVIDIA ASR NIM for non-English languages in the same manner by setting `config.language_code` to the appropriate language code. NVIDIA ASR NIM must be deployed with the model for the required language. For a list of available models, refer to [Supported Models](https://docs.nvidia.com/nim/speech/latest/reference/support-matrix/asr.html#supported-models) in the NVIDIA ASR NIM documentation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Streaming Recognition\n",
        "\n",
        "In case of streaming transcription, the input speech is submitted to the service in chunks, and the final transcript is received incrementally as the input speech is processed.\n",
        "\n",
        "Try generating the transcripts using the Riva ASR APIs for some sample audio clips in English."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Begin by importing some of the necessary libraries, including the Riva Client libraries."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Import required libraries\n",
        "import io\n",
        "from pathlib import Path\n",
        "import grpc\n",
        "import riva.client\n",
        "import IPython.display as ipd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The following URI assumes a local deployment of the NVIDIA ASR NIM 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",
      "metadata": {},
      "source": [
        "# Create a Riva client and connect to the Riva ASR NIM\n",
        "auth = riva.client.Auth(uri='0.0.0.0:50051')\n",
        "asr_service = riva.client.ASRService(auth)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Get list of available ASR models for streaming/online use case\n",
        "print(\"Available ASR models for streaming/online use case\")\n",
        "config_response = asr_service.stub.GetRivaSpeechRecognitionConfig(riva.client.proto.riva_asr_pb2.RivaSpeechRecognitionConfigRequest())\n",
        "for model_config in config_response.model_config:\n",
        "    if model_config.parameters[\"type\"] == \"online\":\n",
        "        print(f\"{model_config.parameters['language_code']} : {model_config.model_name}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# This example uses a .wav file with LINEAR_PCM encoding.\n",
        "input_speech_file = str(Path(\"~/tutorials/audio_samples/en-US_sample.wav\").expanduser())\n",
        "ipd.Audio(input_speech_file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Send streaming requests with the first config and then followed by audio chunks to the server. Receive the responses and print them to the console."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "def read_responses(responses):\n",
        "    try:\n",
        "        final_transcript = \"\"\n",
        "        for response in responses:\n",
        "            if not response.results:\n",
        "                continue\n",
        "            for result in response.results:\n",
        "                if not result.alternatives:\n",
        "                    continue\n",
        "                if result.is_final:\n",
        "                    final_transcript += result.alternatives[0].transcript\n",
        "                    print(f\"FINAL: {result.audio_processed:.2f} : {result.alternatives[0].transcript}\")\n",
        "                else:\n",
        "                    print(f\"PARTIAL: {result.audio_processed:.2f} : {result.alternatives[0].transcript}\")\n",
        "\n",
        "        print(\"Transcript:\", final_transcript)\n",
        "\n",
        "    except grpc.RpcError as error:\n",
        "        print(error.code(), error.details())\n",
        "        return\n",
        "\n",
        "\n",
        "def generate_requests(input_speech_file: str):\n",
        "    print(f\"Transcribing File: {input_speech_file}\")\n",
        "\n",
        "    # Set up an recognition config\n",
        "    streaming_config = riva.client.StreamingRecognitionConfig(\n",
        "        config = riva.client.RecognitionConfig(\n",
        "            language_code = \"en-US\",                    # Language code of the audio clip\n",
        "            max_alternatives = 1,                       # How many top-N hypotheses to return. Only value of 1 is supported.\n",
        "            enable_automatic_punctuation = True,        # Enable punctuation and capitalization\n",
        "            audio_channel_count = 1,                    # Mono - Default\n",
        "            verbatim_transcripts = False,               # Set to True to return verbatim transcripts\n",
        "            profanity_filter = False,                   # Set to True to filter and replace profane words with first letter followed by asterisks (e.g. \"f***\")\n",
        "        ),\n",
        "        interim_results = True\n",
        "    )\n",
        "\n",
        "    # In cases where audio samples are submitted in `.raw` format, you need to set the following parameters appropriately:\n",
        "    # streaming_config.config.encoding = riva.client.AudioEncoding.LINEAR_PCM\n",
        "    # streaming_config.config.sample_rate_hertz = 16000\n",
        "\n",
        "    # First send the config\n",
        "    yield riva.client.proto.riva_asr_pb2.StreamingRecognizeRequest(streaming_config=streaming_config)\n",
        "\n",
        "    # Followed by audio chunks\n",
        "    try:\n",
        "        # stream audio in chunks of 100ms\n",
        "        chunk_size_ms = 100\n",
        "        for audio_chunk in riva.client.AudioChunkFileIterator(input_speech_file, chunk_size_ms):\n",
        "            yield riva.client.proto.riva_asr_pb2.StreamingRecognizeRequest(audio_content=audio_chunk)\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "        return\n",
        "\n",
        "# Get response stream to read transcripts\n",
        "read_responses(asr_service.stub.StreamingRecognize(generate_requests(input_speech_file)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Speaker Diarization\n",
        "\n",
        "Speaker Diarization is supported using Streaming as well."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import riva.client.proto.riva_asr_pb2 as riva_asr_pb2\n",
        "\n",
        "# This example uses a .wav file with LINEAR_PCM encoding.\n",
        "input_speech_file = str(Path(\"~/tutorials/audio_samples/2_speaker_en_US.wav\").expanduser())\n",
        "\n",
        "def read_responses(responses):\n",
        "    try:\n",
        "        final_transcript = \"\"\n",
        "        for response in responses:\n",
        "            if not response.results:\n",
        "                continue\n",
        "            for result in response.results:\n",
        "                if not result.alternatives:\n",
        "                    continue\n",
        "                if result.is_final:\n",
        "                    final_transcript += result.alternatives[0].transcript\n",
        "                    speaker = \"\"\n",
        "                    for word in result.alternatives[0].words:\n",
        "                        if speaker != word.speaker_tag:\n",
        "                            speaker = word.speaker_tag\n",
        "                            print(f\"\\nSpeaker {speaker}: \", end=\"\")\n",
        "                        print(word.word, end = \" \")\n",
        "                else:\n",
        "                    print(f\"PARTIAL: {result.audio_processed:.2f} : {result.alternatives[0].transcript}\")\n",
        "        print(\"\\nTranscript:\", final_transcript)\n",
        "    except grpc.RpcError as error:\n",
        "        print(error.code(), error.details())\n",
        "        return\n",
        "\n",
        "\n",
        "def generate_requests(input_speech_file: str):\n",
        "    print(f\"Transcribing File: {input_speech_file}\")\n",
        "\n",
        "    # Set up an recognition config\n",
        "    streaming_config = riva.client.StreamingRecognitionConfig(\n",
        "        config = riva.client.RecognitionConfig(\n",
        "            language_code = \"en-US\",                    # Language code of the audio clip\n",
        "            max_alternatives = 1,                       # How many top-N hypotheses to return. Only value of 1 is supported.\n",
        "            enable_automatic_punctuation = True,        # Enable punctuation and capitalization\n",
        "            audio_channel_count = 1,                    # Mono - Default\n",
        "            verbatim_transcripts = False,               # Set to True to return verbatim transcripts\n",
        "            profanity_filter = False,                   # Set to True to filter and replace profane words with first letter followed by asterisks (e.g. \"f***\")\n",
        "            enable_word_time_offsets = True,            # Enable word timestamps since they're needed for diarization\n",
        "        ),\n",
        "        interim_results = True\n",
        "    )\n",
        "\n",
        "    # In cases where audio samples are submitted in `.raw` format, you need to set the following parameters appropriately:\n",
        "    riva.client.asr.add_speaker_diarization_to_config(streaming_config.config, diarization_enable=True, diarization_max_speakers=8)\n",
        "    \n",
        "    # First send the config\n",
        "    yield riva.client.proto.riva_asr_pb2.StreamingRecognizeRequest(streaming_config=streaming_config)\n",
        "    \n",
        "    # Followed by audio chunks\n",
        "    try:\n",
        "        # stream audio in chunks of 100ms\n",
        "        chunk_size_ms = 100\n",
        "        for audio_chunk in riva.client.AudioChunkFileIterator(input_speech_file, chunk_size_ms):\n",
        "            yield riva.client.proto.riva_asr_pb2.StreamingRecognizeRequest(audio_content=audio_chunk)\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "        return\n",
        "\n",
        "# Get response stream to read transcripts\n",
        "read_responses(asr_service.stub.StreamingRecognize(generate_requests(input_speech_file)))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "usr",
      "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}