tfm_ranking

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "handling_sparse_features.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MHkcIl4E-Fgf"
      },
      "source": [
        "# Tutorial: TF-Ranking for sparse features\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V8tMYn22vtDV"
      },
      "source": [
        "This tutorial is an end-to-end walkthrough of training a TensorFlow Ranking (TF-Ranking) neural network model which incorporates sparse textual features.\n",
        "\n",
        "A Python script version of this code is available [here](https://github.com/tensorflow/ranking/blob/master/tensorflow_ranking/examples/tf_ranking_tfrecord.py). The script version supports flags for hyperparameters, and advanced use-cases like [Document Interaction Network](https://research.google/pubs/pub49364).\n",
        "\n",
        "TF-Ranking is a library for solving large scale ranking problems using deep learning. TF-Ranking can handle heterogeneous dense and sparse features, and scales up to millions of data points. For more details, please read the technical paper published on [arXiv](https://arxiv.org/abs/1812.00073).\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/ranking/blob/master/tensorflow_ranking/examples/handling_sparse_features.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/ranking/blob/master/tensorflow_ranking/examples/handling_sparse_features.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xmGZHWkdMEdB"
      },
      "source": [
        "## Motivation\n",
        "\n",
        "Learning to Rank (LTR) deals with learning to optimally order a list of examples, given some context. For instance, in search applications, examples are documents and context is the query.\n",
        "\n",
        "These models are usually trained using user relevance feedback, which can be explicit (human ratings) or implicit (clicks).\n",
        "\n",
        "This tutorial demonstrates how to build ranking estimators over sparse features, such as textual data. Textual data is prevalent in several settings for ranking, and plays a significant role is relevance judgment by a user.\n",
        "\n",
        "In three different LTR scenarios, the following textual features provide useful signals for ranking:\n",
        "\n",
        "*   Search: queries and document titles\n",
        "*   Question Answering: questions and answers\n",
        "*   Recommendation: titles of items and their descriptions\n",
        "\n",
        "Hence it is important for LTR models to effectively incorporate textual features."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A5A2UcQM1b7S"
      },
      "source": [
        "## Task: Ranking over Question-Answering data\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JNilCoqq1jJn"
      },
      "source": [
        "### ANTIQUE: A Question Answering Dataset\n",
        "\n",
        "For the purpose of this tutorial, we consider ranking problem over ANTIQUE, a question-answering dataset. Given a query, and a list of answers, the objective it to maximize a rank related metric (say NDCG).\n",
        "\n",
        "[ANTIQUE](http://hamedz.ir/resources/) is a publicly available dataset for open-domain non-factoid question answering, collected over Yahoo! answers.\n",
        "\n",
        "Each question has a list of answers, whose relevance are graded on a scale of 1-5.\n",
        "\n",
        "The list size can vary depending on the query, so we use a fixed \"list size\" of 50, where the list is either truncated or padded with dummy values.\n",
        "\n",
        "This dataset is a suitable one for learning-to-rank scenario. The dataset is split into 2206 queries for training and 200 queries for testing. For more details, please read the technical paper on [arXiv](https://arxiv.org/abs/1905.08957)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sIxFuiy5qh6L"
      },
      "source": [
        "\n",
        "Download training, test data and vocabulary file."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mwxtsi4wqoOJ"
      },
      "source": [
        "!wget -O \"/tmp/vocab.txt\" \"http://ciir.cs.umass.edu/downloads/Antique/tf-ranking/vocab.txt\"\n",
        "!wget -O \"/tmp/train.tfrecords\" \"http://ciir.cs.umass.edu/downloads/Antique/tf-ranking/ELWC/train.tfrecords\"\n",
        "!wget -O \"/tmp/test.tfrecords\" \"http://ciir.cs.umass.edu/downloads/Antique/tf-ranking//ELWC/test.tfrecords\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T7L-IOmOWm3s"
      },
      "source": [
        "Next, we discuss data formats in more detail, and show how to generate and store dummy ranking data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dXA_7oIRWKd9"
      },
      "source": [
        "## Data Formats\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6tuna6Td3_UO"
      },
      "source": [
        "### Data Formats for Ranking\n",
        "\n",
        "For representing ranking data, [protobuffers](https://developers.google.com/protocol-buffers/) are extensible structures suitable for storing data in a serialized format, either locally or in a distributed manner.\n",
        "\n",
        "Ranking usually consists of features corresponding to each of the examples being sorted. In addition, features related to query, user or session are also useful for ranking. We refer to these as context features, as these are independent of the examples.\n",
        "\n",
        "We use the popular [tf.Example](https://www.tensorflow.org/tutorials/load_data/tfrecord) proto to represent the features for context, and each of the examples. We use the protobuffer, **ExampleListWithContext** (ELWC), to store context as a tf.Example proto and the list of examples to be ranked as a list of tf.Example protos.\n",
        "\n",
        "ExampleListWithContext protbuffer is defined [here](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/apis/input.proto#L72)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95QnMu1cyPYA"
      },
      "source": [
        "Let us create some dummy data in ELWC format. We will use this dummy data to show how the proto looks like."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "etVmJYOngOxD"
      },
      "source": [
        "Download and install the TensorFlow Ranking and TensorFlow Serving packages."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aPmhLkMWgPLO"
      },
      "source": [
        "!pip install -q tensorflow_ranking tensorflow-serving-api"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7xriWS-gUx3I"
      },
      "source": [
        "Let us start by importing libraries that will be used throughout this Notebook. We also enable the \"eager execution\" mode for convenience and demonstration purposes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fmlaz2D5Ux3J"
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_ranking as tfr\n",
        "from tensorflow_serving.apis import input_pb2"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ooOmCPHbyd02"
      },
      "source": [
        "from google.protobuf import text_format\n",
        "\n",
        "CONTEXT = text_format.Parse(\n",
        "    \"\"\"\n",
        "    features {\n",
        "      feature {\n",
        "        key: \"query_tokens\"\n",
        "        value { bytes_list { value: [\"this\", \"is\", \"a\", \"relevant\", \"question\"] } }\n",
        "      }\n",
        "    }\"\"\", tf.train.Example())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eE7hpEBBykVS"
      },
      "source": [
        "EXAMPLES = [\n",
        "    text_format.Parse(\n",
        "    \"\"\"\n",
        "    features {\n",
        "      feature {\n",
        "        key: \"document_tokens\"\n",
        "        value { bytes_list { value: [\"this\", \"is\", \"a\", \"relevant\", \"answer\"] } }\n",
        "      }\n",
        "      feature {\n",
        "        key: \"relevance\"\n",
        "        value { int64_list { value: 5 } }\n",
        "      }\n",
        "    }\"\"\", tf.train.Example()),\n",
        "    text_format.Parse(\n",
        "        \"\"\"\n",
        "    features {\n",
        "      feature {\n",
        "        key: \"document_tokens\"\n",
        "        value { bytes_list { value: [\"irrelevant\", \"data\"] } }\n",
        "      }\n",
        "      feature {\n",
        "        key: \"relevance\"\n",
        "        value { int64_list { value: 1 } }\n",
        "      }\n",
        "    }\"\"\", tf.train.Example()),\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iImhGlJyyo--"
      },
      "source": [
        "ELWC = input_pb2.ExampleListWithContext()\n",
        "ELWC.context.CopyFrom(CONTEXT)\n",
        "for example in EXAMPLES:\n",
        "  example_features = ELWC.examples.add()\n",
        "  example_features.CopyFrom(example)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "Wml7IFUP0XoM",
        "outputId": "a25b6962-cd6a-4fb2-9a77-13d906936668"
      },
      "source": [
        "print(ELWC)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "data": {
            "text/plain": [
              "examples {\n",
              "  features {\n",
              "    feature {\n",
              "      key: \"document_tokens\"\n",
              "      value {\n",
              "        bytes_list {\n",
              "          value: \"this\"\n",
              "          value: \"is\"\n",
              "          value: \"a\"\n",
              "          value: \"relevant\"\n",
              "          value: \"answer\"\n",
              "        }\n",
              "      }\n",
              "    }\n",
              "    feature {\n",
              "      key: \"relevance\"\n",
              "      value {\n",
              "        int64_list {\n",
              "          value: 5\n",
              "        }\n",
              "      }\n",
              "    }\n",
              "  }\n",
              "}\n",
              "examples {\n",
              "  features {\n",
              "    feature {\n",
              "      key: \"document_tokens\"\n",
              "      value {\n",
              "        bytes_list {\n",
              "          value: \"irrelevant\"\n",
              "          value: \"data\"\n",
              "        }\n",
              "      }\n",
              "    }\n",
              "    feature {\n",
              "      key: \"relevance\"\n",
              "      value {\n",
              "        int64_list {\n",
              "          value: 1\n",
              "        }\n",
              "      }\n",
              "    }\n",
              "  }\n",
              "}\n",
              "context {\n",
              "  features {\n",
              "    feature {\n",
              "      key: \"query_tokens\"\n",
              "      value {\n",
              "        bytes_list {\n",
              "          value: \"this\"\n",
              "          value: \"is\"\n",
              "          value: \"a\"\n",
              "          value: \"relevant\"\n",
              "          value: \"question\"\n",
              "        }\n",
              "      }\n",
              "    }\n",
              "  }\n",
              "}"
            ]
          },
          "execution_count": 8,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d_PrxvnLTO1i"
      },
      "source": [
        "### Dependencies and Global Variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0B09XkzZYKLV"
      },
      "source": [
        "Here we define the train and test paths, along with model hyperparameters.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SDEhIqPSYJS1"
      },
      "source": [
        "# Store the paths to files containing training and test instances.\n",
        "_TRAIN_DATA_PATH = \"/tmp/train.tfrecords\"\n",
        "_TEST_DATA_PATH = \"/tmp/test.tfrecords\"\n",
        "\n",
        "# Store the vocabulary path for query and document tokens.\n",
        "_VOCAB_PATH = \"/tmp/vocab.txt\"\n",
        "\n",
        "# The maximum number of documents per query in the dataset.\n",
        "# Document lists are padded or truncated to this size.\n",
        "_LIST_SIZE = 50\n",
        "\n",
        "# The document relevance label.\n",
        "_LABEL_FEATURE = \"relevance\"\n",
        "\n",
        "# Padding labels are set negative so that the corresponding examples can be\n",
        "# ignored in loss and metrics.\n",
        "_PADDING_LABEL = -1\n",
        "\n",
        "# Learning rate for optimizer.\n",
        "_LEARNING_RATE = 0.05\n",
        "\n",
        "# Parameters to the scoring function.\n",
        "_BATCH_SIZE = 32\n",
        "_HIDDEN_LAYER_DIMS = [\"64\", \"32\", \"16\"]\n",
        "_DROPOUT_RATE = 0.8\n",
        "_GROUP_SIZE = 1  # Pointwise scoring.\n",
        "\n",
        "# Location of model directory and number of training steps.\n",
        "_MODEL_DIR = \"/tmp/ranking_model_dir\"\n",
        "_NUM_TRAIN_STEPS = 15 * 1000"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFbFBTUh9WXf"
      },
      "source": [
        "## Components of a Ranking Estimator\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pSIfzuaKKOa8"
      },
      "source": [
        "The overall components of a Ranking Estimator are shown below.\n",
        "\n",
        "The key components of the library are:\n",
        "\n",
        "1.   Input Reader\n",
        "2.   Tranform Function\n",
        "3.   Scoring Function\n",
        "4.   Ranking Losses\n",
        "5.   Ranking Metrics\n",
        "6.   Ranking Head\n",
        "7.   Model Builder\n",
        "\n",
        "These are described in more details in the following sections."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pohq1eKbo4TD"
      },
      "source": [
        "### TensorFlow Ranking Architecture\n",
        "\n",
        "![tf_ranking_arch](https://user-images.githubusercontent.com/3262617/60061785-5f107980-96ab-11e9-9849-ace2d117220f.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aQ-VTA56sOTA"
      },
      "source": [
        "### Specifying Features via Feature Columns\n",
        "\n",
        "[Feature Columns](https://www.tensorflow.org/tutorials/structured_data/feature_columns) are TensorFlow abstractions that are used to capture rich information about each feature. It allows for easy transformations for a diverse range of raw features and for interfacing with Estimators.\n",
        "\n",
        "Consistent with our input formats for ranking, such as ELWC format, we create feature columns for context features and example features."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I0TWSuMZ70eI"
      },
      "source": [
        "_EMBEDDING_DIMENSION = 20\n",
        "\n",
        "\n",
        "def context_feature_columns():\n",
        "  \"\"\"Returns context feature names to column definitions.\"\"\"\n",
        "  sparse_column = tf.feature_column.categorical_column_with_vocabulary_file(\n",
        "      key=\"query_tokens\",\n",
        "      vocabulary_file=_VOCAB_PATH)\n",
        "  query_embedding_column = tf.feature_column.embedding_column(\n",
        "      sparse_column, _EMBEDDING_DIMENSION)\n",
        "  return {\"query_tokens\": query_embedding_column}\n",
        "\n",
        "\n",
        "def example_feature_columns():\n",
        "  \"\"\"Returns the example feature columns.\"\"\"\n",
        "  sparse_column = tf.feature_column.categorical_column_with_vocabulary_file(\n",
        "      key=\"document_tokens\",\n",
        "      vocabulary_file=_VOCAB_PATH)\n",
        "  document_embedding_column = tf.feature_column.embedding_column(\n",
        "      sparse_column, _EMBEDDING_DIMENSION)\n",
        "  return {\"document_tokens\": document_embedding_column}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZGJ6rRJyZmiB"
      },
      "source": [
        "### Reading Input Data using *input_fn*\n",
        "\n",
        "The input reader reads in data from persistent storage to produce raw dense and sparse tensors of appropriate type for each feature. Example features are represented by 3-D tensors (where dimensions correspond to queries, examples and feature values). Context features are represented by 2-D tensors (where dimensions correspond to queries and feature values)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "itrqULz5lubN"
      },
      "source": [
        "def input_fn(path, num_epochs=None):\n",
        "  context_feature_spec = tf.feature_column.make_parse_example_spec(\n",
        "        context_feature_columns().values())\n",
        "  label_column = tf.feature_column.numeric_column(\n",
        "        _LABEL_FEATURE, dtype=tf.int64, default_value=_PADDING_LABEL)\n",
        "  example_feature_spec = tf.feature_column.make_parse_example_spec(\n",
        "        list(example_feature_columns().values()) + [label_column])\n",
        "  dataset = tfr.data.build_ranking_dataset(\n",
        "        file_pattern=path,\n",
        "        data_format=tfr.data.ELWC,\n",
        "        batch_size=_BATCH_SIZE,\n",
        "        list_size=_LIST_SIZE,\n",
        "        context_feature_spec=context_feature_spec,\n",
        "        example_feature_spec=example_feature_spec,\n",
        "        reader=tf.data.TFRecordDataset,\n",
        "        shuffle=False,\n",
        "        num_epochs=num_epochs)\n",
        "  features = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()\n",
        "  label = tf.squeeze(features.pop(_LABEL_FEATURE), axis=2)\n",
        "  label = tf.cast(label, tf.float32)\n",
        "\n",
        "  return features, label"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TXNUKT2s8bsQ"
      },
      "source": [
        "### Feature Transformations with *transform_fn*\n",
        "\n",
        "The transform function takes in the raw dense or sparse features from the input reader, applies suitable transformations to return dense representations for each feature. This is important before passing these features to a neural network, as neural networks layers usually take dense features as inputs.\n",
        "\n",
        "The transform function handles any custom feature transformations defined by the user. For handling sparse features, like text data, we provide an easy utlity to create shared embeddings, based on the feature columns."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zQG5UXub82SB"
      },
      "source": [
        "def make_transform_fn():\n",
        "  def _transform_fn(features, mode):\n",
        "    \"\"\"Defines transform_fn.\"\"\"\n",
        "    context_features, example_features = tfr.feature.encode_listwise_features(\n",
        "        features=features,\n",
        "        context_feature_columns=context_feature_columns(),\n",
        "        example_feature_columns=example_feature_columns(),\n",
        "        mode=mode,\n",
        "        scope=\"transform_layer\")\n",
        "\n",
        "    return context_features, example_features\n",
        "  return _transform_fn"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QdVFrZBIeWXL"
      },
      "source": [
        "### Feature Interactions using *scoring_fn*\n",
        "\n",
        "Next, we turn to the scoring function which is arguably at the heart of a TF Ranking model. The idea is to compute a relevance score for a (set of) query-document pair(s). The TF-Ranking model will use training data to learn this function.\n",
        "\n",
        "Here we formulate a scoring function using a feed forward network. The function takes the features of a single example (i.e., query-document pair) and produces a relevance score."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B7Ft1i2oieEY"
      },
      "source": [
        "def make_score_fn():\n",
        "  \"\"\"Returns a scoring function to build `EstimatorSpec`.\"\"\"\n",
        "\n",
        "  def _score_fn(context_features, group_features, mode, params, config):\n",
        "    \"\"\"Defines the network to score a group of documents.\"\"\"\n",
        "    with tf.compat.v1.name_scope(\"input_layer\"):\n",
        "      context_input = [\n",
        "          tf.compat.v1.layers.flatten(context_features[name])\n",
        "          for name in sorted(context_feature_columns())\n",
        "      ]\n",
        "      group_input = [\n",
        "          tf.compat.v1.layers.flatten(group_features[name])\n",
        "          for name in sorted(example_feature_columns())\n",
        "      ]\n",
        "      input_layer = tf.concat(context_input + group_input, 1)\n",
        "\n",
        "    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
        "    cur_layer = input_layer\n",
        "    cur_layer = tf.compat.v1.layers.batch_normalization(\n",
        "      cur_layer,\n",
        "      training=is_training,\n",
        "      momentum=0.99)\n",
        "\n",
        "    for i, layer_width in enumerate(int(d) for d in _HIDDEN_LAYER_DIMS):\n",
        "      cur_layer = tf.compat.v1.layers.dense(cur_layer, units=layer_width)\n",
        "      cur_layer = tf.compat.v1.layers.batch_normalization(\n",
        "        cur_layer,\n",
        "        training=is_training,\n",
        "        momentum=0.99)\n",
        "      cur_layer = tf.nn.relu(cur_layer)\n",
        "      cur_layer = tf.compat.v1.layers.dropout(\n",
        "          inputs=cur_layer, rate=_DROPOUT_RATE, training=is_training)\n",
        "    logits = tf.compat.v1.layers.dense(cur_layer, units=_GROUP_SIZE)\n",
        "    return logits\n",
        "\n",
        "  return _score_fn"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xQlrS6uB8_zl"
      },
      "source": [
        "## Losses, Metrics and Ranking Head"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0gt2bu7kbtS"
      },
      "source": [
        "### Evaluation Metrics\n",
        "\n",
        "We have provided an implementation of several popular Information Retrieval evaluation metrics in the TF Ranking library, which are shown [here](https://github.com/tensorflow/ranking/blob/d8c2e2e64a92923f1448cf5302c92a80bb469a20/tensorflow_ranking/python/metrics.py#L32). The user can also define a custom evaluation metric, as shown in the description below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YkU6o2QjkyXR"
      },
      "source": [
        "def eval_metric_fns():\n",
        "  \"\"\"Returns a dict from name to metric functions.\n",
        "\n",
        "  This can be customized as follows. Care must be taken when handling padded\n",
        "  lists.\n",
        "\n",
        "  def _auc(labels, predictions, features):\n",
        "    is_label_valid = tf_reshape(tf.greater_equal(labels, 0.), [-1, 1])\n",
        "    clean_labels = tf.boolean_mask(tf.reshape(labels, [-1, 1], is_label_valid)\n",
        "    clean_pred = tf.boolean_maks(tf.reshape(predictions, [-1, 1], is_label_valid)\n",
        "    return tf.metrics.auc(clean_labels, tf.sigmoid(clean_pred), ...)\n",
        "  metric_fns[\"auc\"] = _auc\n",
        "\n",
        "  Returns:\n",
        "    A dict mapping from metric name to a metric function with above signature.\n",
        "  \"\"\"\n",
        "  metric_fns = {}\n",
        "  metric_fns.update({\n",
        "      f\"metric/ndcg@{topn}\": tfr.metrics.make_ranking_metric_fn(\n",
        "          tfr.metrics.RankingMetricKey.NDCG, topn=topn)\n",
        "      for topn in [1, 3, 5, 10]\n",
        "  })\n",
        "\n",
        "  return metric_fns"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HfDMGnZY9eVO"
      },
      "source": [
        "### Ranking Losses\n",
        "\n",
        "We provide several popular ranking loss functions as part of the library, which are shown [here](https://github.com/tensorflow/ranking/blob/d8c2e2e64a92923f1448cf5302c92a80bb469a20/tensorflow_ranking/python/losses.py#L35). The user can also define a custom loss function, similar to ones in tfr.losses."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kJSaBkxW9jZy"
      },
      "source": [
        "# Define a loss function. To find a complete list of available\n",
        "# loss functions or to learn how to add your own custom function\n",
        "# please refer to the tensorflow_ranking.losses module.\n",
        "\n",
        "_LOSS = tfr.losses.RankingLossKey.APPROX_NDCG_LOSS\n",
        "loss_fn = tfr.losses.make_loss_fn(_LOSS)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OUoFtZGe-pGD"
      },
      "source": [
        "### Ranking Head\n",
        "\n",
        "In the Estimator workflow, Head is an abstraction that encapsulates losses and corresponding metrics. Head easily interfaces with the Estimator, needing the user to define a scoring function and specify losses and metric computation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3pXnzLh_-rBb"
      },
      "source": [
        "optimizer = tf.compat.v1.train.AdagradOptimizer(\n",
        "    learning_rate=_LEARNING_RATE)\n",
        "\n",
        "\n",
        "def _train_op_fn(loss):\n",
        "  \"\"\"Defines train op used in ranking head.\"\"\"\n",
        "  update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)\n",
        "  minimize_op = optimizer.minimize(\n",
        "      loss=loss, global_step=tf.compat.v1.train.get_global_step())\n",
        "  train_op = tf.group([update_ops, minimize_op])\n",
        "  return train_op\n",
        "\n",
        "\n",
        "ranking_head = tfr.head.create_ranking_head(\n",
        "      loss_fn=loss_fn,\n",
        "      eval_metric_fns=eval_metric_fns(),\n",
        "      train_op_fn=_train_op_fn)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SGJwvyrXk-Yj"
      },
      "source": [
        "## Putting It All Together in a Model Builder\n",
        "\n",
        "We are now ready to put all of the components above together and create an `Estimator` that can be used to train and evaluate a model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XoR9hRWHlCR4"
      },
      "source": [
        "model_fn = tfr.model.make_groupwise_ranking_fn(\n",
        "          group_score_fn=make_score_fn(),\n",
        "          transform_fn=make_transform_fn(),\n",
        "          group_size=_GROUP_SIZE,\n",
        "          ranking_head=ranking_head)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7aky8RFpH-D0"
      },
      "source": [
        "### Train and evaluate the ranker"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7tD0aNuYU8LR"
      },
      "source": [
        "def train_and_eval_fn():\n",
        "  \"\"\"Train and eval function used by `tf.estimator.train_and_evaluate`.\"\"\"\n",
        "  run_config = tf.estimator.RunConfig(\n",
        "      save_checkpoints_steps=1000)\n",
        "  ranker = tf.estimator.Estimator(\n",
        "      model_fn=model_fn,\n",
        "      model_dir=_MODEL_DIR,\n",
        "      config=run_config)\n",
        "\n",
        "  train_input_fn = lambda: input_fn(_TRAIN_DATA_PATH)\n",
        "  eval_input_fn = lambda: input_fn(_TEST_DATA_PATH, num_epochs=1)\n",
        "\n",
        "  train_spec = tf.estimator.TrainSpec(\n",
        "      input_fn=train_input_fn, max_steps=_NUM_TRAIN_STEPS)\n",
        "  eval_spec = tf.estimator.EvalSpec(\n",
        "      name=\"eval\",\n",
        "      input_fn=eval_input_fn,\n",
        "      throttle_secs=15)\n",
        "  return (ranker, train_spec, eval_spec)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uXoA7xwuVmCD"
      },
      "source": [
        "! rm -rf \"/tmp/ranking_model_dir\"  # Clean up the model directory.\n",
        "ranker, train_spec, eval_spec = train_and_eval_fn()\n",
        "tf.estimator.train_and_evaluate(ranker, train_spec, eval_spec)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45WYaJNaGfLM"
      },
      "source": [
        "### Launch TensorBoard"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sHfuUVQ5D1jq"
      },
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir=\"/tmp/ranking_model_dir\" --port 12345"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S1s1BKWSP8p_"
      },
      "source": [
        "A sample tensorboard output is shown here, with the ranking metrics.\n",
        "\n",
        "![tensorboard](https://user-images.githubusercontent.com/3262617/60866646-be0edc00-a1dd-11e9-9599-eefb734ce801.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b1QymoFfDNr7"
      },
      "source": [
        "## Generating Predictions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QqHmehSKDOg1"
      },
      "source": [
        "We show how to generate predictions over the features of a dataset. We assume that the label is not present and needs to be inferred using the ranking model.\n",
        "\n",
        "Similar to the `input_fn` used for training and evaluation,  `predict_input_fn` reads in data in ELWC format and stored as TFRecords to generate features. We set number of epochs to be 1, so that the generator stops iterating when it reaches the end of the dataset. Also the datapoints are not shuffled while reading, so that the behavior of the `predict()` function is deterministic."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vF-4BiTnDRou"
      },
      "source": [
        "def predict_input_fn(path):\n",
        "  context_feature_spec = tf.feature_column.make_parse_example_spec(\n",
        "        context_feature_columns().values())\n",
        "  example_feature_spec = tf.feature_column.make_parse_example_spec(\n",
        "        list(example_feature_columns().values()))\n",
        "  dataset = tfr.data.build_ranking_dataset(\n",
        "        file_pattern=path,\n",
        "        data_format=tfr.data.ELWC,\n",
        "        batch_size=_BATCH_SIZE,\n",
        "        list_size=_LIST_SIZE,\n",
        "        context_feature_spec=context_feature_spec,\n",
        "        example_feature_spec=example_feature_spec,\n",
        "        reader=tf.data.TFRecordDataset,\n",
        "        shuffle=False,\n",
        "        num_epochs=1)\n",
        "  features = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()\n",
        "  return features"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lVpLTPp-DVTi"
      },
      "source": [
        "We generate predictions on the test dataset, where we only consider context and example features and predict the labels. The `predict_input_fn` generates predictions on a batch of datapoints. Batching allows us to iterate over large datasets which cannot be loaded in memory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KVdlImkADTvK"
      },
      "source": [
        "predictions = ranker.predict(input_fn=lambda: predict_input_fn(\"/tmp/test.tfrecords\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXh8knanDdd1"
      },
      "source": [
        "`ranker.predict` returns a generator, which we can iterate over to create predictions, till the generator is exhausted."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cROuS7-8Dd-r"
      },
      "source": [
        "x = next(predictions)\n",
        "assert len(x) == _LIST_SIZE  # Note that this includes padding."
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}