python/notebooks/EntityPredictionDemo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Entity Prediction\n",
    "---\n",
    "\n",
    "#### _Given a user query and context, predict the entity that the user is looking for_\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First... Install ml4ir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.python.org/simple\n",
      "Requirement already satisfied: ml4ir in /Users/ashish.srinivasa/search_relevance/ml4ir/python (0.0.1)\n",
      "Requirement already satisfied: click in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (7.0)\n",
      "Requirement already satisfied: Sphinx in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.3.1)\n",
      "Requirement already satisfied: coverage in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (5.0.3)\n",
      "Requirement already satisfied: awscli in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.17.9)\n",
      "Requirement already satisfied: flake8 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.7.9)\n",
      "Requirement already satisfied: python-dotenv>=0.5.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.10.5)\n",
      "Requirement already satisfied: flake8-black in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.1.1)\n",
      "Requirement already satisfied: flake8-mypy in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (17.8.0)\n",
      "Requirement already satisfied: pre-commit in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.0.1)\n",
      "Requirement already satisfied: mypy in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.761)\n",
      "Requirement already satisfied: appnope==0.1.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.1.0)\n",
      "Requirement already satisfied: attrs==19.3.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (19.3.0)\n",
      "Requirement already satisfied: backcall==0.1.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.1.0)\n",
      "Requirement already satisfied: colorlog==4.0.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.0.2)\n",
      "Requirement already satisfied: dask==2.8.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.8.1)\n",
      "Requirement already satisfied: decorator==4.4.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.4.1)\n",
      "Requirement already satisfied: dill==0.3.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.3.0)\n",
      "Requirement already satisfied: distributed==2.8.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.8.1)\n",
      "Requirement already satisfied: entrypoints==0.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.3)\n",
      "Requirement already satisfied: future==0.18.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.18.2)\n",
      "Requirement already satisfied: hdfs==2.5.8 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.5.8)\n",
      "Requirement already satisfied: ipykernel==5.1.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (5.1.3)\n",
      "Requirement already satisfied: ipywidgets==7.5.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (7.5.1)\n",
      "Requirement already satisfied: ipython==7.11.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (7.11.1)\n",
      "Requirement already satisfied: ipython-genutils==0.2.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.2.0)\n",
      "Requirement already satisfied: jedi==0.15.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.15.2)\n",
      "Requirement already satisfied: joblib==0.14.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.14.0)\n",
      "Requirement already satisfied: json5==0.8.5 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.8.5)\n",
      "Requirement already satisfied: jsonschema==3.2.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.2.0)\n",
      "Requirement already satisfied: jupyter-client==5.3.4 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (5.3.4)\n",
      "Requirement already satisfied: jupyter-core==4.6.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.6.1)\n",
      "Requirement already satisfied: jupyterlab==1.2.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.2.3)\n",
      "Requirement already satisfied: jupyterlab-server==1.0.6 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.0.6)\n",
      "Requirement already satisfied: Keras-Applications==1.0.8 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.0.8)\n",
      "Requirement already satisfied: Keras-Preprocessing==1.1.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.1.0)\n",
      "Requirement already satisfied: lime==0.1.1.36 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.1.1.36)\n",
      "Requirement already satisfied: Markdown==3.1.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.1.1)\n",
      "Requirement already satisfied: MarkupSafe==1.1.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.1.1)\n",
      "Requirement already satisfied: matplotlib==3.1.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.1.2)\n",
      "Requirement already satisfied: nbconvert==5.6.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (5.6.1)\n",
      "Requirement already satisfied: nbformat==4.4.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.4.0)\n",
      "Requirement already satisfied: notebook==6.0.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (6.0.2)\n",
      "Requirement already satisfied: numpy==1.17.4 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.17.4)\n",
      "Requirement already satisfied: oauth2client==3.0.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.0.0)\n",
      "Requirement already satisfied: oauthlib==3.1.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.1.0)\n",
      "Requirement already satisfied: pandas==0.25.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.25.3)\n",
      "Requirement already satisfied: plotly==4.4.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.4.1)\n",
      "Requirement already satisfied: protobuf==3.10.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.10.0)\n",
      "Requirement already satisfied: pycodestyle==2.5.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.5.0)\n",
      "Requirement already satisfied: pyflakes==2.1.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.1.1)\n",
      "Requirement already satisfied: pytest==4.6.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.6.3)\n",
      "Requirement already satisfied: pytest-cov==2.5.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.5.1)\n",
      "Requirement already satisfied: pytest-html==2.1.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/pytest_html-2.1.1-py3.7.egg (from ml4ir) (2.1.1)\n",
      "Requirement already satisfied: python-dateutil==2.8.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.8.1)\n",
      "Requirement already satisfied: parso==0.5.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.5.2)\n",
      "Requirement already satisfied: pexpect==4.7.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.7.0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: pickleshare==0.7.5 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.7.5)\n",
      "Requirement already satisfied: prompt-toolkit==3.0.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (3.0.2)\n",
      "Requirement already satisfied: ptyprocess==0.6.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.6.0)\n",
      "Requirement already satisfied: Pygments==2.5.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.5.2)\n",
      "Requirement already satisfied: PyYAML==5.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (5.1)\n",
      "Requirement already satisfied: requests==2.22.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.22.0)\n",
      "Requirement already satisfied: requests-oauthlib==1.3.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.3.0)\n",
      "Requirement already satisfied: rsa==4.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.0)\n",
      "Requirement already satisfied: scikit-image==0.16.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.16.2)\n",
      "Requirement already satisfied: scikit-learn==0.21.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.21.3)\n",
      "Requirement already satisfied: scipy==1.3.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.3.2)\n",
      "Requirement already satisfied: seaborn==0.9.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.9.0)\n",
      "Requirement already satisfied: six==1.14.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.14.0)\n",
      "Requirement already satisfied: sklearn==0.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.0)\n",
      "Requirement already satisfied: swifter==0.296 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.296)\n",
      "Requirement already satisfied: tblib==1.5.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.5.0)\n",
      "Requirement already satisfied: tensorboard==2.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.0.1)\n",
      "Requirement already satisfied: tensorflow==2.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.0.1)\n",
      "Requirement already satisfied: tensorflow-estimator==2.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.0.1)\n",
      "Requirement already satisfied: tensorflow-hub==0.7.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.7.0)\n",
      "Requirement already satisfied: tensorflow-metadata==0.15.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.15.1)\n",
      "Requirement already satisfied: tensorflow-probability==0.8.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.8.0)\n",
      "Requirement already satisfied: tensorflow-ranking==0.2.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.2.0)\n",
      "Requirement already satisfied: tensorflow-serving-api==2.0.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.0.0)\n",
      "Requirement already satisfied: tensorflow-text==2.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (2.0.1)\n",
      "Requirement already satisfied: tensorflow-transform==0.15.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.15.0)\n",
      "Requirement already satisfied: traitlets==4.3.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (4.3.3)\n",
      "Requirement already satisfied: urllib3==1.25.7 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (1.25.7)\n",
      "Requirement already satisfied: wandb==0.8.36 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.8.36)\n",
      "Requirement already satisfied: wcwidth==0.1.8 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ml4ir) (0.1.8)\n",
      "Requirement already satisfied: packaging in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (20.1)\n",
      "Requirement already satisfied: babel!=2.0,>=1.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (2.8.0)\n",
      "Requirement already satisfied: sphinxcontrib-htmlhelp in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (1.0.2)\n",
      "Requirement already satisfied: sphinxcontrib-devhelp in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (1.0.1)\n",
      "Requirement already satisfied: Jinja2>=2.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (2.11.1)\n",
      "Requirement already satisfied: snowballstemmer>=1.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (2.0.0)\n",
      "Requirement already satisfied: alabaster<0.8,>=0.7 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (0.7.12)\n",
      "Requirement already satisfied: sphinxcontrib-applehelp in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (1.0.1)\n",
      "Requirement already satisfied: docutils>=0.12 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (0.16)\n",
      "Requirement already satisfied: sphinxcontrib-serializinghtml in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (1.1.3)\n",
      "Requirement already satisfied: sphinxcontrib-qthelp in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (1.0.2)\n",
      "Requirement already satisfied: setuptools in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (45.1.0)\n",
      "Requirement already satisfied: sphinxcontrib-jsmath in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (1.0.1)\n",
      "Requirement already satisfied: imagesize in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Sphinx->ml4ir) (1.2.0)\n",
      "Requirement already satisfied: s3transfer<0.4.0,>=0.3.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from awscli->ml4ir) (0.3.2)\n",
      "Requirement already satisfied: colorama<0.4.2,>=0.2.5 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from awscli->ml4ir) (0.4.1)\n",
      "Requirement already satisfied: botocore==1.14.9 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from awscli->ml4ir) (1.14.9)\n",
      "Requirement already satisfied: mccabe<0.7.0,>=0.6.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from flake8->ml4ir) (0.6.1)\n",
      "Requirement already satisfied: black>=19.3b0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from flake8-black->ml4ir) (19.10b0)\n",
      "Requirement already satisfied: aspy.yaml in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pre-commit->ml4ir) (1.3.0)\n",
      "Requirement already satisfied: identify>=1.0.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pre-commit->ml4ir) (1.4.11)\n",
      "Requirement already satisfied: importlib-metadata; python_version < \"3.8\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pre-commit->ml4ir) (1.5.0)\n",
      "Requirement already satisfied: nodeenv>=0.11.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pre-commit->ml4ir) (1.3.5)\n",
      "Requirement already satisfied: cfgv>=2.0.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pre-commit->ml4ir) (2.0.1)\n",
      "Requirement already satisfied: virtualenv>=15.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pre-commit->ml4ir) (16.7.9)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: toml in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pre-commit->ml4ir) (0.10.0)\n",
      "Requirement already satisfied: typed-ast<1.5.0,>=1.4.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from mypy->ml4ir) (1.4.1)\n",
      "Requirement already satisfied: mypy-extensions<0.5.0,>=0.4.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from mypy->ml4ir) (0.4.3)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from mypy->ml4ir) (3.7.4.1)\n",
      "Requirement already satisfied: tornado>=5 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from distributed==2.8.1->ml4ir) (6.0.3)\n",
      "Requirement already satisfied: cloudpickle>=0.2.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from distributed==2.8.1->ml4ir) (1.2.2)\n",
      "Requirement already satisfied: zict>=0.1.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from distributed==2.8.1->ml4ir) (1.0.0)\n",
      "Requirement already satisfied: msgpack in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from distributed==2.8.1->ml4ir) (0.6.2)\n",
      "Requirement already satisfied: sortedcontainers!=2.0.0,!=2.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from distributed==2.8.1->ml4ir) (2.1.0)\n",
      "Requirement already satisfied: psutil>=5.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from distributed==2.8.1->ml4ir) (5.6.7)\n",
      "Requirement already satisfied: toolz>=0.7.4 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from distributed==2.8.1->ml4ir) (0.10.0)\n",
      "Requirement already satisfied: docopt in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from hdfs==2.5.8->ml4ir) (0.6.2)\n",
      "Requirement already satisfied: widgetsnbextension~=3.5.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from ipywidgets==7.5.1->ml4ir) (3.5.1)\n",
      "Requirement already satisfied: pyrsistent>=0.14.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from jsonschema==3.2.0->ml4ir) (0.15.7)\n",
      "Requirement already satisfied: pyzmq>=13 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from jupyter-client==5.3.4->ml4ir) (18.1.1)\n",
      "Requirement already satisfied: h5py in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from Keras-Applications==1.0.8->ml4ir) (2.10.0)\n",
      "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from matplotlib==3.1.2->ml4ir) (2.4.6)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from matplotlib==3.1.2->ml4ir) (1.1.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from matplotlib==3.1.2->ml4ir) (0.10.0)\n",
      "Requirement already satisfied: defusedxml in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from nbconvert==5.6.1->ml4ir) (0.6.0)\n",
      "Requirement already satisfied: bleach in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from nbconvert==5.6.1->ml4ir) (3.1.0)\n",
      "Requirement already satisfied: mistune<2,>=0.8.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from nbconvert==5.6.1->ml4ir) (0.8.4)\n",
      "Requirement already satisfied: testpath in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from nbconvert==5.6.1->ml4ir) (0.4.4)\n",
      "Requirement already satisfied: pandocfilters>=1.4.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from nbconvert==5.6.1->ml4ir) (1.4.2)\n",
      "Requirement already satisfied: prometheus-client in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from notebook==6.0.2->ml4ir) (0.7.1)\n",
      "Requirement already satisfied: terminado>=0.8.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from notebook==6.0.2->ml4ir) (0.8.3)\n",
      "Requirement already satisfied: Send2Trash in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from notebook==6.0.2->ml4ir) (1.5.0)\n",
      "Requirement already satisfied: httplib2>=0.9.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from oauth2client==3.0.0->ml4ir) (0.17.0)\n",
      "Requirement already satisfied: pyasn1>=0.1.7 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from oauth2client==3.0.0->ml4ir) (0.4.8)\n",
      "Requirement already satisfied: pyasn1-modules>=0.0.5 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from oauth2client==3.0.0->ml4ir) (0.2.8)\n",
      "Requirement already satisfied: pytz>=2017.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pandas==0.25.3->ml4ir) (2019.3)\n",
      "Requirement already satisfied: retrying>=1.3.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from plotly==4.4.1->ml4ir) (1.3.3)\n",
      "Requirement already satisfied: atomicwrites>=1.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pytest==4.6.3->ml4ir) (1.3.0)\n",
      "Requirement already satisfied: py>=1.5.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pytest==4.6.3->ml4ir) (1.8.1)\n",
      "Requirement already satisfied: pluggy<1.0,>=0.12 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pytest==4.6.3->ml4ir) (0.13.1)\n",
      "Requirement already satisfied: more-itertools>=4.0.0; python_version > \"2.7\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from pytest==4.6.3->ml4ir) (8.2.0)\n",
      "Requirement already satisfied: pytest-metadata in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/pytest_metadata-1.9.0-py3.7.egg (from pytest-html==2.1.1->ml4ir) (1.9.0)\n",
      "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from requests==2.22.0->ml4ir) (3.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from requests==2.22.0->ml4ir) (2019.11.28)\n",
      "Requirement already satisfied: idna<2.9,>=2.5 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from requests==2.22.0->ml4ir) (2.8)\n",
      "Requirement already satisfied: PyWavelets>=0.4.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from scikit-image==0.16.2->ml4ir) (1.1.1)\n",
      "Requirement already satisfied: imageio>=2.3.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from scikit-image==0.16.2->ml4ir) (2.6.1)\n",
      "Requirement already satisfied: pillow>=4.3.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from scikit-image==0.16.2->ml4ir) (7.0.0)\n",
      "Requirement already satisfied: networkx>=2.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from scikit-image==0.16.2->ml4ir) (2.4)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: tqdm>=4.33.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from swifter==0.296->ml4ir) (4.42.1)\n",
      "Requirement already satisfied: numba in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from swifter==0.296->ml4ir) (0.48.0)\n",
      "Requirement already satisfied: grpcio>=1.24.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorboard==2.0.1->ml4ir) (1.26.0)\n",
      "Requirement already satisfied: absl-py>=0.4 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorboard==2.0.1->ml4ir) (0.9.0)\n",
      "Requirement already satisfied: werkzeug>=0.11.15 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorboard==2.0.1->ml4ir) (0.16.1)\n",
      "Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorboard==2.0.1->ml4ir) (0.34.2)\n",
      "Requirement already satisfied: google-auth<2,>=1.6.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorboard==2.0.1->ml4ir) (1.11.0)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorboard==2.0.1->ml4ir) (0.4.1)\n",
      "Requirement already satisfied: wrapt>=1.11.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow==2.0.1->ml4ir) (1.11.2)\n",
      "Requirement already satisfied: astor>=0.6.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow==2.0.1->ml4ir) (0.8.1)\n",
      "Requirement already satisfied: opt-einsum>=2.3.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow==2.0.1->ml4ir) (3.1.0)\n",
      "Requirement already satisfied: google-pasta>=0.1.6 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow==2.0.1->ml4ir) (0.1.8)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow==2.0.1->ml4ir) (1.1.0)\n",
      "Requirement already satisfied: gast==0.2.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow==2.0.1->ml4ir) (0.2.2)\n",
      "Requirement already satisfied: googleapis-common-protos in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow-metadata==0.15.1->ml4ir) (1.51.0)\n",
      "Requirement already satisfied: pydot<2,>=1.2 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow-transform==0.15.0->ml4ir) (1.4.1)\n",
      "Requirement already satisfied: apache-beam[gcp]<3,>=2.16 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow-transform==0.15.0->ml4ir) (2.19.0)\n",
      "Requirement already satisfied: tfx-bsl<0.16,>=0.15 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from tensorflow-transform==0.15.0->ml4ir) (0.15.3)\n",
      "Requirement already satisfied: shortuuid>=0.5.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (1.0.1)\n",
      "Requirement already satisfied: docker-pycreds>=0.4.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (0.4.0)\n",
      "Requirement already satisfied: configparser>=3.8.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (5.0.0)\n",
      "Requirement already satisfied: GitPython>=1.0.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (3.1.3)\n",
      "Requirement already satisfied: subprocess32>=3.5.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (3.5.4)\n",
      "Requirement already satisfied: gql==0.2.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (0.2.0)\n",
      "Requirement already satisfied: sentry-sdk>=0.4.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (0.14.4)\n",
      "Requirement already satisfied: nvidia-ml-py3>=7.352.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (7.352.0)\n",
      "Requirement already satisfied: watchdog>=0.8.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from wandb==0.8.36->ml4ir) (0.10.2)\n",
      "Requirement already satisfied: jmespath<1.0.0,>=0.7.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from botocore==1.14.9->awscli->ml4ir) (0.9.4)\n",
      "Requirement already satisfied: appdirs in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from black>=19.3b0->flake8-black->ml4ir) (1.4.3)\n",
      "Requirement already satisfied: regex in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from black>=19.3b0->flake8-black->ml4ir) (2020.1.8)\n",
      "Requirement already satisfied: pathspec<1,>=0.6 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from black>=19.3b0->flake8-black->ml4ir) (0.7.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from importlib-metadata; python_version < \"3.8\"->pre-commit->ml4ir) (2.1.0)\n",
      "Requirement already satisfied: heapdict in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from zict>=0.1.3->distributed==2.8.1->ml4ir) (1.0.1)\n",
      "Requirement already satisfied: webencodings in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from bleach->nbconvert==5.6.1->ml4ir) (0.5.1)\n",
      "Requirement already satisfied: llvmlite<0.32.0,>=0.31.0dev0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from numba->swifter==0.296->ml4ir) (0.31.0)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from google-auth<2,>=1.6.3->tensorboard==2.0.1->ml4ir) (4.0.0)\n",
      "Requirement already satisfied: typing<3.8.0,>=3.7.0; python_version < \"3.8.0\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (3.7.4.1)\n",
      "Requirement already satisfied: crcmod<2.0,>=1.7 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.7)\n",
      "Collecting pyarrow<0.16.0,>=0.15.1; python_version >= \"3.0\" or platform_system != \"Windows\"\n",
      "  Using cached pyarrow-0.15.1-cp37-cp37m-macosx_10_6_intel.whl (36.6 MB)\n",
      "Requirement already satisfied: fastavro<0.22,>=0.21.4 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (0.21.24)\n",
      "Requirement already satisfied: pymongo<4.0.0,>=3.8.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (3.10.1)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: avro-python3<2.0.0,>=1.8.1; python_version >= \"3.0\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.9.1)\n",
      "Requirement already satisfied: mock<3.0.0,>=1.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (2.0.0)\n",
      "Requirement already satisfied: google-apitools<0.5.29,>=0.5.28; extra == \"gcp\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (0.5.28)\n",
      "Requirement already satisfied: google-cloud-pubsub<1.1.0,>=0.39.0; extra == \"gcp\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.0.2)\n",
      "Requirement already satisfied: google-cloud-core<2,>=0.28.1; extra == \"gcp\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.3.0)\n",
      "Requirement already satisfied: google-cloud-bigtable<1.1.0,>=0.31.1; extra == \"gcp\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.0.0)\n",
      "Requirement already satisfied: google-cloud-bigquery<1.18.0,>=1.6.0; extra == \"gcp\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.17.1)\n",
      "Requirement already satisfied: google-cloud-datastore<1.8.0,>=1.7.1; extra == \"gcp\" in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.7.4)\n",
      "Requirement already satisfied: gitdb<5,>=4.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from GitPython>=1.0.0->wandb==0.8.36->ml4ir) (4.0.5)\n",
      "Requirement already satisfied: promise<3,>=2.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from gql==0.2.0->wandb==0.8.36->ml4ir) (2.3)\n",
      "Requirement already satisfied: graphql-core<2,>=0.5.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from gql==0.2.0->wandb==0.8.36->ml4ir) (1.1)\n",
      "Requirement already satisfied: pathtools>=0.1.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from watchdog>=0.8.3->wandb==0.8.36->ml4ir) (0.1.2)\n",
      "Requirement already satisfied: pbr>=0.11 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from mock<3.0.0,>=1.0.1->apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (5.4.4)\n",
      "Requirement already satisfied: fasteners>=0.14 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from google-apitools<0.5.29,>=0.5.28; extra == \"gcp\"->apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (0.15)\n",
      "Requirement already satisfied: google-api-core[grpc]<2.0.0dev,>=1.14.0 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from google-cloud-pubsub<1.1.0,>=0.39.0; extra == \"gcp\"->apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.16.0)\n",
      "Requirement already satisfied: grpc-google-iam-v1<0.13dev,>=0.12.3 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from google-cloud-pubsub<1.1.0,>=0.39.0; extra == \"gcp\"->apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (0.12.3)\n",
      "Requirement already satisfied: google-resumable-media<0.5.0dev,>=0.3.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from google-cloud-bigquery<1.18.0,>=1.6.0; extra == \"gcp\"->apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (0.4.1)\n",
      "Requirement already satisfied: smmap<4,>=3.0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from gitdb<5,>=4.0.1->GitPython>=1.0.0->wandb==0.8.36->ml4ir) (3.0.4)\n",
      "Requirement already satisfied: monotonic>=0.1 in /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages (from fasteners>=0.14->google-apitools<0.5.29,>=0.5.28; extra == \"gcp\"->apache-beam[gcp]<3,>=2.16->tensorflow-transform==0.15.0->ml4ir) (1.5)\n",
      "\u001b[31mERROR: tfx-bsl 0.15.3 has requirement absl-py<0.9,>=0.7, but you'll have absl-py 0.9.0 which is incompatible.\u001b[0m\n",
      "\u001b[31mERROR: tfx-bsl 0.15.3 has requirement apache-beam[gcp]<2.17,>=2.16, but you'll have apache-beam 2.19.0 which is incompatible.\u001b[0m\n",
      "\u001b[31mERROR: tfx-bsl 0.15.3 has requirement pyarrow<0.15.0,>=0.14.0, but you'll have pyarrow 0.15.1 which is incompatible.\u001b[0m\n",
      "\u001b[31mERROR: tensorflow-transform 0.15.0 has requirement absl-py<0.9,>=0.7, but you'll have absl-py 0.9.0 which is incompatible.\u001b[0m\n",
      "\u001b[31mERROR: apache-beam 2.19.0 has requirement dill<0.3.2,>=0.3.1.1, but you'll have dill 0.3.0 which is incompatible.\u001b[0m\n",
      "\u001b[31mERROR: apache-beam 2.19.0 has requirement httplib2<=0.12.0,>=0.8, but you'll have httplib2 0.17.0 which is incompatible.\u001b[0m\n",
      "Installing collected packages: pyarrow\n",
      "  Attempting uninstall: pyarrow\n",
      "    Found existing installation: pyarrow 0.14.1\n",
      "    Uninstalling pyarrow-0.14.1:\n",
      "      Successfully uninstalled pyarrow-0.14.1\n",
      "Successfully installed pyarrow-0.15.1\n"
     ]
    }
   ],
   "source": [
    "!pip install ml4ir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 0: Looking at the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:root:Logger is initialized...\n",
      "INFO:root:Loading dataframe from path : ../ml4ir/applications/classification/tests/data/csv/train/file_0.csv\n",
      "INFO:root:(700, 5)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query_key</th>\n",
       "      <th>query_text</th>\n",
       "      <th>domain_id</th>\n",
       "      <th>user_context</th>\n",
       "      <th>entity_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>query_id_0</td>\n",
       "      <td>yourself</td>\n",
       "      <td>Y</td>\n",
       "      <td>EEE,BBB,AAA,GGG,FFF,FFF,AAA,CCC,CCC,FFF,FFF,DD...</td>\n",
       "      <td>AAA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>query_id_1</td>\n",
       "      <td>struck entire the come thanks</td>\n",
       "      <td>B</td>\n",
       "      <td>CCC,CCC,AAA</td>\n",
       "      <td>DDD</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>query_id_2</td>\n",
       "      <td>sick unfold am prince you</td>\n",
       "      <td>Q</td>\n",
       "      <td>DDD,FFF,AAA,GGG,GGG,HHH,GGG,FFF,AAA,CCC,BBB,HH...</td>\n",
       "      <td>AAA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>query_id_3</td>\n",
       "      <td>bitter twelve if upon of him</td>\n",
       "      <td>U</td>\n",
       "      <td>AAA,FFF,DDD,GGG,AAA,EEE,HHH,DDD,HHH,CCC,CCC,HHH</td>\n",
       "      <td>DDD</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>query_id_4</td>\n",
       "      <td>tragedy</td>\n",
       "      <td>O</td>\n",
       "      <td>AAA,EEE,FFF,EEE,GGG,GGG,AAA</td>\n",
       "      <td>EEE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>query_id_5</td>\n",
       "      <td>quiet</td>\n",
       "      <td>W</td>\n",
       "      <td>BBB,GGG,AAA</td>\n",
       "      <td>AAA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>query_id_6</td>\n",
       "      <td>friends bid thee hamlet most</td>\n",
       "      <td>P</td>\n",
       "      <td>FFF,GGG</td>\n",
       "      <td>AAA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>query_id_7</td>\n",
       "      <td>bid his for entire dane hear</td>\n",
       "      <td>L</td>\n",
       "      <td>BBB,AAA,GGG,EEE,CCC,EEE,GGG,FFF</td>\n",
       "      <td>EEE</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>query_id_8</td>\n",
       "      <td>guard platform watch</td>\n",
       "      <td>Z</td>\n",
       "      <td>GGG,GGG,AAA,CCC,HHH,AAA,FFF,HHH,HHH,EEE,GGG,AA...</td>\n",
       "      <td>AAA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>query_id_9</td>\n",
       "      <td>had stand whos</td>\n",
       "      <td>I</td>\n",
       "      <td>AAA,BBB,DDD,DDD,AAA,DDD,EEE,DDD,BBB,FFF,BBB,FF...</td>\n",
       "      <td>DDD</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    query_key                     query_text domain_id  \\\n",
       "0  query_id_0                       yourself         Y   \n",
       "1  query_id_1  struck entire the come thanks         B   \n",
       "2  query_id_2      sick unfold am prince you         Q   \n",
       "3  query_id_3   bitter twelve if upon of him         U   \n",
       "4  query_id_4                        tragedy         O   \n",
       "5  query_id_5                          quiet         W   \n",
       "6  query_id_6   friends bid thee hamlet most         P   \n",
       "7  query_id_7   bid his for entire dane hear         L   \n",
       "8  query_id_8           guard platform watch         Z   \n",
       "9  query_id_9                 had stand whos         I   \n",
       "\n",
       "                                        user_context entity_id  \n",
       "0  EEE,BBB,AAA,GGG,FFF,FFF,AAA,CCC,CCC,FFF,FFF,DD...       AAA  \n",
       "1                                        CCC,CCC,AAA       DDD  \n",
       "2  DDD,FFF,AAA,GGG,GGG,HHH,GGG,FFF,AAA,CCC,BBB,HH...       AAA  \n",
       "3    AAA,FFF,DDD,GGG,AAA,EEE,HHH,DDD,HHH,CCC,CCC,HHH       DDD  \n",
       "4                        AAA,EEE,FFF,EEE,GGG,GGG,AAA       EEE  \n",
       "5                                        BBB,GGG,AAA       AAA  \n",
       "6                                            FFF,GGG       AAA  \n",
       "7                    BBB,AAA,GGG,EEE,CCC,EEE,GGG,FFF       EEE  \n",
       "8  GGG,GGG,AAA,CCC,HHH,AAA,FFF,HHH,HHH,EEE,GGG,AA...       AAA  \n",
       "9  AAA,BBB,DDD,DDD,AAA,DDD,EEE,DDD,BBB,FFF,BBB,FF...       DDD  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ml4ir.base.io.file_io import FileIO\n",
    "from ml4ir.base.io.local_io import LocalIO\n",
    "import glob\n",
    "import logging\n",
    "import pandas as pd\n",
    "import os\n",
    "import tensorflow as tf\n",
    "\n",
    "# Pandas options\n",
    "pd.set_option('display.max_rows', None)\n",
    "pd.set_option('display.max_columns', None)\n",
    "\n",
    "# Setup logging\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.DEBUG)\n",
    "tf.get_logger().setLevel('INFO')\n",
    "logging.debug(\"Logger is initialized...\")\n",
    "\n",
    "# Define FileIO\n",
    "file_io: FileIO = LocalIO(logger)\n",
    "\n",
    "# Load data\n",
    "CSV_DATA_DIR = '../ml4ir/applications/classification/tests/data/csv'\n",
    "\n",
    "df = file_io.read_df(os.path.join(CSV_DATA_DIR, \"train\", \"file_0.csv\"))\n",
    "\n",
    "logger.info(df.shape)\n",
    "\n",
    "df.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the FeatureConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:root:{\n",
      "    \"query_key\": {\n",
      "        \"name\": \"query_key\",\n",
      "        \"node_name\": \"query_key\",\n",
      "        \"trainable\": false,\n",
      "        \"dtype\": \"string\",\n",
      "        \"log_at_inference\": true,\n",
      "        \"feature_layer_info\": {\n",
      "            \"type\": \"numeric\",\n",
      "            \"shape\": null\n",
      "        },\n",
      "        \"serving_info\": {\n",
      "            \"name\": \"query_key\",\n",
      "            \"required\": false\n",
      "        },\n",
      "        \"default_value\": \"\"\n",
      "    },\n",
      "    \"label\": {\n",
      "        \"name\": \"entity_id\",\n",
      "        \"node_name\": \"entity_id\",\n",
      "        \"trainable\": false,\n",
      "        \"dtype\": \"string\",\n",
      "        \"shape\": [\n",
      "            1,\n",
      "            null\n",
      "        ],\n",
      "        \"log_at_inference\": true,\n",
      "        \"preprocessing_info\": [\n",
      "            {\n",
      "                \"fn\": \"one_hot_vectorize_label\",\n",
      "                \"args\": {\n",
      "                    \"vocabulary_file\": \"../ml4ir/applications/classification/tests/data/csv/../configs/entity_id_vocab.csv\",\n",
      "                    \"num_oov_buckets\": 1\n",
      "                }\n",
      "            }\n",
      "        ],\n",
      "        \"feature_layer_info\": {\n",
      "            \"type\": \"numeric\",\n",
      "            \"fn\": \"categorical_indicator_with_vocabulary_file\",\n",
      "            \"args\": {\n",
      "                \"vocabulary_file\": \"../ml4ir/applications/classification/tests/data/csv/../configs/entity_id_vocab.csv\",\n",
      "                \"num_oov_buckets\": 1\n",
      "            }\n",
      "        },\n",
      "        \"serving_info\": {\n",
      "            \"name\": \"entity_id\",\n",
      "            \"required\": false\n",
      "        },\n",
      "        \"default_value\": \"\"\n",
      "    },\n",
      "    \"features\": [\n",
      "        {\n",
      "            \"name\": \"query_text\",\n",
      "            \"node_name\": \"query_text\",\n",
      "            \"trainable\": true,\n",
      "            \"dtype\": \"string\",\n",
      "            \"log_at_inference\": true,\n",
      "            \"feature_layer_info\": {\n",
      "                \"type\": \"numeric\",\n",
      "                \"shape\": null,\n",
      "                \"fn\": \"bytes_sequence_to_encoding_bilstm\",\n",
      "                \"args\": {\n",
      "                    \"encoding_type\": \"bilstm\",\n",
      "                    \"encoding_size\": 128,\n",
      "                    \"embedding_size\": 128,\n",
      "                    \"max_length\": 20\n",
      "                }\n",
      "            },\n",
      "            \"preprocessing_info\": [\n",
      "                {\n",
      "                    \"fn\": \"preprocess_text\",\n",
      "                    \"args\": {\n",
      "                        \"remove_punctuation\": true,\n",
      "                        \"to_lower\": true\n",
      "                    }\n",
      "                }\n",
      "            ],\n",
      "            \"serving_info\": {\n",
      "                \"name\": \"query_text\",\n",
      "                \"required\": true\n",
      "            },\n",
      "            \"default_value\": \"\"\n",
      "        },\n",
      "        {\n",
      "            \"name\": \"domain_id\",\n",
      "            \"node_name\": \"domain_id\",\n",
      "            \"trainable\": true,\n",
      "            \"dtype\": \"string\",\n",
      "            \"log_at_inference\": true,\n",
      "            \"is_group_metric_key\": true,\n",
      "            \"feature_layer_info\": {\n",
      "                \"type\": \"numeric\",\n",
      "                \"shape\": null,\n",
      "                \"fn\": \"categorical_embedding_with_vocabulary_file\",\n",
      "                \"args\": {\n",
      "                    \"vocabulary_file\": \"../ml4ir/applications/classification/tests/data/csv/../configs/domain_id_vocab.csv\",\n",
      "                    \"embedding_size\": 64,\n",
      "                    \"default_value\": -1,\n",
      "                    \"num_oov_buckets\": 1\n",
      "                }\n",
      "            },\n",
      "            \"serving_info\": {\n",
      "                \"name\": \"domain_id\",\n",
      "                \"required\": true\n",
      "            },\n",
      "            \"default_value\": \"\"\n",
      "        },\n",
      "        {\n",
      "            \"name\": \"user_context\",\n",
      "            \"node_name\": \"user_context\",\n",
      "            \"trainable\": true,\n",
      "            \"dtype\": \"string\",\n",
      "            \"shape\": [\n",
      "                1,\n",
      "                20\n",
      "            ],\n",
      "            \"log_at_inference\": true,\n",
      "            \"is_group_metric_key\": true,\n",
      "            \"preprocessing_info\": [\n",
      "                {\n",
      "                    \"fn\": \"split_string\",\n",
      "                    \"args\": {\n",
      "                        \"split_char\": \",\",\n",
      "                        \"max_length\": 20\n",
      "                    }\n",
      "                }\n",
      "            ],\n",
      "            \"feature_layer_info\": {\n",
      "                \"type\": \"numeric\",\n",
      "                \"shape\": null,\n",
      "                \"fn\": \"smart_scope_embedding_bilstm_encoding\",\n",
      "                \"args\": {\n",
      "                    \"vocabulary_file\": \"../ml4ir/applications/classification/tests/data/csv/../configs/entity_id_vocab.csv\",\n",
      "                    \"embedding_size\": 64,\n",
      "                    \"encoding_size\": 64,\n",
      "                    \"num_oov_buckets\": 1,\n",
      "                    \"max_length\": 20\n",
      "                }\n",
      "            },\n",
      "            \"serving_info\": {\n",
      "                \"name\": \"user_context\",\n",
      "                \"required\": true\n",
      "            },\n",
      "            \"default_value\": \"\"\n",
      "        }\n",
      "    ]\n",
      "}\n",
      "INFO:root:Feature config loaded successfully\n",
      "INFO:root:Trainable Features : \n",
      "query_text\n",
      "domain_id\n",
      "user_context\n",
      "INFO:root:Label : entity_id\n",
      "INFO:root:Metadata Features : \n",
      "query_key\n",
      "entity_id\n"
     ]
    }
   ],
   "source": [
    "# Set up the feature configurations\n",
    "from ml4ir.base.features.feature_config import FeatureConfig, ExampleFeatureConfig\n",
    "from ml4ir.base.config.keys import TFRecordTypeKey\n",
    "import json\n",
    "import yaml\n",
    "\n",
    "feature_config_yaml = '''\n",
    "query_key: \n",
    "  name: query_key\n",
    "  node_name: query_key\n",
    "  trainable: false\n",
    "  dtype: string\n",
    "  log_at_inference: true\n",
    "  feature_layer_info:\n",
    "    type: numeric\n",
    "    shape: null\n",
    "  serving_info:\n",
    "    name: query_key\n",
    "    required: false\n",
    "  default_value: \"\"\n",
    "label:\n",
    "  name: entity_id\n",
    "  node_name: entity_id\n",
    "  trainable: false\n",
    "  dtype: string\n",
    "  shape: \n",
    "    - 1\n",
    "    - null\n",
    "  log_at_inference: true\n",
    "  preprocessing_info:\n",
    "    - fn: one_hot_vectorize_label\n",
    "      args:\n",
    "        vocabulary_file: {0}\n",
    "        num_oov_buckets: 1\n",
    "  feature_layer_info:\n",
    "    type: numeric\n",
    "    fn: categorical_indicator_with_vocabulary_file\n",
    "    args:\n",
    "      vocabulary_file: {0}\n",
    "      num_oov_buckets: 1\n",
    "  serving_info:\n",
    "    name: entity_id\n",
    "    required: false\n",
    "  default_value: \"\"\n",
    "features:\n",
    "  - name: query_text\n",
    "    node_name: query_text\n",
    "    trainable: true\n",
    "    dtype: string\n",
    "    log_at_inference: true\n",
    "    feature_layer_info:\n",
    "      type: numeric\n",
    "      shape: null\n",
    "      fn: bytes_sequence_to_encoding_bilstm\n",
    "      args:\n",
    "        encoding_type: bilstm\n",
    "        encoding_size: 128\n",
    "        embedding_size: 128\n",
    "        max_length: 20\n",
    "    preprocessing_info:\n",
    "      - fn: preprocess_text\n",
    "        args:\n",
    "          remove_punctuation: true\n",
    "          to_lower: true\n",
    "    serving_info:\n",
    "      name: query_text\n",
    "      required: true\n",
    "    default_value: \"\"\n",
    "  - name: domain_id\n",
    "    node_name: domain_id\n",
    "    trainable: true\n",
    "    dtype: string\n",
    "    log_at_inference: true\n",
    "    is_group_metric_key: true\n",
    "    feature_layer_info:\n",
    "      type: numeric\n",
    "      shape: null\n",
    "      # fn: categorical_embedding_with_hash_buckets\n",
    "      # args:\n",
    "      #   num_hash_buckets: 4\n",
    "      #   hash_bucket_size: 64\n",
    "      #   embedding_size: 32\n",
    "      #   merge_mode: concat\n",
    "      fn: categorical_embedding_with_vocabulary_file\n",
    "      args:\n",
    "        vocabulary_file: {1}\n",
    "        embedding_size: 64\n",
    "        default_value: -1\n",
    "        num_oov_buckets: 1\n",
    "    serving_info:\n",
    "      name: domain_id\n",
    "      required: true\n",
    "    default_value: \"\"\n",
    "  - name: user_context\n",
    "    node_name: user_context\n",
    "    trainable: true\n",
    "    dtype: string\n",
    "    shape:\n",
    "      - 1\n",
    "      - {2}\n",
    "    log_at_inference: true\n",
    "    is_group_metric_key: true\n",
    "    preprocessing_info:\n",
    "      - fn: split_string\n",
    "        args:\n",
    "          split_char: \",\"\n",
    "          max_length: {2}\n",
    "    feature_layer_info:\n",
    "      type: numeric\n",
    "      shape: null\n",
    "      # fn: categorical_sequence_bilstm_embedding\n",
    "      # args:\n",
    "      #   num_hash_buckets: 4\n",
    "      #   hash_bucket_size: 64\n",
    "      #   embedding_size: 32\n",
    "      #   merge_mode: concat\n",
    "      fn: smart_scope_embedding_bilstm_encoding\n",
    "      args:\n",
    "        vocabulary_file: {0}\n",
    "        embedding_size: 64\n",
    "        encoding_size: 64\n",
    "        num_oov_buckets: 1\n",
    "        max_length: {2}\n",
    "    serving_info:\n",
    "      name: user_context\n",
    "      required: true\n",
    "    default_value: \"\"\n",
    "'''.format(\n",
    "    os.path.join(CSV_DATA_DIR, '../configs/vocabulary', 'entity_id.csv'),\n",
    "    os.path.join(CSV_DATA_DIR, '../configs/vocabulary', 'domain_id.csv'),\n",
    "    20\n",
    ")\n",
    "feature_config: ExampleFeatureConfig = FeatureConfig.get_instance(\n",
    "    tfrecord_type=TFRecordTypeKey.EXAMPLE,\n",
    "    feature_config_dict=yaml.safe_load(feature_config_yaml),\n",
    "    logger=logger)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Load the RelevanceDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Loading dataframe from path : ../ml4ir/applications/classification/tests/data/csv/../configs/entity_id_vocab.csv\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/tensorflow_core/python/feature_column/feature_column_v2.py:4276: IndicatorColumn._variable_shape (from tensorflow.python.feature_column.feature_column_v2) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "The old _FeatureColumn APIs are being deprecated. Please use the new FeatureColumn APIs instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/tensorflow_core/python/feature_column/feature_column_v2.py:4276: IndicatorColumn._variable_shape (from tensorflow.python.feature_column.feature_column_v2) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "The old _FeatureColumn APIs are being deprecated. Please use the new FeatureColumn APIs instead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/tensorflow_core/python/feature_column/feature_column_v2.py:4331: IdentityCategoricalColumn._num_buckets (from tensorflow.python.feature_column.feature_column_v2) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "The old _FeatureColumn APIs are being deprecated. Please use the new FeatureColumn APIs instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/tensorflow_core/python/feature_column/feature_column_v2.py:4331: IdentityCategoricalColumn._num_buckets (from tensorflow.python.feature_column.feature_column_v2) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "The old _FeatureColumn APIs are being deprecated. Please use the new FeatureColumn APIs instead.\n",
      "INFO:root:1 files found under ../ml4ir/applications/classification/tests/data/csv/train\n",
      "INFO:root:Reading 1 files from [../ml4ir/applications/classification/tests/data/csv/train/file_0.csv, ..\n",
      "INFO:root:Loading dataframe from path : ../ml4ir/applications/classification/tests/data/csv/train/file_0.csv\n",
      "INFO:root:Writing SequenceExample protobufs to : ../ml4ir/applications/classification/tests/data/csv/tfrecord/train/file_0.tfrecord\n",
      "INFO:root:1 files found under ../ml4ir/applications/classification/tests/data/csv/tfrecord/train\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'query_key': FixedLenFeature(shape=[], dtype='string', default_value=''), 'entity_id': FixedLenFeature(shape=[], dtype='string', default_value=''), 'query_text': FixedLenFeature(shape=[], dtype='string', default_value=''), 'domain_id': FixedLenFeature(shape=[], dtype='string', default_value=''), 'user_context': FixedLenFeature(shape=[], dtype='string', default_value='')}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Created TFRecordDataset from SequenceExample protobufs from 1 files : ['../ml4ir/applications/classification/tests/data/\n",
      "INFO:root:1 files found under ../ml4ir/applications/classification/tests/data/csv/validation\n",
      "INFO:root:Reading 1 files from [../ml4ir/applications/classification/tests/data/csv/validation/file_0.csv, ..\n",
      "INFO:root:Loading dataframe from path : ../ml4ir/applications/classification/tests/data/csv/validation/file_0.csv\n",
      "INFO:root:Writing SequenceExample protobufs to : ../ml4ir/applications/classification/tests/data/csv/tfrecord/validation/file_0.tfrecord\n",
      "INFO:root:1 files found under ../ml4ir/applications/classification/tests/data/csv/tfrecord/validation\n",
      "INFO:root:Created TFRecordDataset from SequenceExample protobufs from 1 files : ['../ml4ir/applications/classification/tests/data/\n",
      "INFO:root:1 files found under ../ml4ir/applications/classification/tests/data/csv/test\n",
      "INFO:root:Reading 1 files from [../ml4ir/applications/classification/tests/data/csv/test/file_0.csv, ..\n",
      "INFO:root:Loading dataframe from path : ../ml4ir/applications/classification/tests/data/csv/test/file_0.csv\n",
      "INFO:root:Writing SequenceExample protobufs to : ../ml4ir/applications/classification/tests/data/csv/tfrecord/test/file_0.tfrecord\n",
      "INFO:root:1 files found under ../ml4ir/applications/classification/tests/data/csv/tfrecord/test\n",
      "INFO:root:Created TFRecordDataset from SequenceExample protobufs from 1 files : ['../ml4ir/applications/classification/tests/data/\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'query_key': FixedLenFeature(shape=[], dtype='string', default_value=''), 'entity_id': FixedLenFeature(shape=[], dtype='string', default_value=''), 'query_text': FixedLenFeature(shape=[], dtype='string', default_value=''), 'domain_id': FixedLenFeature(shape=[], dtype='string', default_value=''), 'user_context': FixedLenFeature(shape=[], dtype='string', default_value='')}\n",
      "{'query_key': FixedLenFeature(shape=[], dtype='string', default_value=''), 'entity_id': FixedLenFeature(shape=[], dtype='string', default_value=''), 'query_text': FixedLenFeature(shape=[], dtype='string', default_value=''), 'domain_id': FixedLenFeature(shape=[], dtype='string', default_value=''), 'user_context': FixedLenFeature(shape=[], dtype='string', default_value='')}\n",
      "<BatchDataset shapes: ({query_key: (128, 1), query_text: (128, 1), domain_id: (128, 1), user_context: (128, 1, 20)}, (128, 1, 10)), types: ({query_key: tf.string, query_text: tf.string, domain_id: tf.string, user_context: tf.string}, tf.float32)>\n",
      "<BatchDataset shapes: ({query_key: (128, 1), query_text: (128, 1), domain_id: (128, 1), user_context: (128, 1, 20)}, (128, 1, 10)), types: ({query_key: tf.string, query_text: tf.string, domain_id: tf.string, user_context: tf.string}, tf.float32)>\n",
      "<BatchDataset shapes: ({query_key: (128, 1), query_text: (128, 1), domain_id: (128, 1), user_context: (128, 1, 20)}, (128, 1, 10)), types: ({query_key: tf.string, query_text: tf.string, domain_id: tf.string, user_context: tf.string}, tf.float32)>\n"
     ]
    }
   ],
   "source": [
    "from ml4ir.base.data.relevance_dataset import RelevanceDataset\n",
    "from ml4ir.base.config.keys import DataFormatKey\n",
    "from ml4ir.base.features.feature_fns.categorical import categorical_indicator_with_vocabulary_file\n",
    "from tensorflow import image\n",
    "from tensorflow import print as tfprint\n",
    "import tensorflow as tf\n",
    "\n",
    "\n",
    "def get_one_hot_vectorizer(feature_info, file_io: FileIO):\n",
    "    label_str = tf.keras.Input(shape=(1,), dtype=tf.string)\n",
    "    label_one_hot = categorical_indicator_with_vocabulary_file(label_str, feature_info, file_io)\n",
    "    one_hot_vectorizer = tf.keras.Model(inputs=label_str, outputs=label_one_hot)\n",
    "    \n",
    "    @tf.function\n",
    "    def one_hot_vectorize(feature_tensor, **kwargs):\n",
    "        return tf.squeeze(one_hot_vectorizer(feature_tensor), axis=[0])\n",
    "    \n",
    "    return one_hot_vectorize\n",
    "\n",
    "@tf.function\n",
    "def split_string(feature_tensor, split_char=\",\", max_length=20, **kwargs):\n",
    "    tokens = tf.strings.split(feature_tensor, sep=split_char).to_tensor()\n",
    "    padded_tokens = image.pad_to_bounding_box(tf.expand_dims(tokens[:, :max_length], axis=-1),\n",
    "                                              offset_height=0,\n",
    "                                              offset_width=0,\n",
    "                                              target_height=1,\n",
    "                                              target_width=max_length)\n",
    "    padded_tokens = tf.squeeze(padded_tokens, axis=-1)\n",
    "    return padded_tokens\n",
    "\n",
    "preprocessing_keys_to_fns = {\n",
    "    \"one_hot_vectorize_label\": get_one_hot_vectorizer(feature_config.get_label(), file_io),\n",
    "    \"split_string\": split_string\n",
    "}\n",
    "    \n",
    "relevance_dataset = RelevanceDataset(\n",
    "        data_dir=CSV_DATA_DIR,\n",
    "        data_format=DataFormatKey.CSV,\n",
    "        feature_config=feature_config,\n",
    "        tfrecord_type=TFRecordTypeKey.EXAMPLE,\n",
    "        batch_size=128,\n",
    "        preprocessing_keys_to_fns=preprocessing_keys_to_fns,\n",
    "        file_io=file_io,\n",
    "        logger=logger\n",
    "    )\n",
    "\n",
    "tfprint(relevance_dataset.train)\n",
    "tfprint(relevance_dataset.validation)\n",
    "tfprint(relevance_dataset.test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "~~ query_key ~~\n",
      "tf.Tensor(\n",
      "[[b'query_id_0']\n",
      " [b'query_id_1']\n",
      " [b'query_id_2']\n",
      " [b'query_id_3']\n",
      " [b'query_id_4']], shape=(5, 1), dtype=string)\n",
      "\n",
      "~~ query_text ~~\n",
      "tf.Tensor(\n",
      "[[b'yourself']\n",
      " [b'struck entire the come thanks']\n",
      " [b'sick unfold am prince you']\n",
      " [b'bitter twelve if upon of him']\n",
      " [b'tragedy']], shape=(5, 1), dtype=string)\n",
      "\n",
      "~~ domain_id ~~\n",
      "tf.Tensor(\n",
      "[[b'Y']\n",
      " [b'B']\n",
      " [b'Q']\n",
      " [b'U']\n",
      " [b'O']], shape=(5, 1), dtype=string)\n",
      "\n",
      "~~ user_context ~~\n",
      "tf.Tensor(\n",
      "[[[b'EEE' b'BBB' b'AAA' b'GGG' b'FFF' b'FFF' b'AAA' b'CCC' b'CCC' b'FFF'\n",
      "   b'FFF' b'DDD' b'CCC' b'AAA' b'' b'' b'' b'' b'' b'']]\n",
      "\n",
      " [[b'CCC' b'CCC' b'AAA' b'' b'' b'' b'' b'' b'' b'' b'' b'' b'' b'' b''\n",
      "   b'' b'' b'' b'' b'']]\n",
      "\n",
      " [[b'DDD' b'FFF' b'AAA' b'GGG' b'GGG' b'HHH' b'GGG' b'FFF' b'AAA' b'CCC'\n",
      "   b'BBB' b'HHH' b'EEE' b'CCC' b'FFF' b'FFF' b'' b'' b'' b'']]\n",
      "\n",
      " [[b'AAA' b'FFF' b'DDD' b'GGG' b'AAA' b'EEE' b'HHH' b'DDD' b'HHH' b'CCC'\n",
      "   b'CCC' b'HHH' b'' b'' b'' b'' b'' b'' b'' b'']]\n",
      "\n",
      " [[b'AAA' b'EEE' b'FFF' b'EEE' b'GGG' b'GGG' b'AAA' b'' b'' b'' b'' b''\n",
      "   b'' b'' b'' b'' b'' b'' b'' b'']]], shape=(5, 1, 20), dtype=string)\n",
      "\n",
      "~~ entity ~~\n",
      "tf.Tensor(\n",
      "[[[0. 1. 0. 0. 0. 0. 0. 0. 0. 0.]]\n",
      "\n",
      " [[0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]]\n",
      "\n",
      " [[0. 1. 0. 0. 0. 0. 0. 0. 0. 0.]]\n",
      "\n",
      " [[0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]]\n",
      "\n",
      " [[0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]]], shape=(5, 1, 10), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "batch = next(iter(relevance_dataset.train))\n",
    "for col in batch[0]:\n",
    "    print(\"\\n~~ {} ~~\".format(col))\n",
    "    print(batch[0][col][:5])\n",
    "print(\"\\n~~ {} ~~\".format(\"entity\"))\n",
    "print(batch[1][:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Define the InteractionModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ml4ir.base.model.scoring.interaction_model import InteractionModel, UnivariateInteractionModel\n",
    "from ml4ir.base.config.keys import TFRecordTypeKey\n",
    "from ml4ir.base.features.feature_fns.categorical import VocabLookup\n",
    "\n",
    "from tensorflow import feature_column\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "\n",
    "# Define custom feature layer ops\n",
    "def smart_scope_embedding_bilstm_encoding(feature_tensor, feature_info, file_io: FileIO):\n",
    "    args = feature_info.get(\"feature_layer_info\")[\"args\"]\n",
    "    \n",
    "    vocabulary_df = file_io.read_df(args[\"vocabulary_file\"])\n",
    "    vocabulary_keys = vocabulary_df[\"key\"].fillna(feature_info[\"default_value\"]).values\n",
    "    vocabulary_ids = (\n",
    "        vocabulary_df[\"id\"].values if \"id\" in vocabulary_df else list(range(len(vocabulary_keys)))\n",
    "    )\n",
    "\n",
    "    num_oov_buckets = args.get(\"num_oov_buckets\", 1)\n",
    "    vocabulary_size = len(set(vocabulary_ids))\n",
    "    lookup_table = VocabLookup(\n",
    "        vocabulary_keys=vocabulary_keys,\n",
    "        vocabulary_ids=vocabulary_ids,\n",
    "        num_oov_buckets=num_oov_buckets,\n",
    "        feature_name=feature_info.get(\"node_name\", feature_info[\"name\"]),\n",
    "    )\n",
    "    categorical_indices = lookup_table(feature_tensor)\n",
    "    categorical_embeddings = layers.Embedding(\n",
    "                                    input_dim=vocabulary_size + num_oov_buckets,\n",
    "                                    output_dim=args[\"embedding_size\"],\n",
    "                                    mask_zero=True,\n",
    "                                    input_length=args.get(\"max_length\")\n",
    "                                )(categorical_indices)\n",
    "\n",
    "    encoding = layers.Bidirectional(\n",
    "                    layers.LSTM(\n",
    "                        units=int(args[\"encoding_size\"] / 2), return_sequences=False\n",
    "                    ),\n",
    "                    merge_mode=\"concat\",\n",
    "                )(tf.squeeze(categorical_embeddings, axis=1))\n",
    "    encoding = tf.expand_dims(encoding, name=\"smart_scope_encoding\", axis=1)\n",
    "    \n",
    "    return encoding\n",
    "\n",
    "feature_layer_fns = {\n",
    "    \"smart_scope_embedding_bilstm_encoding\": smart_scope_embedding_bilstm_encoding,\n",
    "}\n",
    "\n",
    "interaction_model: InteractionModel = UnivariateInteractionModel(\n",
    "                                            feature_config=feature_config,\n",
    "                                            feature_layer_keys_to_fns=feature_layer_fns,\n",
    "                                            tfrecord_type=TFRecordTypeKey.EXAMPLE,\n",
    "                                            file_io=file_io)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Define the Loss and Scoring Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Reading YAML file from : ../ml4ir/applications/classification/tests/data/csv/../configs/model_config.yaml\n",
      "INFO:root:{\n",
      "    \"architecture_key\": \"dnn\",\n",
      "    \"layers\": [\n",
      "        {\n",
      "            \"type\": \"dense\",\n",
      "            \"name\": \"first_dense\",\n",
      "            \"units\": 256,\n",
      "            \"activation\": \"relu\"\n",
      "        },\n",
      "        {\n",
      "            \"type\": \"dropout\",\n",
      "            \"name\": \"first_dropout\",\n",
      "            \"rate\": 0.3\n",
      "        },\n",
      "        {\n",
      "            \"type\": \"dense\",\n",
      "            \"name\": \"second_dense\",\n",
      "            \"units\": 64,\n",
      "            \"activation\": \"relu\"\n",
      "        },\n",
      "        {\n",
      "            \"type\": \"dropout\",\n",
      "            \"name\": \"second_dropout\",\n",
      "            \"rate\": 0.0\n",
      "        },\n",
      "        {\n",
      "            \"type\": \"dense\",\n",
      "            \"name\": \"final_dense\",\n",
      "            \"units\": 10,\n",
      "            \"activation\": null\n",
      "        }\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from ml4ir.base.model.scoring.scoring_model import ScorerBase, RelevanceScorer\n",
    "from ml4ir.base.model.losses.loss_base import RelevanceLossBase\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras import losses\n",
    "from ml4ir.base.features.feature_fns.categorical import categorical_indicator_with_vocabulary_file\n",
    "\n",
    "class CustomCategoricalCrossEntropy(RelevanceLossBase):\n",
    "        \n",
    "    def get_loss_fn(self, **kwargs):\n",
    "        \"\"\"\n",
    "        Define a softmax cross entropy loss\n",
    "\n",
    "        \"\"\"\n",
    "        cce = losses.CategoricalCrossentropy(reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)\n",
    "\n",
    "        def _loss_fn(y_true, y_pred):\n",
    "            # NOTE: Can use any of the metadata features to qualify your loss here\n",
    "            return cce(y_true, y_pred)\n",
    "\n",
    "        return _loss_fn\n",
    "\n",
    "    def get_final_activation_op(self, output_name):\n",
    "        return lambda logits, mask: layers.Activation(\"softmax\", name=output_name)(logits)\n",
    "\n",
    "scorer: ScorerBase = RelevanceScorer.from_model_config_file(\n",
    "    model_config_file=os.path.join(CSV_DATA_DIR, '../configs/model_config.yaml'),\n",
    "    interaction_model=interaction_model,\n",
    "    loss=CustomCategoricalCrossEntropy(),\n",
    "    output_name=\"relevance_score\",\n",
    "    file_io=file_io)\n",
    "    \n",
    "logger.info(json.dumps(scorer.model_config, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Define Metrics and Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras import metrics as kmetrics\n",
    "from ml4ir.applications.ranking.model.metrics.metric_factory import get_metric\n",
    "\n",
    "\n",
    "# metrics = ['categorical_accuracy', kmetrics.Precision, get_metric(\"MRR\"), get_metric(\"ACR\")]\n",
    "metrics = ['categorical_accuracy', kmetrics.Precision]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.optimizers import Optimizer\n",
    "from ml4ir.base.model.optimizer import get_optimizer\n",
    "from ml4ir.base.config.keys import OptimizerKey\n",
    "\n",
    "optimizer: Optimizer = get_optimizer(\n",
    "                optimizer_key=OptimizerKey.ADAM,\n",
    "                learning_rate=0.01,\n",
    "                learning_rate_decay=0.94,\n",
    "                learning_rate_decay_steps=1000,\n",
    "                gradient_clip_value=50,\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6: Putting it all together!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Loading dataframe from path : ../ml4ir/applications/classification/tests/data/csv/../configs/domain_id_vocab.csv\n",
      "INFO:root:Loading dataframe from path : ../ml4ir/applications/classification/tests/data/csv/../configs/entity_id_vocab.csv\n",
      "INFO:root:Model: \"model_1\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "query_text (InputLayer)         [(None, 1)]          0                                            \n",
      "__________________________________________________________________________________________________\n",
      "user_context (InputLayer)       [(None, 1, 20)]      0                                            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_DecodePaddedRaw (Te [(None, 1, 20)]      0           query_text[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "vocab_lookup_2 (VocabLookup)    (None, 1, 20)        0           user_context[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "domain_id (InputLayer)          [(None, 1)]          0                                            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Squeeze (TensorFlow [(None, 20)]         0           tf_op_layer_DecodePaddedRaw[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "embedding_1 (Embedding)         (None, 1, 20, 64)    640         vocab_lookup_2[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "vocab_lookup_1 (VocabLookup)    (None, 1)            0           domain_id[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "embedding (Embedding)           (None, 20, 128)      32768       tf_op_layer_Squeeze[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Squeeze_1 (TensorFl [(None, 20, 64)]     0           embedding_1[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "domain_id_embedding (DenseFeatu (None, 64)           1728        vocab_lookup_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "bidirectional (Bidirectional)   (None, 128)          98816       embedding[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "bidirectional_1 (Bidirectional) (None, 64)           24832       tf_op_layer_Squeeze_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_ExpandDims_2 (Tenso [(None, 1, 64)]      0           domain_id_embedding[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_ExpandDims_1 (Tenso [(None, 1, 128)]     0           bidirectional[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_smart_scope_encodin [(None, 1, 64)]      0           bidirectional_1[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_train_features (Ten [(None, 1, 256)]     0           tf_op_layer_ExpandDims_2[0][0]   \n",
      "                                                                 tf_op_layer_ExpandDims_1[0][0]   \n",
      "                                                                 tf_op_layer_smart_scope_encoding[\n",
      "__________________________________________________________________________________________________\n",
      "first_dense (Dense)             (None, 1, 256)       65792       tf_op_layer_train_features[0][0] \n",
      "__________________________________________________________________________________________________\n",
      "first_dropout (Dropout)         (None, 1, 256)       0           first_dense[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "second_dense (Dense)            (None, 1, 64)        16448       first_dropout[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "second_dropout (Dropout)        (None, 1, 64)        0           second_dense[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "final_dense (Dense)             (None, 1, 10)        650         second_dropout[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "query_key (InputLayer)          [(None, 1)]          0                                            \n",
      "__________________________________________________________________________________________________\n",
      "relevance_score (Activation)    (None, 1, 10)        0           final_dense[0][0]                \n",
      "==================================================================================================\n",
      "Total params: 241,674\n",
      "Trainable params: 241,674\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "from ml4ir.base.model.relevance_model import RelevanceModel\n",
    "from ml4ir.base.config.keys import OptimizerKey\n",
    "\n",
    "relevance_model = RelevanceModel(\n",
    "        feature_config=feature_config,\n",
    "        scorer=scorer,\n",
    "        metrics=metrics,\n",
    "        optimizer=optimizer,\n",
    "        tfrecord_type=TFRecordTypeKey.EXAMPLE,\n",
    "        output_name=\"entity_prediction_score\",\n",
    "        file_io=file_io,\n",
    "        logger=logger\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 7: Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Training Model\n",
      "INFO:root:Starting Epoch : 1\n",
      "INFO:root:{}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:[epoch: 1 | batch: 0] {'batch': 0, 'size': 128, 'loss': 2.2914057, 'categorical_accuracy': 0.140625, 'precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      1/Unknown - 7s 7s/step - loss: 2.2914 - categorical_accuracy: 0.1406 - precision: 0.0000e+00WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (0.436761). Check your callbacks.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Method (on_train_batch_end) is slow compared to the batch update (0.436761). Check your callbacks.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      5/Unknown - 9s 2s/step - loss: 1.9248 - categorical_accuracy: 0.2094 - precision: 0.0000e+00"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Evaluating Model\n",
      "INFO:root:Completed evaluating model\n",
      "INFO:root:None\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 00001: val_categorical_accuracy improved from -inf to 0.18750, saving model to ../models/checkpoint.tf\n",
      "WARNING:tensorflow:From /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1781: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/ashish.srinivasa/search_relevance/ml4ir/python/env/.ranking_venv3/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1781: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ../models/checkpoint.tf/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ../models/checkpoint.tf/assets\n",
      "INFO:root:End of Epoch 1\n",
      "INFO:root:{'loss': 1.9248095273971557, 'categorical_accuracy': 0.209375, 'precision': 0.0, 'val_loss': 1.7251148223876953, 'val_categorical_accuracy': 0.1875, 'val_precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "5/5 [==============================] - 59s 12s/step - loss: 1.9248 - categorical_accuracy: 0.2094 - precision: 0.0000e+00 - val_loss: 0.0000e+00 - val_categorical_accuracy: 0.0000e+00 - val_precision: 0.0000e+00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Starting Epoch : 2\n",
      "INFO:root:{}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:[epoch: 2 | batch: 0] {'batch': 0, 'size': 128, 'loss': 1.6742113, 'categorical_accuracy': 0.28125, 'precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4/5 [=======================>......] - ETA: 0s - loss: 1.6689 - categorical_accuracy: 0.2500 - precision: 0.0000e+00"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Evaluating Model\n",
      "INFO:root:Completed evaluating model\n",
      "INFO:root:None\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 00002: val_categorical_accuracy improved from 0.18750 to 0.19531, saving model to ../models/checkpoint.tf\n",
      "INFO:tensorflow:Assets written to: ../models/checkpoint.tf/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ../models/checkpoint.tf/assets\n",
      "INFO:root:End of Epoch 2\n",
      "INFO:root:{'loss': 1.6542035341262817, 'categorical_accuracy': 0.259375, 'precision': 0.0, 'val_loss': 1.633194923400879, 'val_categorical_accuracy': 0.1953125, 'val_precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "5/5 [==============================] - 45s 9s/step - loss: 1.6542 - categorical_accuracy: 0.2594 - precision: 0.0000e+00 - val_loss: 1.6332 - val_categorical_accuracy: 0.1953 - val_precision: 0.0000e+00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Starting Epoch : 3\n",
      "INFO:root:{}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:[epoch: 3 | batch: 0] {'batch': 0, 'size': 128, 'loss': 1.5382508, 'categorical_accuracy': 0.2890625, 'precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4/5 [=======================>......] - ETA: 0s - loss: 1.5667 - categorical_accuracy: 0.2812 - precision: 0.0000e+00"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Evaluating Model\n",
      "INFO:root:Completed evaluating model\n",
      "INFO:root:None\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 00003: val_categorical_accuracy improved from 0.19531 to 0.21875, saving model to ../models/checkpoint.tf\n",
      "INFO:tensorflow:Assets written to: ../models/checkpoint.tf/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ../models/checkpoint.tf/assets\n",
      "INFO:root:End of Epoch 3\n",
      "INFO:root:{'loss': 1.5646453619003295, 'categorical_accuracy': 0.2859375, 'precision': 0.4, 'val_loss': 1.6941783428192139, 'val_categorical_accuracy': 0.21875, 'val_precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "5/5 [==============================] - 49s 10s/step - loss: 1.5646 - categorical_accuracy: 0.2859 - precision: 0.4000 - val_loss: 1.6942 - val_categorical_accuracy: 0.2188 - val_precision: 0.0000e+00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Starting Epoch : 4\n",
      "INFO:root:{}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:[epoch: 4 | batch: 0] {'batch': 0, 'size': 128, 'loss': 1.4908521, 'categorical_accuracy': 0.3359375, 'precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4/5 [=======================>......] - ETA: 0s - loss: 1.5390 - categorical_accuracy: 0.3027 - precision: 0.2000"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Evaluating Model\n",
      "INFO:root:Completed evaluating model\n",
      "INFO:root:None\n",
      "INFO:root:End of Epoch 4\n",
      "INFO:root:{'loss': 1.5353684186935426, 'categorical_accuracy': 0.3046875, 'precision': 0.30769232, 'val_loss': 1.6862261295318604, 'val_categorical_accuracy': 0.1640625, 'val_precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 00004: val_categorical_accuracy did not improve from 0.21875\n",
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "5/5 [==============================] - 1s 217ms/step - loss: 1.5354 - categorical_accuracy: 0.3047 - precision: 0.3077 - val_loss: 1.6862 - val_categorical_accuracy: 0.1641 - val_precision: 0.0000e+00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Starting Epoch : 5\n",
      "INFO:root:{}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:[epoch: 5 | batch: 0] {'batch': 0, 'size': 128, 'loss': 1.4402865, 'categorical_accuracy': 0.4296875, 'precision': 0.0}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4/5 [=======================>......] - ETA: 0s - loss: 1.4802 - categorical_accuracy: 0.3711 - precision: 0.3333    "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Evaluating Model\n",
      "INFO:root:Completed evaluating model\n",
      "INFO:root:None\n",
      "INFO:root:End of Epoch 5\n",
      "INFO:root:{'loss': 1.4775662422180176, 'categorical_accuracy': 0.36875, 'precision': 0.52380955, 'val_loss': 1.7673120498657227, 'val_categorical_accuracy': 0.15625, 'val_precision': 0.1875}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 00005: val_categorical_accuracy did not improve from 0.21875\n",
      "Restoring model weights from the end of the best epoch.\n",
      "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r",
      "5/5 [==============================] - 1s 240ms/step - loss: 1.4776 - categorical_accuracy: 0.3688 - precision: 0.5238 - val_loss: 1.7673 - val_categorical_accuracy: 0.1562 - val_precision: 0.1875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Completed training model\n",
      "INFO:root:None\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 00005: early stopping\n"
     ]
    }
   ],
   "source": [
    "if not os.path.exists('../models'):\n",
    "    os.makedirs('../models')\n",
    "if not os.path.exists('../logs'):\n",
    "    os.makedirs('../logs')\n",
    "\n",
    "relevance_model.fit(relevance_dataset, \n",
    "                    num_epochs=5, \n",
    "                    models_dir='../models',\n",
    "                    logs_dir='../logs',\n",
    "                    monitor_metric='val_categorical_accuracy',\n",
    "                    monitor_mode='max')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 8: Save the Model to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ../models/entity_prediction/final/default/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ../models/entity_prediction/final/default/assets\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'query_text': FixedLenFeature(shape=[], dtype='string', default_value=''), 'domain_id': FixedLenFeature(shape=[], dtype='string', default_value=''), 'user_context': FixedLenFeature(shape=[], dtype='string', default_value='')}\n",
      "INFO:tensorflow:Assets written to: ../models/entity_prediction/final/tfrecord/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ../models/entity_prediction/final/tfrecord/assets\n",
      "INFO:root:Final model saved to : ../models/entity_prediction/final\n"
     ]
    }
   ],
   "source": [
    "MODEL_DIR = '../models/entity_prediction'\n",
    "if not os.path.exists(MODEL_DIR):\n",
    "    os.makedirs(MODEL_DIR)\n",
    "\n",
    "preprocessing_keys_to_fns = {\n",
    "    \"split_string\": split_string\n",
    "}\n",
    "relevance_model.save(\n",
    "    models_dir=MODEL_DIR,\n",
    "    preprocessing_keys_to_fns=preprocessing_keys_to_fns,\n",
    "    required_fields_only=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:\n",
      "\n",
      "signature_def['__saved_model_init_op']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['__saved_model_init_op'] tensor_info:\n",
      "        dtype: DT_INVALID\n",
      "        shape: unknown_rank\n",
      "        name: NoOp\n",
      "  Method name is: \n",
      "\n",
      "signature_def['serving_tfrecord']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['protos'] tensor_info:\n",
      "        dtype: DT_STRING\n",
      "        shape: (-1)\n",
      "        name: serving_tfrecord_protos:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['entity_prediction_score'] tensor_info:\n",
      "        dtype: DT_FLOAT\n",
      "        shape: (-1, 1, 10)\n",
      "        name: StatefulPartitionedCall_2:0\n",
      "  Method name is: tensorflow/serving/predict\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --dir ../models/entity_prediction/final/tfrecord/ --all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 9: Load the Model and make predictions!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example proto: \n",
      "b'\\n\\xbe\\x01\\n\\x14\\n\\tentity_id\\x12\\x07\\n\\x05\\n\\x03AAA\\n \\n\\nquery_text\\x12\\x12\\n\\x10\\n\\x0ea nay act hour\\n\\x12\\n\\tdomain_id\\x12\\x05\\n\\x03\\n\\x01G\\n\\x1b\\n\\tquery_key\\x12\\x0e\\n\\x0c\\n\\nquery_id_0\\nS\\n\\x0cuser_context\\x12C\\nA\\n?BBB,FFF,HHH,HHH,CCC,HHH,DDD,FFF,EEE,CCC,BBB,CCC,AAA,HHH,BBB,FFF'\n",
      "---------------------------------------\n",
      "\n",
      "\n",
      "Looking inside the proto:\n",
      "features {\n",
      "  feature {\n",
      "    key: \"domain_id\"\n",
      "    value {\n",
      "      bytes_list {\n",
      "        value: \"G\"\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "  feature {\n",
      "    key: \"entity_id\"\n",
      "    value {\n",
      "      bytes_list {\n",
      "        value: \"AAA\"\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "  feature {\n",
      "    key: \"query_key\"\n",
      "    value {\n",
      "      bytes_list {\n",
      "        value: \"query_id_0\"\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "  feature {\n",
      "    key: \"query_text\"\n",
      "    value {\n",
      "      bytes_list {\n",
      "        value: \"a nay act hour\"\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "  feature {\n",
      "    key: \"user_context\"\n",
      "    value {\n",
      "      bytes_list {\n",
      "        value: \"BBB,FFF,HHH,HHH,CCC,HHH,DDD,FFF,EEE,CCC,BBB,CCC,AAA,HHH,BBB,FFF\"\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n",
      "---------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "Predictions:\n",
      "{'entity_prediction_score': <tf.Tensor: id=395960, shape=(1, 1, 10), dtype=float32, numpy=\n",
      "array([[[3.0117732e-05, 1.9878587e-01, 2.9541317e-01, 1.9052488e-01,\n",
      "         6.4797543e-02, 2.5025532e-01, 1.9767896e-05, 5.1846171e-05,\n",
      "         5.1124658e-05, 7.0347443e-05]]], dtype=float32)>}\n",
      "---------------------------------------\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.keras import models as kmodels\n",
    "from tensorflow import data\n",
    "\n",
    "model = kmodels.load_model(\n",
    "    os.path.join(MODEL_DIR, 'final/tfrecord/'),\n",
    "    compile=False)\n",
    "infer_fn = model.signatures[\"serving_tfrecord\"]\n",
    "\n",
    "dataset = data.TFRecordDataset(\n",
    "    glob.glob(os.path.join(CSV_DATA_DIR, \"tfrecord\", \"test\", \"*.tfrecord\")))\n",
    "protos = next(iter(dataset.batch(5)))\n",
    "\n",
    "print(\"Example proto: \\n{}\".format(protos[0]))\n",
    "print(\"---------------------------------------\")\n",
    "\n",
    "print(\"\\n\\nLooking inside the proto:\")\n",
    "e = tf.train.Example()\n",
    "e.ParseFromString(protos[0].numpy())\n",
    "print(e)\n",
    "print(\"---------------------------------------\")\n",
    "\n",
    "print(\"\\n\\n\\nPredictions:\")\n",
    "print(infer_fn(protos=protos[:1]))\n",
    "print(\"---------------------------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Directory deleted : ../ml4ir/applications/classification/tests/data/csv/tfrecord\n"
     ]
    }
   ],
   "source": [
    "# Clean up directories\n",
    "# NOTE: Run only if you don't want to make any more predictions\n",
    "file_io.rm_dir(os.path.join(CSV_DATA_DIR, \"tfrecord\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![thanks](images/thats_all_folks.gif)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}