Skip to content

TensorFlow Ranking v0.4.2

Compare
Choose a tag to compare
@HongleiZhuang HongleiZhuang released this 22 Jul 23:55
· 150 commits to master since this release

This is the 0.4.2 release of TensorFlow Ranking. The main changes are the TFR-BERT module based on the Orbit framework in tf-models, which facilitates users to write customized training loops. The new components are:

TFR-BERT in Orbit

  • tfr.keras.task: This module contains the general boilerplate code to train TF-Ranking models in the Orbit framework. Particularly, there are:
    • RankingDataLoader, which parses an ELWC formatted data record into tensors
    • RankingTask, which specifies the behaviors of each training and evaluation step, as well as the training losses and evaluation metrics.
    • In addition, there are config data classes like RankingDataConfig and RankingTaskConfig to store configurations for above classes.
  • tfr.keras.premade.tfrbert_task: This module contains the TFR-BERT specification of the TF-Ranking Orbit task.
    • TFRBertDataLoader, which subclasses the RankingDataLoader and further specifies the feature specs of a TFR-BERT model.
    • TFRBertScorer and TFRBertModelBuilder, which defines a model builder that can create a TFR-BERT ranking model as a Keras model, based on tf-models’ implementation of BERT encoder.
    • TFRBertTask, which is a subclass of RankingTask. It defines the build_model behavior. It also defines the initialization method which would load an pretrained BERT checkpoint to initialize the encoder. It also provides the function to output the prediction results along with query ids and document ids.
    • In addition, there are config data classes like TFRBertDataConfig, TFRBertModelConfig and TFRBertConfig which stores configurations for above classes.
  • examples/keras/tfrbert_antique_train.py: This file provides an example of training a TFR-BERT model on the Antique data set. There is also an .yaml file where users can specify parameter configurations.

Dependencies: The following packages will be installed as required when installing tensorflow-ranking.

  • tf-models-official >= 2.5.0
  • tensorflow-serving-api>= 2.0.0, < 3.0.0
  • tensorflow==2.5.0.