{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/riva_tts_tts-python-basics/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Streaming Speech-to-Speech Translation with Speech NIM\n",
    "\n",
    "This notebook streams English audio to an NVIDIA NMT NIM, which coordinates remote ASR and TTS NIM services to return synthesized Latin American Spanish audio.\n",
    "\n",
    "Before running the notebook, start the three-container network described in [Run Streaming Speech-to-Speech Translation](../speech-to-speech-translation.md). Keep `udhr-english.wav` and `udhr-spanish.wav` in the same directory as this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install the Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U nvidia-riva-client numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wave\n",
    "\n",
    "import IPython.display as ipd\n",
    "import numpy as np\n",
    "import riva.client\n",
    "import riva.client.proto.riva_asr_pb2 as riva_asr_pb2\n",
    "import riva.client.proto.riva_nmt_pb2 as riva_nmt_pb2\n",
    "\n",
    "print(\"Using Riva Python Client version:\", riva.client.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure the Example\n",
    "\n",
    "The NMT NIM is the client-facing service for the integrated pipeline. It reaches the ASR and TTS services through the remote-service addresses configured in the NMT container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NMT_URI = \"localhost:50051\"\n",
    "MODEL_NAME = \"megatronnmt_any_any_1b\"\n",
    "SOURCE_LANGUAGE = \"en-US\"\n",
    "TARGET_LANGUAGE = \"es-US\"\n",
    "TARGET_VOICE = \"Magpie-Multilingual.ES-US.Isabela\"\n",
    "OUTPUT_SAMPLE_RATE = 16_000\n",
    "CHUNK_DURATION_MS = 100\n",
    "SOURCE_AUDIO = \"udhr-english.wav\"\n",
    "REFERENCE_AUDIO = \"udhr-spanish.wav\"\n",
    "GENERATED_AUDIO = \"translated-output.wav\"\n",
    "\n",
    "with wave.open(SOURCE_AUDIO, \"rb\") as source_file:\n",
    "    source_sample_rate = source_file.getframerate()\n",
    "\n",
    "CHUNK_N_FRAMES = source_sample_rate * CHUNK_DURATION_MS // 1_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auth = riva.client.Auth(uri=NMT_URI)\n",
    "nmt_client = riva.client.NeuralMachineTranslationClient(auth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Play the Source Audio\n",
    "\n",
    "The source audio is a reading of Article 1 of the Universal Declaration of Human Rights in English:\n",
    "\n",
    "> All human beings are born free and equal in dignity and rights. They are endowed with reason and conscience, and should act toward one another in a spirit of brotherhood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(SOURCE_AUDIO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure the S2S Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asr_config = riva_asr_pb2.StreamingRecognitionConfig(\n",
    "    config=riva_asr_pb2.RecognitionConfig(\n",
    "        language_code=SOURCE_LANGUAGE,\n",
    "        max_alternatives=1,\n",
    "        enable_automatic_punctuation=True,\n",
    "    ),\n",
    "    interim_results=True,\n",
    ")\n",
    "\n",
    "translation_config = riva_nmt_pb2.TranslationConfig(\n",
    "    source_language_code=SOURCE_LANGUAGE,\n",
    "    target_language_code=TARGET_LANGUAGE,\n",
    "    model_name=MODEL_NAME,\n",
    ")\n",
    "\n",
    "tts_config = riva_nmt_pb2.SynthesizeSpeechConfig(\n",
    "    language_code=TARGET_LANGUAGE,\n",
    "    encoding=riva.client.AudioEncoding.LINEAR_PCM,\n",
    "    sample_rate_hz=OUTPUT_SAMPLE_RATE,\n",
    "    voice_name=TARGET_VOICE,\n",
    ")\n",
    "\n",
    "s2s_config = riva.client.StreamingTranslateSpeechToSpeechConfig(\n",
    "    asr_config=asr_config,\n",
    "    translation_config=translation_config,\n",
    "    tts_config=tts_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stream the Audio\n",
    "\n",
    "`AudioChunkFileIterator` sends the source file in real-time chunks. The NMT client returns synthesized audio chunks from the integrated pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_chunks = riva.client.AudioChunkFileIterator(\n",
    "    SOURCE_AUDIO,\n",
    "    CHUNK_N_FRAMES,\n",
    "    delay_callback=riva.client.sleep_audio_length,\n",
    ")\n",
    "\n",
    "responses = nmt_client.streaming_s2s_response_generator(\n",
    "    audio_chunks=audio_chunks,\n",
    "    streaming_config=s2s_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "received_chunks = []\n",
    "\n",
    "for index, response in enumerate(responses):\n",
    "    audio_samples = np.frombuffer(response.speech.audio, dtype=np.int16)\n",
    "    if audio_samples.size == 0:\n",
    "        continue\n",
    "\n",
    "    print(f\"Received audio chunk {index}\")\n",
    "    received_chunks.append(audio_samples)\n",
    "\n",
    "if not received_chunks:\n",
    "    raise RuntimeError(\"The S2S pipeline returned no audio. Check the container logs.\")\n",
    "\n",
    "translated_audio = np.concatenate(received_chunks)\n",
    "\n",
    "with wave.open(GENERATED_AUDIO, \"wb\") as output_file:\n",
    "    output_file.setnchannels(1)\n",
    "    output_file.setsampwidth(2)\n",
    "    output_file.setframerate(OUTPUT_SAMPLE_RATE)\n",
    "    output_file.writeframes(translated_audio.tobytes())\n",
    "\n",
    "print(f\"Saved translated audio to {GENERATED_AUDIO}\")\n",
    "ipd.Audio(translated_audio, rate=OUTPUT_SAMPLE_RATE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare with the Reference Output\n",
    "\n",
    "Play the reference output captured with the original example. Its Spanish translation is:\n",
    "\n",
    "> Todos los seres humanos nacen libres e iguales en dignidad y derechos; están dotados de razón y conciencia y deben actuar unos con otros en espíritu de hermandad.\n",
    "\n",
    "Generated translations and speech can differ across model and container versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(REFERENCE_AUDIO)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
