{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6277db36",
   "metadata": {},
   "outputs": [],
   "source": [
    "# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
    "# SPDX-License-Identifier: MIT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c97d9d54",
   "metadata": {},
   "source": [
    "# NLP API tutorial\n",
    "\n",
    "This tutorial demonstates how to use Python Riva API.\n",
    "\n",
    "## <font color=\"blue\">Server</font>\n",
    "\n",
    "Before running client part of Riva, please set up a server. The simplest\n",
    "way to do this is to follow\n",
    "[quick start guide](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/quick-start-guide.html#local-deployment-using-quick-start-scripts).\n",
    "\n",
    "\n",
    "## <font color=\"blue\">Authentication</font>\n",
    "\n",
    "Before using Riva services you will need to establish connection with a server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd57f561",
   "metadata": {},
   "outputs": [],
   "source": [
    "import riva.client\n",
    "\n",
    "uri = \"localhost:50051\"  # Default value\n",
    "\n",
    "auth = riva.client.Auth(uri=uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5319eae",
   "metadata": {},
   "source": [
    "## <font color=\"blue\">Setting up service</font>\n",
    "\n",
    "To instantiate a service pass `riva.client.Auth` instance to a constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb75ba38",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp_service = riva.client.NLPService(auth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "267f4a06",
   "metadata": {},
   "source": [
    "You may find full response field description in [documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/riva_asr.proto.html?highlight=max%20alternatives#riva-proto-riva-asr-proto)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e01f300b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from riva.client.proto.riva_nlp_pb2 import (\n",
    "    AnalyzeIntentResponse,\n",
    "    NaturalQueryResponse,\n",
    "    TextClassResponse,\n",
    "    TextTransformResponse,\n",
    "    TokenClassResponse,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c87b82b",
   "metadata": {},
   "source": [
    "## <font color=\"blue\">Text classification</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "539c95f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_class_queries = [\"A hurricane is approaching Japan.\", \"What is the weather on Wednesday in Moscow?\"]\n",
    "text_class_model = \"riva_intent_weather\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35834e64",
   "metadata": {},
   "outputs": [],
   "source": [
    "response: TextClassResponse = nlp_service.classify_text(text_class_queries, text_class_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a9432ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "796cab62",
   "metadata": {},
   "outputs": [],
   "source": [
    "detected_intent = response.results[0].labels[0].class_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f13be276",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(detected_intent)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3695715e",
   "metadata": {},
   "source": [
    "You may use a function `riva.client.extract_most_probable_text_class_and_confidence()`. The function returns a list of most probable intents and their scores for all queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61ecf049",
   "metadata": {},
   "outputs": [],
   "source": [
    "classes, probs = riva.client.extract_most_probable_text_class_and_confidence(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14975bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(classes)\n",
    "print(probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97e891c8",
   "metadata": {},
   "source": [
    "## <font color=\"blue\">Token classification</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83ba98f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "response: TokenClassResponse = nlp_service.classify_tokens(text_class_queries[1], text_class_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ec2e730",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16eb717e",
   "metadata": {},
   "outputs": [],
   "source": [
    "token = response.results[0].results[0].token\n",
    "class_name = response.results[0].results[0].label[0].class_name\n",
    "class_score = response.results[0].results[0].label[0].score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0910b208",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(token, class_name, class_score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffdd9348",
   "metadata": {},
   "source": [
    "You may use a function `riva.client.extract_most_probable_token_classification_predictions()`. The function returns:\n",
    " - list of tokens lists for all elements of a batch,\n",
    " - list of most probable classes lists for all elements of a batch,\n",
    " - list of most probable classes confidences lits for all elements of a batch,\n",
    " - list of token span starts lists for all elements of a batch,\n",
    " - list of token span ends lists for all elements of a batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76532974",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tokens, class_names, confidences, starts, ends = riva.client.extract_most_probable_token_classification_predictions(\n",
    "    response\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaead038",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"First batch element tokens:\", tokens[0])\n",
    "print(\"First batch element first token class name:\", class_names[0][0])\n",
    "print(confidences)\n",
    "print(starts)\n",
    "print(ends)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9e98246",
   "metadata": {},
   "source": [
    "> Spans do not work properly for batches which contain more than 1 element."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef5254e8",
   "metadata": {},
   "source": [
    "## <font color=\"blue\">Punctuation and Capitalization</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb7d7a9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Batches of sizes greater than 1 are not working with riva v2.2.0 and v2.2.1\n",
    "\n",
    "punctuation_queries = [\n",
    "    \"by the early 20th century the gar complained more and more about the younger generation\",\n",
    "#     \"boa Vista is the capital of the brazilian state of roraima situated on the western bank of \"\n",
    "#     \"the branco river the city lies 220 km from brazil's border with venezuela.\",\n",
    "]\n",
    "response: TextTransformResponse = nlp_service.punctuate_text(punctuation_queries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd2361a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b480ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_query_result = response.text[0]\n",
    "print(first_query_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e796a68a",
   "metadata": {},
   "source": [
    "## <font color=\"blue\">Intent analysis</font>\n",
    "\n",
    "Accepts an input string and returns the most likely intent as well as slots relevant to that intent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4638b508",
   "metadata": {},
   "outputs": [],
   "source": [
    "options = riva.client.AnalyzeIntentOptions(lang='en-US')\n",
    "intent_query = \"How is the weather today in New England?\"\n",
    "response: AnalyzeIntentResponse = nlp_service.analyze_intent(intent_query, options)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f330a101",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8580421",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"intent name:\", response.intent.class_name)\n",
    "print(\"intent score:\", response.intent.score)\n",
    "print(\"domain name:\", response.domain.class_name)\n",
    "print(\"domain score:\", response.domain.score)\n",
    "print(\"first slot token:\", response.slots[0].token)\n",
    "print(\"first slot most probable label name:\", response.slots[0].label[0].class_name)\n",
    "print(\"first slot most probable label score:\", response.slots[0].label[0].score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b146f07",
   "metadata": {},
   "source": [
    "## <font color=\"blue\">Question answering</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a6cce30",
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_query = \"How many gigatons of carbon dioxide was released in 2005?\"\n",
    "qa_context = (\n",
    "    \"In 2010 the Amazon rainforest experienced another severe drought, in some ways more extreme than the \"\n",
    "    \"2005 drought. The affected region was approximate 1,160,000 square miles (3,000,000 km2) of \"\n",
    "    \"rainforest, compared to 734,000 square miles (1,900,000 km2) in 2005. The 2010 drought had three \"\n",
    "    \"epicenters where vegetation died off, whereas in 2005 the drought was focused on the southwestern \"\n",
    "    \"part. The findings were published in the journal Science. In a typical year the Amazon absorbs 1.5 \"\n",
    "    \"gigatons of carbon dioxide; during 2005 instead 5 gigatons were released and in 2010 8 gigatons were \"\n",
    "    \"released.\"\n",
    ")\n",
    "response: NaturalQueryResponse = nlp_service.natural_query(qa_query, qa_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a65243d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "602518ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = response.results[0].answer\n",
    "score = response.results[0].score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d798d3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "198f2b2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64125714",
   "metadata": {},
   "source": [
    "## <font color=\"blue\">Asynchronous calls</font>\n",
    "\n",
    "Any of the above methods can be used in asynchronous manner. For this you need set parameter `future=True`. Then instead of response the methods will return future objects. Responses can be retrieved by calling `result()` on future objects.\n",
    "\n",
    "Following example demonstrates how latency is reduced via using asynchronous calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a35d96f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f4135ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_class_queries = [\"A hurricane is approaching Japan.\", \"What is the weather on Wednesday in Moscow?\"]\n",
    "text_class_model = \"riva_intent_weather\"\n",
    "# Batches of sizes greater than 1 are not working with riva v2.2.0 and v2.2.1\n",
    "\n",
    "punctuation_queries = [\n",
    "    \"by the early 20th century the gar complained more and more about the younger generation\",\n",
    "#     \"boa Vista is the capital of the brazilian state of roraima situated on the western bank of \"\n",
    "#     \"the branco river the city lies 220 km from brazil's border with venezuela.\",\n",
    "]\n",
    "\n",
    "options = riva.client.AnalyzeIntentOptions(lang='en-US')\n",
    "intent_query = \"How is the weather today in New England?\"\n",
    "\n",
    "qa_query = \"How many gigatons of carbon dioxide was released in 2005?\"\n",
    "qa_context = (\n",
    "    \"In 2010 the Amazon rainforest experienced another severe drought, in some ways more extreme than the \"\n",
    "    \"2005 drought. The affected region was approximate 1,160,000 square miles (3,000,000 km2) of \"\n",
    "    \"rainforest, compared to 734,000 square miles (1,900,000 km2) in 2005. The 2010 drought had three \"\n",
    "    \"epicenters where vegetation died off, whereas in 2005 the drought was focused on the southwestern \"\n",
    "    \"part. The findings were published in the journal Science. In a typical year the Amazon absorbs 1.5 \"\n",
    "    \"gigatons of carbon dioxide; during 2005 instead 5 gigatons were released and in 2010 8 gigatons were \"\n",
    "    \"released.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c451bc58",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_class_reference = ['weather.alert', 'weather.weather']\n",
    "token_class_reference = ['wednesday', 'moscow', '?']\n",
    "punctuation_reference = (\n",
    "    \"By the early 20th century, the Gar complained more and more about the younger generation.\"\n",
    ")\n",
    "intent_analysis_reference = \"weather.weather\"\n",
    "natural_query_reference = \"5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4e72ccd",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_repeats = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99b68e87",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "def check_text_classification(ref: List[str], resp: TextClassResponse, repeat_idx: int) -> None:\n",
    "    labels, _ = riva.client.extract_most_probable_text_class_and_confidence(resp)\n",
    "    assert labels == ref, f\"On repeat {repeat_idx} expected text classification results {ref}, but got {labels}.\"\n",
    "\n",
    "def check_token_classification(ref: List[str], resp: TokenClassResponse, repeat_idx: int) -> None:\n",
    "    tokens = riva.client.extract_most_probable_token_classification_predictions(resp)[0][0]\n",
    "    assert tokens == token_class_reference, (\n",
    "        f\"On repeat {repeat_idx} expected to find token classification tokens {ref}, but got {tokens}.\"\n",
    "    )\n",
    "\n",
    "def check_punctuation(ref: str, resp: TextTransformResponse, repeat_idx: int) -> None:\n",
    "    output = resp.text[0]\n",
    "    assert output == ref, (\n",
    "        f\"On repeat {repeat_idx} expected punctuated output '{ref}', but got '{output}'.\"\n",
    "    )\n",
    "\n",
    "def check_intent_analysis(ref: str, resp: AnalyzeIntentResponse, repeat_idx: int) -> None:\n",
    "    output = resp.intent.class_name\n",
    "    assert output == ref, f\"On repeat {repeat_idx} expected intent is '{ref}', but got '{output}'.\"\n",
    "\n",
    "def check_natural_query(ref: str, resp: NaturalQueryResponse, repeat_idx: int) -> None:\n",
    "    answer = resp.results[0].answer\n",
    "    assert answer == ref, f\"On repeat {repeat_idx} expected answer is '{ref}', but got '{answer}'.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e21b04f5",
   "metadata": {},
   "source": [
    "Synchronous:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d54d12",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time()\n",
    "for repeat_idx in range(num_repeats):\n",
    "    text_class_response = nlp_service.classify_text(text_class_queries, text_class_model)\n",
    "    check_text_classification(text_class_reference, text_class_response, repeat_idx)\n",
    "    \n",
    "    token_class_response = nlp_service.classify_tokens(text_class_queries[1], text_class_model)\n",
    "    check_token_classification(token_class_reference, token_class_response, repeat_idx)\n",
    "    \n",
    "    punctuation_response = nlp_service.punctuate_text(punctuation_queries)\n",
    "    check_punctuation(punctuation_reference, punctuation_response, repeat_idx)\n",
    "    \n",
    "    intent_analysis_response = nlp_service.analyze_intent(intent_query, options)\n",
    "    check_intent_analysis(intent_analysis_reference, intent_analysis_response, repeat_idx)\n",
    "    \n",
    "    natural_query_response = nlp_service.natural_query(qa_query, qa_context)\n",
    "    check_natural_query(natural_query_reference, natural_query_response, repeat_idx)\n",
    "    \n",
    "print(f\"time spent on synchronous calls: {time() - start_time:.2f} sec\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02a5d5c6",
   "metadata": {},
   "source": [
    "Asynchronous:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ca17218",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time()\n",
    "futures = []\n",
    "for _ in range(num_repeats):\n",
    "    repeat_futures = []\n",
    "    repeat_futures.append(nlp_service.classify_text(text_class_queries, text_class_model, future=True))\n",
    "    repeat_futures.append(nlp_service.classify_tokens(text_class_queries[1], text_class_model, future=True))\n",
    "    repeat_futures.append(nlp_service.punctuate_text(punctuation_queries, future=True))\n",
    "    repeat_futures.append(nlp_service.analyze_intent(intent_query, options, future=True))\n",
    "    repeat_futures.append(nlp_service.natural_query(qa_query, qa_context, future=True))\n",
    "    futures.append(repeat_futures)\n",
    "for repeat_idx, repeat_futures in enumerate(futures):\n",
    "    check_text_classification(text_class_reference, repeat_futures[0].result(), repeat_idx)\n",
    "    check_token_classification(token_class_reference, repeat_futures[1].result(), repeat_idx)\n",
    "    check_punctuation(punctuation_reference, repeat_futures[2].result(), repeat_idx)\n",
    "    check_intent_analysis(intent_analysis_reference, repeat_futures[3].result(), repeat_idx)\n",
    "    check_natural_query(natural_query_reference, repeat_futures[4].result(), repeat_idx)\n",
    "        \n",
    "print(f\"time spent on async calls: {time() - start_time:.2f}sec\")"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}