#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2022 The TensorFlow Authors.
End of explanation
"""


!pip install -q tensorflow-ranking tf-models-official

"""
Explanation: Passage Ranking using TFR-BERT
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/ranking/tutorials/tfr_bert"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/ranking/blob/master/docs/tutorials/tfr_bert.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/ranking/blob/master/docs/tutorials/tfr_bert.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/ranking/docs/tutorials/tfr_bert.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>

TensorFlow Ranking can handle heterogeneous dense and sparse features, and scales up to millions of data points. However, building and deploying a learning to rank model to operate at scale creates additional challenges beyond simply designing a model. The Ranking library provides workflow utility classes for building distributed training for large-scale ranking applications. For more information about these features, see the TensorFlow Ranking Overview.
This tutorial shows you how to build a ranking model that uses BERT for scoring. BERT is a highly effective pretrained module to effective encode textual features into contextualized word embeddings. We use BERT to initialize the ranking model and finetune the model using a ranking loss.
Note: An advanced version of this code is also available as a Python script.
ANTIQUE dataset
In this tutorial, you will build a ranking model for ANTIQUE, a question-answering dataset using BERT as the scoring function. Bidirectional Encoder Representations from Transformers (BERT) is a transformer-based machine learning technique which has proven to be effective in many natural language processing (NLP) tasks. Recent work on TFR-BERT has shown BERT to be an effective scoring function for learning-to-rank tasks.
Given a query, and a list of answers, the objective of the ranking model is to rank the answers with optimal rank related metrics, such as NDCG. For more details about ranking metrics, review evaluation measures offline metrics.
ANTIQUE is a publicly available dataset for open-domain non-factoid question answering, collected from Yahoo! answers.
Each question has a list of answers, whose relevance are graded on a scale of 0-4, 0 for irrelevant and 4 for fully relevant.
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 default values.
The dataset is split into 2206 queries for training and 200 queries for testing. For more details, please read the technical paper on arXiv.
Setup
Download and install the TensorFlow Ranking and TensorFlow Model Garden packages.
End of explanation
"""


import os
import tensorflow as tf
import tensorflow_ranking as tfr
from official.nlp.configs import encoders
from tensorflow_ranking.extension.premade import tfrbert_task

"""
Explanation: Import TensorFlow Ranking and useful libraries through the notebook.
End of explanation
"""


!wget -O "/tmp/train.tfrecords" "https://ciir.cs.umass.edu/downloads/Antique/tf-ranking/antique_train_seq_64_elwc.tfrecords"
!wget -O "/tmp/test.tfrecords" "https://ciir.cs.umass.edu/downloads/Antique/tf-ranking/antique_test_seq_64_elwc.tfrecords"

!mkdir -p /tmp/tfrbert
!wget "https://storage.googleapis.com/cloud-tpu-checkpoints/bert/v3/uncased_L-12_H-768_A-12.tar.gz" -P "/tmp/tfrbert"
!mkdir -p /tmp/tfrbert/uncased_L-12_H-768_A-12
!tar -xvf /tmp/tfrbert/uncased_L-12_H-768_A-12.tar.gz --strip-components 3 -C "/tmp/tfrbert/uncased_L-12_H-768_A-12/"

"""
Explanation: Data preparation
Download training and test data.
End of explanation
"""


SEQ_LENGTH = 64
context_feature_spec = {}
example_feature_spec = {
    'input_word_ids': tf.io.FixedLenFeature(
        shape=(SEQ_LENGTH,), dtype=tf.int64,
        default_value=[0] * SEQ_LENGTH),
    'input_mask': tf.io.FixedLenFeature(
        shape=(SEQ_LENGTH,), dtype=tf.int64,
        default_value=[0] * SEQ_LENGTH),
    'input_type_ids': tf.io.FixedLenFeature(
        shape=(SEQ_LENGTH,), dtype=tf.int64,
        default_value=[0] * SEQ_LENGTH)}
label_spec = (
    "relevance",
    tf.io.FixedLenFeature(shape=(1,), dtype=tf.int64, default_value=-1)
)

"""
Explanation: Overview of TFR-BERT in Orbit
BERT-based ranking models (TFR-BERT) have been shown to be effective for learning-to-rank tasks when using raw textual features for query and passages in MSMARCO passage ranking dataset.
Orbit is a flexible, lightweight library designed to make it easy to write custom training loops in TensorFlow. TensorFlow Ranking provides support for implementing ranking models, particularly for BERT based ranking models using Orbit.
Create a Ranking Task for TFR-BERT
We create a ranking task for TFR-BERT model which can be trained using Orbit. The steps to build this are:

Define Feature Specifications
Define datasets
Setup data and task configurations

Specify Features
Feature Specification are TensorFlow abstractions to capture information about each feature. These help developers and model researchers understand and use a model.
Create feature specifications for context features, example features, and labels, consistent with the input formats for ranking, such as ELWC format.
End of explanation
"""


# Set up data config
# We use a small list size here for demo purposes only. Users can use a larger
# list size on a machine with more memory to train TFR-BERT.
train_data_config = tfrbert_task.TFRBertDataConfig(
    input_path="/tmp/train.tfrecords",
    is_training=True,
    global_batch_size=8,
    list_size=2,
    dataset_fn='tfrecord',
    seq_length=64)

validation_data_config = tfrbert_task.TFRBertDataConfig(
    input_path="/tmp/test.tfrecords",
    is_training=False,
    global_batch_size=8,
    list_size=2,
    dataset_fn='tfrecord',
    seq_length=64)

"""
Explanation: Note: the default_value of label_spec feature is set to -1 to take care of the padding items to be masked out.
Define Datasets
We define data configurations for training and validation data, which specifies parameters such as path, batch size, and dataset format. These configurations are used to create training and validation datasets.
End of explanation
"""


# Set up task config
task_config = tfrbert_task.TFRBertConfig(
    init_checkpoint='/tmp/tfrbert/uncased_L-12_H-768_A-12/bert_model.ckpt',
    train_data=train_data_config,
    validation_data=validation_data_config,
    model=tfrbert_task.TFRBertModelConfig(
        encoder=encoders.EncoderConfig(
            bert=encoders.BertEncoderConfig(num_layers=12))))

# Set up TFRBertTask
task = tfrbert_task.TFRBertTask(
    task_config,
    label_spec=label_spec,
    dataset_fn=tf.data.TFRecordDataset,
    logging_dir='/tmp/model_dir')

"""
Explanation: Define Task
We define a task configuration which defines the training and validation dataset along with the model. This configuration creates a TFRBertTask object that can be trained using Orbit.
End of explanation
"""


metrics = task.build_metrics()
model = task.build_model()
task.initialize(model)
train_dataset = task.build_inputs(task_config.train_data)
vali_dataset = task.build_inputs(task_config.validation_data)
train_iterator = iter(train_dataset)
vali_iterator = iter(vali_dataset)
optimizer = tf.keras.optimizers.Adam(lr=1e-6)

NUM_TRAIN_STEPS = 100
EVAL_STEPS = 10
for train_step in range(NUM_TRAIN_STEPS):
  task.train_step(next(train_iterator), model, optimizer, metrics=metrics)
  train_metrics = {m.name: m.result().numpy() for m in metrics}
  print("Training metrics for epoch: " + str(train_step) + " ", train_metrics)

  if train_step % EVAL_STEPS == 0:
    task.validation_step(next(train_iterator), model, metrics=metrics)
    vali_metrics = {m.name: m.result().numpy() for m in metrics}
    print("Validation metrics for epoch: " + str(train_step) + " ",
          vali_metrics)

"""
Explanation: Train and evaluate the model
We define the training loop here to train and evaluate the model. We define the metrics, create train and eval datasets and train the model for a specific number of training steps.
End of explanation
"""


