TensorFlow Ranking v0.4.2
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 tensorsRankingTask
, 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
andRankingTaskConfig
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 theRankingDataLoader
and further specifies the feature specs of a TFR-BERT model.TFRBertScorer
andTFRBertModelBuilder
, 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 ofRankingTask
. It defines thebuild_model
behavior. It also defines theinitialization
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
andTFRBertConfig
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
.