site/en/hub/tutorials/retrieval_with_tf_hub_universal_encoder_qa.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VFMCdVJIIraw"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "code",
        "id": "ZxMYj8OpIrCp"
      },
      "outputs": [],
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0fO2R2BBKx3l"
      },
      "source": [
        "# Multilingual Universal Sentence Encoder Q&A Retrieval\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/retrieval_with_tf_hub_universal_encoder_qa\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/hub/tutorials/retrieval_with_tf_hub_universal_encoder_qa.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/hub/tutorials/retrieval_with_tf_hub_universal_encoder_qa.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/hub/tutorials/retrieval_with_tf_hub_universal_encoder_qa.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://tfhub.dev/s?q=google%2Funiversal-sentence-encoder-multilingual-qa%2F3%20OR%20google%2Funiversal-sentence-encoder-qa%2F3\"><img src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" />See TF Hub models</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zsDm_WgMNlJQ"
      },
      "source": [
        "This is a demo for using [Universal Encoder Multilingual Q&A model](https://tfhub.dev/google/universal-sentence-encoder-multilingual-qa/3) for question-answer retrieval of text, illustrating the use of **question_encoder** and **response_encoder** of the model. We use sentences from [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) paragraphs as the demo dataset, each sentence and its context (the text surrounding the sentence) is encoded into high dimension embeddings with the **response_encoder**. These embeddings are stored in an index built using the [simpleneighbors](https://pypi.org/project/simpleneighbors/) library for question-answer retrieval.\n",
        "\n",
        "On retrieval a random question is selected from the [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) dataset and encoded into high dimension embedding with the **question_encoder** and query the  simpleneighbors index returning a list of approximate nearest neighbors in semantic space."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U0eOW2LTWiLg"
      },
      "source": [
        "### More models\n",
        "You can find all currently hosted text embedding models [here](https://tfhub.dev/s?module-type=text-embedding) and all models that have been trained on SQuAD as well [here](https://tfhub.dev/s?dataset=squad)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ORy-KvWXGXBo"
      },
      "source": [
        "## Setup\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "x00t_uJCEbeb"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "#@title Setup Environment\n",
        "# Install the latest Tensorflow version.\n",
        "!pip install -q \"tensorflow-text==2.11.*\"\n",
        "!pip install -q simpleneighbors[annoy]\n",
        "!pip install -q nltk\n",
        "!pip install -q tqdm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DmeFAuVsyWxg"
      },
      "outputs": [],
      "source": [
        "#@title Setup common imports and functions\n",
        "import json\n",
        "import nltk\n",
        "import os\n",
        "import pprint\n",
        "import random\n",
        "import simpleneighbors\n",
        "import urllib\n",
        "from IPython.display import HTML, display\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_hub as hub\n",
        "from tensorflow_text import SentencepieceTokenizer\n",
        "\n",
        "nltk.download('punkt')\n",
        "\n",
        "\n",
        "def download_squad(url):\n",
        "  return json.load(urllib.request.urlopen(url))\n",
        "\n",
        "def extract_sentences_from_squad_json(squad):\n",
        "  all_sentences = []\n",
        "  for data in squad['data']:\n",
        "    for paragraph in data['paragraphs']:\n",
        "      sentences = nltk.tokenize.sent_tokenize(paragraph['context'])\n",
        "      all_sentences.extend(zip(sentences, [paragraph['context']] * len(sentences)))\n",
        "  return list(set(all_sentences)) # remove duplicates\n",
        "\n",
        "def extract_questions_from_squad_json(squad):\n",
        "  questions = []\n",
        "  for data in squad['data']:\n",
        "    for paragraph in data['paragraphs']:\n",
        "      for qas in paragraph['qas']:\n",
        "        if qas['answers']:\n",
        "          questions.append((qas['question'], qas['answers'][0]['text']))\n",
        "  return list(set(questions))\n",
        "\n",
        "def output_with_highlight(text, highlight):\n",
        "  output = \"<li> \"\n",
        "  i = text.find(highlight)\n",
        "  while True:\n",
        "    if i == -1:\n",
        "      output += text\n",
        "      break\n",
        "    output += text[0:i]\n",
        "    output += '<b>'+text[i:i+len(highlight)]+'</b>'\n",
        "    text = text[i+len(highlight):]\n",
        "    i = text.find(highlight)\n",
        "  return output + \"</li>\\n\"\n",
        "\n",
        "def display_nearest_neighbors(query_text, answer_text=None):\n",
        "  query_embedding = model.signatures['question_encoder'](tf.constant([query_text]))['outputs'][0]\n",
        "  search_results = index.nearest(query_embedding, n=num_results)\n",
        "\n",
        "  if answer_text:\n",
        "    result_md = '''\n",
        "    <p>Random Question from SQuAD:</p>\n",
        "    <p>&nbsp;&nbsp;<b>%s</b></p>\n",
        "    <p>Answer:</p>\n",
        "    <p>&nbsp;&nbsp;<b>%s</b></p>\n",
        "    ''' % (query_text , answer_text)\n",
        "  else:\n",
        "    result_md = '''\n",
        "    <p>Question:</p>\n",
        "    <p>&nbsp;&nbsp;<b>%s</b></p>\n",
        "    ''' % query_text\n",
        "\n",
        "  result_md += '''\n",
        "    <p>Retrieved sentences :\n",
        "    <ol>\n",
        "  '''\n",
        "\n",
        "  if answer_text:\n",
        "    for s in search_results:\n",
        "      result_md += output_with_highlight(s, answer_text)\n",
        "  else:\n",
        "    for s in search_results:\n",
        "      result_md += '<li>' + s + '</li>\\n'\n",
        "\n",
        "  result_md += \"</ol>\"\n",
        "  display(HTML(result_md))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1kbkT8i3FL_C"
      },
      "source": [
        "Run the following code block to download and extract the SQuAD dataset into:\n",
        "\n",
        "* **sentences** is a list of (text, context) tuples - each paragraph from the SQuAD dataset are split into sentences using nltk library and the sentence and paragraph text forms the (text, context) tuple.\n",
        "* **questions** is a list of (question, answer) tuples.\n",
        "\n",
        "Note: You can use this demo to index the SQuAD train dataset or the smaller dev dataset (1.1 or 2.0) by selecting the **squad_url** below.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "iYqV2GAty_Eh"
      },
      "outputs": [],
      "source": [
        "#@title Download and extract SQuAD data\n",
        "squad_url = 'https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json' #@param [\"https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json\", \"https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v2.0.json\", \"https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\", \"https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\"]\n",
        "\n",
        "squad_json = download_squad(squad_url)\n",
        "sentences = extract_sentences_from_squad_json(squad_json)\n",
        "questions = extract_questions_from_squad_json(squad_json)\n",
        "print(\"%s sentences, %s questions extracted from SQuAD %s\" % (len(sentences), len(questions), squad_url))\n",
        "\n",
        "print(\"\\nExample sentence and context:\\n\")\n",
        "sentence = random.choice(sentences)\n",
        "print(\"sentence:\\n\")\n",
        "pprint.pprint(sentence[0])\n",
        "print(\"\\ncontext:\\n\")\n",
        "pprint.pprint(sentence[1])\n",
        "print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9x3u-2uSGbDf"
      },
      "source": [
        "The following code block setup the tensorflow graph **g** and **session** with the [Universal Encoder Multilingual Q&A model](https://tfhub.dev/google/universal-sentence-encoder-multilingual-qa/3)'s **question_encoder** and **response_encoder** signatures."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "44I0uCRQRiFO"
      },
      "outputs": [],
      "source": [
        "#@title Load model from tensorflow hub\n",
        "module_url = \"https://tfhub.dev/google/universal-sentence-encoder-multilingual-qa/3\" #@param [\"https://tfhub.dev/google/universal-sentence-encoder-multilingual-qa/3\", \"https://tfhub.dev/google/universal-sentence-encoder-qa/3\"]\n",
        "model = hub.load(module_url)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SCQpDmTZG0O6"
      },
      "source": [
        "The following code block compute the embeddings for all the text, context tuples and store them in a [simpleneighbors](https://pypi.org/project/simpleneighbors/) index using the **response_encoder**.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FwDUryIfSLp2"
      },
      "outputs": [],
      "source": [
        "#@title Compute embeddings and build simpleneighbors index\n",
        "batch_size = 100\n",
        "\n",
        "encodings = model.signatures['response_encoder'](\n",
        "  input=tf.constant([sentences[0][0]]),\n",
        "  context=tf.constant([sentences[0][1]]))\n",
        "index = simpleneighbors.SimpleNeighbors(\n",
        "    len(encodings['outputs'][0]), metric='angular')\n",
        "\n",
        "print('Computing embeddings for %s sentences' % len(sentences))\n",
        "slices = zip(*(iter(sentences),) * batch_size)\n",
        "num_batches = int(len(sentences) / batch_size)\n",
        "for s in tqdm(slices, total=num_batches):\n",
        "  response_batch = list([r for r, c in s])\n",
        "  context_batch = list([c for r, c in s])\n",
        "  encodings = model.signatures['response_encoder'](\n",
        "    input=tf.constant(response_batch),\n",
        "    context=tf.constant(context_batch)\n",
        "  )\n",
        "  for batch_index, batch in enumerate(response_batch):\n",
        "    index.add_one(batch, encodings['outputs'][batch_index])\n",
        "\n",
        "index.build()\n",
        "print('simpleneighbors index for %s sentences built.' % len(sentences))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZkNcjoPzHJpP"
      },
      "source": [
        "On retrieval, the question is encoded using the **question_encoder** and the question embedding is used to query the simpleneighbors index."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "J0xTw2w3UViK"
      },
      "outputs": [],
      "source": [
        "#@title Retrieve nearest neighbors for a random question from SQuAD\n",
        "num_results = 25 #@param {type:\"slider\", min:5, max:40, step:1}\n",
        "\n",
        "query = random.choice(questions)\n",
        "display_nearest_neighbors(query[0], query[1])"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "VFMCdVJIIraw"
      ],
      "name": "retrieval_with_tf_hub_universal_encoder_qa.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}