From 77f280075e220946d0935b44b64770f4dabb37d0 Mon Sep 17 00:00:00 2001
From: Rama Pasumarthi <ramakumar@google.com>
Date: Mon, 24 May 2021 19:02:02 -0400
Subject: [PATCH] Update version for release 0.4.0

PiperOrigin-RevId: 375577552
---
 .../g3doc/api_docs/python/tfr.md              |   2 +-
 .../g3doc/api_docs/python/tfr/_api_cache.json |  26 +-
 .../g3doc/api_docs/python/tfr/_toc.yaml       |   2 +
 .../g3doc/api_docs/python/tfr/all_symbols.md  |   1 +
 .../g3doc/api_docs/python/tfr/keras.md        |   2 +-
 .../python/tfr/keras/layers/FlattenList.md    |   1 -
 .../api_docs/python/tfr/keras/metrics.md      |   7 +
 .../python/tfr/keras/metrics/ARPMetric.md     |   4 +-
 .../tfr/keras/metrics/AlphaDCGMetric.md       |   6 +-
 .../python/tfr/keras/metrics/DCGMetric.md     |   6 +-
 .../python/tfr/keras/metrics/MRRMetric.md     |   6 +-
 .../metrics/MeanAveragePrecisionMetric.md     |   6 +-
 .../python/tfr/keras/metrics/NDCGMetric.md    |   6 +-
 .../python/tfr/keras/metrics/OPAMetric.md     |   4 +-
 .../tfr/keras/metrics/PrecisionIAMetric.md    |   7 +-
 .../tfr/keras/metrics/PrecisionMetric.md      |   6 +-
 .../tfr/keras/metrics/RankingMetricKey.md     |   2 +-
 .../python/tfr/keras/metrics/RecallMetric.md  |   6 +-
 .../keras/metrics/default_keras_metrics.md    |   3 +-
 .../api_docs/python/tfr/keras/metrics/get.md  |   2 +-
 .../g3doc/api_docs/python/tfr/keras/model.md  |   9 +-
 .../tfr/keras/model/AbstractModelBuilder.md   | 353 +---------------
 .../python/tfr/keras/model/DNNScorer.md       |   5 +-
 .../keras/model/FeatureSpecInputCreator.md    |   5 +-
 .../python/tfr/keras/model/GAMScorer.md       |   5 +-
 .../python/tfr/keras/model/InputCreator.md    |   5 +-
 .../python/tfr/keras/model/ModelBuilder.md    |  21 +-
 .../tfr/keras/model/ModelBuilderWithMask.md   | 378 ++++++++++++++++++
 .../python/tfr/keras/model/Preprocessor.md    |   6 +-
 .../tfr/keras/model/PreprocessorWithSpec.md   |   4 +-
 .../api_docs/python/tfr/keras/model/Scorer.md |   4 +-
 .../tfr/keras/model/TypeSpecInputCreator.md   |   5 +-
 .../tfr/keras/model/UnivariateScorer.md       |   4 +-
 .../keras/pipeline/AbstractDatasetBuilder.md  |   6 -
 .../tfr/keras/pipeline/AbstractPipeline.md    |   5 -
 .../tfr/keras/pipeline/BaseDatasetBuilder.md  |   8 +-
 .../tfr/keras/pipeline/DatasetHparams.md      |   3 +-
 .../tfr/keras/pipeline/ModelFitPipeline.md    |  16 +-
 .../pipeline/MultiLabelDatasetBuilder.md      |   8 +-
 .../tfr/keras/pipeline/MultiTaskPipeline.md   |  23 +-
 .../tfr/keras/pipeline/NullDatasetBuilder.md  |   9 +-
 .../tfr/keras/pipeline/PipelineHparams.md     |  18 +-
 .../keras/pipeline/SimpleDatasetBuilder.md    |   8 +-
 .../tfr/keras/pipeline/SimplePipeline.md      |  23 +-
 .../api_docs/python/tfr/metrics/recall.md     |   3 +-
 tensorflow_ranking/python/version.py          |   4 +-
 tensorflow_ranking/tools/pip_package/setup.py |   4 +-
 47 files changed, 559 insertions(+), 488 deletions(-)
 create mode 100644 tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilderWithMask.md

diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr.md b/tensorflow_ranking/g3doc/api_docs/python/tfr.md
index aeec5df..aa9e9b7 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr.md
@@ -56,7 +56,7 @@ library.
 __version__<a id="__version__"></a>
 </td>
 <td>
-`'0.3.3.dev'`
+`'0.4.0.dev'`
 </td>
 </tr>
 </table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json b/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json
index 6ec5e95..0b92023 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json
@@ -981,6 +981,7 @@
     "tfr.keras.model.AbstractModelBuilder.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.model.AbstractModelBuilder.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.keras.model.AbstractModelBuilder.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.model.AbstractModelBuilder.__init__": "tfr.keras.losses.RankingLossKey.__init__",
     "tfr.keras.model.AbstractModelBuilder.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
     "tfr.keras.model.AbstractModelBuilder.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
     "tfr.keras.model.AbstractModelBuilder.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
@@ -1023,7 +1024,14 @@
     "tfr.keras.model.ModelBuilder.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
     "tfr.keras.model.ModelBuilder.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.keras.model.ModelBuilder.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
-    "tfr.keras.model.ModelBuilder.build": "tfr.keras.model.AbstractModelBuilder.build",
+    "tfr.keras.model.ModelBuilder.build": "tfr.keras.model.ModelBuilderWithMask.build",
+    "tfr.keras.model.ModelBuilderWithMask.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.model.ModelBuilderWithMask.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.model.ModelBuilderWithMask.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.model.ModelBuilderWithMask.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.model.ModelBuilderWithMask.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.model.ModelBuilderWithMask.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.model.ModelBuilderWithMask.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.model.Preprocessor.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.model.Preprocessor.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.keras.model.Preprocessor.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -2577,9 +2585,6 @@
     "tfr.keras.model.AbstractModelBuilder.__ne__": true,
     "tfr.keras.model.AbstractModelBuilder.__new__": true,
     "tfr.keras.model.AbstractModelBuilder.build": true,
-    "tfr.keras.model.AbstractModelBuilder.create_inputs": true,
-    "tfr.keras.model.AbstractModelBuilder.preprocess": true,
-    "tfr.keras.model.AbstractModelBuilder.score": true,
     "tfr.keras.model.DNNScorer": false,
     "tfr.keras.model.DNNScorer.__call__": true,
     "tfr.keras.model.DNNScorer.__eq__": true,
@@ -2633,6 +2638,19 @@
     "tfr.keras.model.ModelBuilder.create_inputs": true,
     "tfr.keras.model.ModelBuilder.preprocess": true,
     "tfr.keras.model.ModelBuilder.score": true,
+    "tfr.keras.model.ModelBuilderWithMask": false,
+    "tfr.keras.model.ModelBuilderWithMask.__eq__": true,
+    "tfr.keras.model.ModelBuilderWithMask.__ge__": true,
+    "tfr.keras.model.ModelBuilderWithMask.__gt__": true,
+    "tfr.keras.model.ModelBuilderWithMask.__init__": true,
+    "tfr.keras.model.ModelBuilderWithMask.__le__": true,
+    "tfr.keras.model.ModelBuilderWithMask.__lt__": true,
+    "tfr.keras.model.ModelBuilderWithMask.__ne__": true,
+    "tfr.keras.model.ModelBuilderWithMask.__new__": true,
+    "tfr.keras.model.ModelBuilderWithMask.build": true,
+    "tfr.keras.model.ModelBuilderWithMask.create_inputs": true,
+    "tfr.keras.model.ModelBuilderWithMask.preprocess": true,
+    "tfr.keras.model.ModelBuilderWithMask.score": true,
     "tfr.keras.model.Preprocessor": false,
     "tfr.keras.model.Preprocessor.__call__": true,
     "tfr.keras.model.Preprocessor.__eq__": true,
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/_toc.yaml b/tensorflow_ranking/g3doc/api_docs/python/tfr/_toc.yaml
index 9101b77..21c5adf 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/_toc.yaml
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/_toc.yaml
@@ -223,6 +223,8 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/model/InputCreator
     - title: ModelBuilder
       path: /ranking/api_docs/python/tfr/keras/model/ModelBuilder
+    - title: ModelBuilderWithMask
+      path: /ranking/api_docs/python/tfr/keras/model/ModelBuilderWithMask
     - title: Preprocessor
       path: /ranking/api_docs/python/tfr/keras/model/Preprocessor
     - title: PreprocessorWithSpec
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/all_symbols.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/all_symbols.md
index aab4a34..d852147 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/all_symbols.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/all_symbols.md
@@ -100,6 +100,7 @@
 *   <a href="../tfr/keras/model/GAMScorer.md"><code>tfr.keras.model.GAMScorer</code></a>
 *   <a href="../tfr/keras/model/InputCreator.md"><code>tfr.keras.model.InputCreator</code></a>
 *   <a href="../tfr/keras/model/ModelBuilder.md"><code>tfr.keras.model.ModelBuilder</code></a>
+*   <a href="../tfr/keras/model/ModelBuilderWithMask.md"><code>tfr.keras.model.ModelBuilderWithMask</code></a>
 *   <a href="../tfr/keras/model/Preprocessor.md"><code>tfr.keras.model.Preprocessor</code></a>
 *   <a href="../tfr/keras/model/PreprocessorWithSpec.md"><code>tfr.keras.model.PreprocessorWithSpec</code></a>
 *   <a href="../tfr/keras/model/Scorer.md"><code>tfr.keras.model.Scorer</code></a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md
index 1b16000..ac4ae06 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md
@@ -37,7 +37,7 @@ in Keras.
 
 [`metrics`](../tfr/keras/metrics.md) module: Keras metrics in TF-Ranking.
 
-[`model`](../tfr/keras/model.md) module: Ranking Model utilities and classes in
+[`model`](../tfr/keras/model.md) module: Ranking model utilities and classes in
 tfr.keras.
 
 [`network`](../tfr/keras/network.md) module: Ranking Networks in Keras.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/FlattenList.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/FlattenList.md
index f70f4ca..ba27b55 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/FlattenList.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/FlattenList.md
@@ -67,7 +67,6 @@ flattened_context_features = {'context_feature_1':
 flattened_example_features = {'example_feature_1':
                               [[1], [0], [1], [0], [0], [0]]}
 ```
-
 `context_feature_1` is repeated by list_size=3 times. `example_feature_1` is
 flattened and padded with the invalid terms replaced by valid terms in each
 query in a circular way.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md
index 144238f..8fc3aeb 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md
@@ -20,6 +20,13 @@ description: Keras metrics in TF-Ranking.
 
 Keras metrics in TF-Ranking.
 
+NOTE: For metrics that compute a ranking, ties are broken randomly. This means
+that metrics may be stochastic if items with equal scores are provided.
+
+WARNING: Some metrics (e.g. Recall or MRR) are not well-defined when there are
+no relevant items (e.g. if `y_true` has a row of only zeroes). For these cases,
+the TF-Ranking metrics will evaluate to `0`.
+
 ## Classes
 
 [`class ARPMetric`](../../tfr/keras/metrics/ARPMetric.md): Average relevance
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md
index a2206ec..540c88f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md
@@ -28,7 +28,7 @@ description: Average relevance position (ARP).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L235-L278">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L243-L286">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -783,7 +783,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
index 7009174..4c30256 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
@@ -28,7 +28,7 @@ description: Alpha discounted cumulative gain (alphaDCG).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L780-L902">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L788-L910">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -731,7 +731,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L894-L902">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L902-L910">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -889,7 +889,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md
index 4489822..02cc5e4 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md
@@ -28,7 +28,7 @@ description: Discounted cumulative gain (DCG).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L695-L776">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L703-L784">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -644,7 +644,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L768-L776">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L776-L784">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -802,7 +802,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md
index e0f45a6..2c7ae57 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md
@@ -28,7 +28,7 @@ description: Mean reciprocal rank (MRR).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L171-L231">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L179-L239">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -637,7 +637,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L226-L231">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L234-L239">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -795,7 +795,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
index f1aadcd..6bd99e5 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
@@ -28,7 +28,7 @@ description: Mean average precision (MAP).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L525-L602">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L533-L610">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -645,7 +645,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L596-L602">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L604-L610">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -803,7 +803,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md
index 3dae1e6..2784cea 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md
@@ -28,7 +28,7 @@ description: Normalized discounted cumulative gain (NDCG).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L606-L691">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L614-L699">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -649,7 +649,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L683-L691">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L691-L699">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -807,7 +807,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md
index 972f0c7..7e31c00 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md
@@ -28,7 +28,7 @@ description: Ordered pair accuracy (OPA).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L906-L957">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L914-L965">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -791,7 +791,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
index 2ca805a..5d53633 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
@@ -28,7 +28,7 @@ description: Precision-IA@k (Pre-IA@k).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L431-L521">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L439-L529">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -108,7 +108,6 @@ https://www.microsoft.com/en-us/research/publication/diversifying-search-results
 [clarke2009]: https://trec.nist.gov/pubs/trec18/papers/ENT09.OVERVIEW.pdf
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -695,7 +694,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L516-L521">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L524-L529">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -853,7 +852,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
index 4b8a563..e488eef 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
@@ -28,7 +28,7 @@ description: Precision@k (P@k).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L282-L352">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L290-L360">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -640,7 +640,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L347-L352">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L355-L360">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -798,7 +798,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RankingMetricKey.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RankingMetricKey.md
index 63b7da5..5acfad5 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RankingMetricKey.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RankingMetricKey.md
@@ -20,7 +20,7 @@ description: Ranking metric key strings.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L15-L42">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L23-L50">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RecallMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RecallMetric.md
index 1391fe3..a74d8e0 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RecallMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RecallMetric.md
@@ -28,7 +28,7 @@ description: Recall@k (R@k).
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L357-L427">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L365-L435">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -640,7 +640,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L422-L427">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L430-L435">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -798,7 +798,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L145-L167">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L153-L175">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md
index 7293439..a8a842c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md
@@ -11,7 +11,7 @@ description: Returns a list of ranking metrics.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L106-L128">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L114-L136">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -29,7 +29,6 @@ Returns a list of ranking metrics.
 <!-- Placeholder for "Used in" -->
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/get.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/get.md
index fb28e31..671b666 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/get.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/get.md
@@ -11,7 +11,7 @@ description: Factory method to get a list of ranking metrics.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L45-L103">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L53-L111">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md
index fddae8e..12df7fe 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md
@@ -1,4 +1,4 @@
-description: Ranking Model utilities and classes in tfr.keras.
+description: Ranking model utilities and classes in tfr.keras.
 
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.model" />
@@ -18,12 +18,12 @@ description: Ranking Model utilities and classes in tfr.keras.
 </td>
 </table>
 
-Ranking Model utilities and classes in tfr.keras.
+Ranking model utilities and classes in tfr.keras.
 
 ## Classes
 
 [`class AbstractModelBuilder`](../../tfr/keras/model/AbstractModelBuilder.md):
-Interface to build a ranking `tf.keras.Model`.
+Interface to build a `tf.keras.Model` for ranking.
 
 [`class DNNScorer`](../../tfr/keras/model/DNNScorer.md): Univariate scorer using
 DNN.
@@ -40,6 +40,9 @@ input creator.
 [`class ModelBuilder`](../../tfr/keras/model/ModelBuilder.md): Builds a
 `tf.keras.Model`.
 
+[`class ModelBuilderWithMask`](../../tfr/keras/model/ModelBuilderWithMask.md):
+Interface to build a `tf.keras.Model` for ranking with a mask Tensor.
+
 [`class Preprocessor`](../../tfr/keras/model/Preprocessor.md): Interface for
 feature preprocessing.
 
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/AbstractModelBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/AbstractModelBuilder.md
index 923aed6..75e2428 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/AbstractModelBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/AbstractModelBuilder.md
@@ -1,13 +1,9 @@
-description: Interface to build a ranking tf.keras.Model.
+description: Interface to build a tf.keras.Model for ranking.
 
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.model.AbstractModelBuilder" />
 <meta itemprop="path" content="Stable" />
-<meta itemprop="property" content="__init__"/>
 <meta itemprop="property" content="build"/>
-<meta itemprop="property" content="create_inputs"/>
-<meta itemprop="property" content="preprocess"/>
-<meta itemprop="property" content="score"/>
 </div>
 
 # tfr.keras.model.AbstractModelBuilder
@@ -16,355 +12,50 @@ description: Interface to build a ranking tf.keras.Model.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L67-L269">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L67-L97">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
 </td>
 </table>
 
-Interface to build a ranking `tf.keras.Model`.
-
-<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>tfr.keras.model.AbstractModelBuilder(
-    mask_feature_name: str,
-    name: Optional[str] = None
-)
-</code></pre>
+Interface to build a `tf.keras.Model` for ranking.
 
 <!-- Placeholder for "Used in" -->
 
-The `AbstractModelBuilder` class is an abstract class to build a ranking model
-in tfr.keras. All the boilerplate codes related to constructing a
-`tf.keras.Model` are integrated in the ModelBuilder class. This class is mostly
-designed to be passed to tfr.keras.pipeline and called to build the model under
-the strategy scope, so as all variables in the model, optimizers, and metrics.
-
-To be implemented by subclasses:
+The `AbstractModelBuilder` serves as the interface between model building and
+training. The training pipeline just calls the `build()` method to get the model
+constructed in the strategy scope used in the training pipeline, so for all
+variables in the model, optimizers, and metrics. See `ModelFitPipeline` in
+`pipeline.py` for example.
 
-*   `create_inputs()`: Contains the logic to create `tf.keras.Input` for context
-    and example inputs and mask for valid list items.
-*   `preprocess()`: Contains the logic to preprocess context and example inputs.
-*   `score()`: Contains the logic to score examples in list and return outputs.
-
-Example subclass implementation:
+The `build()` method is to be implemented in a subclass. The simplest example is
+just to define everything inside the build function when you define a
+tf.keras.Model.
 
 ```python
-class SimpleModelBuilder(AbstractModelBuilder):
-
-  def __init__(self, context_feature_spec, example_feature_spec,
-               mask_feature_name, name=None):
-    self._context_feature_spec = context_feature_spec
-    self._example_feature_spec = example_feature_spec
-    self._mask_feature_name = mask_feature_name
-    self._name = name
-
-  def create_inputs(self):
-    context_inputs = {
-        name: tf.keras.Input(
-            shape=tuple(spec.shape),
-            name=name,
-            dtype=spec.dtype
-        ) for name, spec in self._context_feature_spec.items()
-    }
-    example_inputs = {
-        name: tf.keras.Input(
-            shape=(None,) + tuple(spec.shape),
-            name=name,
-            dtype=spec.dtype
-        ) for name, spec in self._example_feature_spec.items()
-    }
-    mask = tf.keras.Input(
-        name=self._mask_feature_name, shape=(None,), dtype=tf.bool)
-    return context_inputs, example_inputs, mask
-
-  def preprocess(self, context_inputs, example_inputs, mask):
-    context_features = {
-        name: tf.math.log1p(
-            tf.abs(tensor)) for name, tensor in context_inputs.items()
-    }
-    example_features = {
-        name: tf.math.log1p(
-            tf.abs(tensor)) for name, tensor in example_inputs.items()
-    }
-    return context_features, example_features
+class MyModelBuilder(AbstractModelBuilder):
 
-  def score(self, context_features, example_features, mask):
-    x = tf.concat([tensor for tensor in example_features.values()], -1)
-    return tf.keras.layers.Dense(1)(x)
+  def build(self) -> tf.keras.Model:
+    inputs = ...
+    outputs = ...
+    return tf.keras.Model(inputs=inputs, outputs=outputs)
 ```
 
-<!-- Tabular view -->
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
-
-<tr>
-<td>
-`mask_feature_name`
-</td>
-<td>
-name of 2D mask boolean feature.
-</td>
-</tr><tr>
-<td>
-`name`
-</td>
-<td>
-(optional) name of the Model.
-</td>
-</tr>
-</table>
+The `MyModelBuilder` should work with `ModelFitPipeline`. To make the model
+building more structured for ranking problems, we also define subclasses like
+`ModelBuilderWithMask` in the following.
 
 ## Methods
 
 <h3 id="build"><code>build</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L241-L269">View
-source</a>
-
-<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>build() -> tf.keras.Model
-</code></pre>
-
-Builds a Keras Model for Ranking Pipeline.
-
-#### Example usage:
-
-```python
-model_builder = SimpleModelBuilder(
-    {},
-    {"example_feature_1": tf.io.FixedLenFeature(
-        shape=(1,), dtype=tf.float32, default_value=0.0)},
-    "list_mask", "model_builder")
-model = model_builder.build()
-```
-
-<!-- Tabular view -->
-
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Returns</th></tr>
-<tr class="alt">
-<td colspan="2">
-A `tf.keras.Model`.
-</td>
-</tr>
-
-</table>
-
-<h3 id="create_inputs"><code>create_inputs</code></h3>
-
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L146-L167">View
-source</a>
-
-<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>@abc.abstractmethod</code>
-<code>create_inputs() -> Tuple[TensorDict, TensorDict, tf.Tensor]
-</code></pre>
-
-Creates context and example inputs.
-
-#### Example usage:
-
-```python
-model_builder = SimpleModelBuilder(
-    {},
-    {"example_feature_1": tf.io.FixedLenFeature(
-        shape=(1,), dtype=tf.float32, default_value=0.0)},
-    "list_mask", "model_builder")
-context_inputs, example_inputs, mask = model_builder.create_inputs()
-```
-
-<!-- Tabular view -->
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Returns</th></tr>
-<tr class="alt">
-<td colspan="2">
-A tuple of
-</td>
-</tr>
-<tr>
-<td>
-`context_inputs`
-</td>
-<td>
-maps from context feature keys to Keras Input.
-</td>
-</tr><tr>
-<td>
-`example_inputs`
-</td>
-<td>
-maps from example feature keys to Keras Input.
-</td>
-</tr><tr>
-<td>
-`mask`
-</td>
-<td>
-Keras Input for the mask feature.
-</td>
-</tr>
-</table>
-
-<h3 id="preprocess"><code>preprocess</code></h3>
-
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L169-L203">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L94-L97">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@abc.abstractmethod</code>
-<code>preprocess(
-    context_inputs: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
-    example_inputs: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
-    mask: tf.Tensor
-) -> Tuple[TensorDict, TensorDict]
-</code></pre>
-
-Preprocesses context and example inputs.
-
-#### Example usage:
-
-```python
-model_builder = SimpleModelBuilder(
-    {},
-    {"example_feature_1": tf.io.FixedLenFeature(
-        shape=(1,), dtype=tf.float32, default_value=0.0)},
-    "list_mask", "model_builder")
-context_inputs, example_inputs, mask = model_builder.create_inputs()
-context_features, example_features = model_builder.preprocess(
-    context_inputs, example_inputs, mask)
-```
-
-<!-- Tabular view -->
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Args</th></tr>
-
-<tr>
-<td>
-`context_inputs`
-</td>
-<td>
-maps context feature keys to `tf.keras.Input`.
-</td>
-</tr><tr>
-<td>
-`example_inputs`
-</td>
-<td>
-maps example feature keys to `tf.keras.Input`.
-</td>
-</tr><tr>
-<td>
-`mask`
-</td>
-<td>
-[batch_size, list_size]-tensor of mask for valid examples.
-</td>
-</tr>
-</table>
-
-<!-- Tabular view -->
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Returns</th></tr>
-<tr class="alt">
-<td colspan="2">
-A tuple of
-</td>
-</tr>
-<tr>
-<td>
-`context_features`
-</td>
-<td>
-maps from context feature keys to [batch_size,
-feature_dims]-tensors of preprocessed context features.
-</td>
-</tr><tr>
-<td>
-`example_features`
-</td>
-<td>
-maps from example feature keys to [batch_size,
-list_size, feature_dims]-tensors of preprocessed example features.
-</td>
-</tr>
-</table>
-
-<h3 id="score"><code>score</code></h3>
-
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L205-L239">View
-source</a>
-
-<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>@abc.abstractmethod</code>
-<code>score(
-    context_features: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
-    example_features: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
-    mask: tf.Tensor
-) -> Union[TensorLike, TensorDict]
+<code>build() -> tf.keras.Model
 </code></pre>
 
-Scores all examples and returns outputs.
-
-#### Example usage:
-
-```python
-model_builder = SimpleModelBuilder(
-    {},
-    {"example_feature_1": tf.io.FixedLenFeature(
-        shape=(1,), dtype=tf.float32, default_value=0.0)},
-    "list_mask", "model_builder")
-context_inputs, example_inputs, mask = model_builder.create_inputs()
-context_features, example_features = model_builder.preprocess(
-    context_inputs, example_inputs, mask)
-scores = model_builder.score(context_features, example_features)
-```
-
-<!-- Tabular view -->
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Args</th></tr>
-
-<tr>
-<td>
-`context_features`
-</td>
-<td>
-maps from context feature keys to [batch_size,
-feature_dims]-tensors of preprocessed context features.
-</td>
-</tr><tr>
-<td>
-`example_features`
-</td>
-<td>
-maps from example feature keys to [batch_size,
-list_size, feature_dims]-tensors of preprocessed example features.
-</td>
-</tr><tr>
-<td>
-`mask`
-</td>
-<td>
-[batch_size, list_size]-tensor of mask for valid examples.
-</td>
-</tr>
-</table>
-
-<!-- Tabular view -->
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Returns</th></tr>
-<tr class="alt">
-<td colspan="2">
-A [batch_size, list_size]-tensor of logits or a dict mapping task name to
-logits in the multi-task setting.
-</td>
-</tr>
-
-</table>
+The build method to be implemented by a subclass.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/DNNScorer.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/DNNScorer.md
index 1e3802c..3bd4f6c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/DNNScorer.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/DNNScorer.md
@@ -13,7 +13,7 @@ description: Univariate scorer using DNN.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L736-L773">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L767-L804">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -41,7 +41,6 @@ scorer=DNNScorer(hidden_layer_dims=[16])
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -62,7 +61,7 @@ arguments.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L711-L733">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L742-L764">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/FeatureSpecInputCreator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/FeatureSpecInputCreator.md
index 6716674..e81f496 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/FeatureSpecInputCreator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/FeatureSpecInputCreator.md
@@ -13,7 +13,7 @@ description: InputCreator with feature specs.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L386-L454">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L417-L485">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -45,7 +45,6 @@ input_creator=FeatureSpecInputCreator(
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -73,7 +72,7 @@ corresponding example feature specs.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L419-L454">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L450-L485">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/GAMScorer.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/GAMScorer.md
index c7f1e33..c105ac5 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/GAMScorer.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/GAMScorer.md
@@ -13,7 +13,7 @@ description: Univariate scorer using GAM.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L776-L821">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L807-L852">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -45,7 +45,6 @@ scorer=GAMScorer(hidden_layer_dims=[16])
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -65,7 +64,7 @@ A dict of keyward arguments for GAM layers. Please see
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L711-L733">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L742-L764">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/InputCreator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/InputCreator.md
index 5b0ad35..0b6ecf0 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/InputCreator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/InputCreator.md
@@ -12,7 +12,7 @@ description: Interface for input creator.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L354-L383">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L385-L414">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -44,7 +44,7 @@ class SimpleInputCreator(InputCreator):
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L375-L383">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L406-L414">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -55,7 +55,6 @@ source</a>
 Invokes the `InputCreator` instance.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilder.md
index 912005c..6d9fdb5 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilder.md
@@ -16,7 +16,7 @@ description: Builds a tf.keras.Model.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L272-L351">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L303-L382">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -26,6 +26,7 @@ description: Builds a tf.keras.Model.
 Builds a `tf.keras.Model`.
 
 Inherits From:
+[`ModelBuilderWithMask`](../../../tfr/keras/model/ModelBuilderWithMask.md),
 [`AbstractModelBuilder`](../../../tfr/keras/model/AbstractModelBuilder.md)
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -40,7 +41,7 @@ Inherits From:
 
 <!-- Placeholder for "Used in" -->
 
-This class implements the `AbstractModelBuilder` by delegating the class
+This class implements the `ModelBuilderWithMask` by delegating the class
 behaviors to the following implementors that can be specified by callers:
 
 *   input_creator: A callable or a class like `InputCreator` to implement
@@ -67,7 +68,6 @@ model_builder = ModelBuilder(
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -116,7 +116,7 @@ name of 2D mask boolean feature.
 
 <h3 id="build"><code>build</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L241-L269">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L272-L300">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -137,7 +137,6 @@ model = model_builder.build()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -151,18 +150,18 @@ A `tf.keras.Model`.
 
 <h3 id="create_inputs"><code>create_inputs</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L328-L333">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L359-L364">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>create_inputs() -> Tuple[TensorDict, TensorDict, tf.Tensor]
 </code></pre>
 
-See `AbstractModelBuilder`.
+See `ModelBuilderWithMask`.
 
 <h3 id="preprocess"><code>preprocess</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L335-L342">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L366-L373">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -173,11 +172,11 @@ source</a>
 ) -> Tuple[TensorDict, TensorDict]
 </code></pre>
 
-See `AbstractModelBuilder`.
+See `ModelBuilderWithMask`.
 
 <h3 id="score"><code>score</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L344-L351">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L375-L382">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -188,4 +187,4 @@ source</a>
 ) -> Union[TensorLike, TensorDict]
 </code></pre>
 
-See `AbstractModelBuilder`.
+See `ModelBuilderWithMask`.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilderWithMask.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilderWithMask.md
new file mode 100644
index 0000000..1e53d31
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/ModelBuilderWithMask.md
@@ -0,0 +1,378 @@
+description: Interface to build a tf.keras.Model for ranking with a mask Tensor.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.model.ModelBuilderWithMask" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="build"/>
+<meta itemprop="property" content="create_inputs"/>
+<meta itemprop="property" content="preprocess"/>
+<meta itemprop="property" content="score"/>
+</div>
+
+# tfr.keras.model.ModelBuilderWithMask
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L100-L300">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Interface to build a `tf.keras.Model` for ranking with a mask Tensor.
+
+Inherits From:
+[`AbstractModelBuilder`](../../../tfr/keras/model/AbstractModelBuilder.md)
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.model.ModelBuilderWithMask(
+    mask_feature_name: str,
+    name: Optional[str] = None
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+The `ModelBuilderWithMask` class is an abstract class to build a ranking model
+based on dense Tensors and a mask Tensor to indicate the padded ones. All the
+boilerplate codes related to constructing a `tf.keras.Model` are integrated in
+the ModelBuilder class.
+
+To be implemented by subclasses:
+
+*   `create_inputs()`: Contains the logic to create `tf.keras.Input` for context
+    and example inputs and mask for valid list items.
+*   `preprocess()`: Contains the logic to preprocess context and example inputs.
+*   `score()`: Contains the logic to score examples in list and return outputs.
+
+Example subclass implementation:
+
+```python
+class SimpleModelBuilder(ModelBuilderWithMask):
+
+  def __init__(self, context_feature_spec, example_feature_spec,
+               mask_feature_name, name=None):
+    self._context_feature_spec = context_feature_spec
+    self._example_feature_spec = example_feature_spec
+    self._mask_feature_name = mask_feature_name
+    self._name = name
+
+  def create_inputs(self):
+    context_inputs = {
+        name: tf.keras.Input(
+            shape=tuple(spec.shape),
+            name=name,
+            dtype=spec.dtype
+        ) for name, spec in self._context_feature_spec.items()
+    }
+    example_inputs = {
+        name: tf.keras.Input(
+            shape=(None,) + tuple(spec.shape),
+            name=name,
+            dtype=spec.dtype
+        ) for name, spec in self._example_feature_spec.items()
+    }
+    mask = tf.keras.Input(
+        name=self._mask_feature_name, shape=(None,), dtype=tf.bool)
+    return context_inputs, example_inputs, mask
+
+  def preprocess(self, context_inputs, example_inputs, mask):
+    context_features = {
+        name: tf.math.log1p(
+            tf.abs(tensor)) for name, tensor in context_inputs.items()
+    }
+    example_features = {
+        name: tf.math.log1p(
+            tf.abs(tensor)) for name, tensor in example_inputs.items()
+    }
+    return context_features, example_features
+
+  def score(self, context_features, example_features, mask):
+    x = tf.concat([tensor for tensor in example_features.values()], -1)
+    return tf.keras.layers.Dense(1)(x)
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`mask_feature_name`
+</td>
+<td>
+name of 2D mask boolean feature.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+(optional) name of the Model.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="build"><code>build</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L272-L300">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>build() -> tf.keras.Model
+</code></pre>
+
+Builds a Keras Model for Ranking Pipeline.
+
+#### Example usage:
+
+```python
+model_builder = SimpleModelBuilder(
+    {},
+    {"example_feature_1": tf.io.FixedLenFeature(
+        shape=(1,), dtype=tf.float32, default_value=0.0)},
+    "list_mask", "model_builder")
+model = model_builder.build()
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A `tf.keras.Model`.
+</td>
+</tr>
+
+</table>
+
+<h3 id="create_inputs"><code>create_inputs</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L177-L198">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@abc.abstractmethod</code>
+<code>create_inputs() -> Tuple[TensorDict, TensorDict, tf.Tensor]
+</code></pre>
+
+Creates context and example inputs.
+
+#### Example usage:
+
+```python
+model_builder = SimpleModelBuilder(
+    {},
+    {"example_feature_1": tf.io.FixedLenFeature(
+        shape=(1,), dtype=tf.float32, default_value=0.0)},
+    "list_mask", "model_builder")
+context_inputs, example_inputs, mask = model_builder.create_inputs()
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A tuple of
+</td>
+</tr>
+<tr>
+<td>
+`context_inputs`
+</td>
+<td>
+maps from context feature keys to Keras Input.
+</td>
+</tr><tr>
+<td>
+`example_inputs`
+</td>
+<td>
+maps from example feature keys to Keras Input.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+Keras Input for the mask feature.
+</td>
+</tr>
+</table>
+
+<h3 id="preprocess"><code>preprocess</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L200-L234">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@abc.abstractmethod</code>
+<code>preprocess(
+    context_inputs: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
+    example_inputs: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
+    mask: tf.Tensor
+) -> Tuple[TensorDict, TensorDict]
+</code></pre>
+
+Preprocesses context and example inputs.
+
+#### Example usage:
+
+```python
+model_builder = SimpleModelBuilder(
+    {},
+    {"example_feature_1": tf.io.FixedLenFeature(
+        shape=(1,), dtype=tf.float32, default_value=0.0)},
+    "list_mask", "model_builder")
+context_inputs, example_inputs, mask = model_builder.create_inputs()
+context_features, example_features = model_builder.preprocess(
+    context_inputs, example_inputs, mask)
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`context_inputs`
+</td>
+<td>
+maps context feature keys to `tf.keras.Input`.
+</td>
+</tr><tr>
+<td>
+`example_inputs`
+</td>
+<td>
+maps example feature keys to `tf.keras.Input`.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+[batch_size, list_size]-tensor of mask for valid examples.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A tuple of
+</td>
+</tr>
+<tr>
+<td>
+`context_features`
+</td>
+<td>
+maps from context feature keys to [batch_size,
+feature_dims]-tensors of preprocessed context features.
+</td>
+</tr><tr>
+<td>
+`example_features`
+</td>
+<td>
+maps from example feature keys to [batch_size,
+list_size, feature_dims]-tensors of preprocessed example features.
+</td>
+</tr>
+</table>
+
+<h3 id="score"><code>score</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L236-L270">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@abc.abstractmethod</code>
+<code>score(
+    context_features: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
+    example_features: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
+    mask: tf.Tensor
+) -> Union[TensorLike, TensorDict]
+</code></pre>
+
+Scores all examples and returns outputs.
+
+#### Example usage:
+
+```python
+model_builder = SimpleModelBuilder(
+    {},
+    {"example_feature_1": tf.io.FixedLenFeature(
+        shape=(1,), dtype=tf.float32, default_value=0.0)},
+    "list_mask", "model_builder")
+context_inputs, example_inputs, mask = model_builder.create_inputs()
+context_features, example_features = model_builder.preprocess(
+    context_inputs, example_inputs, mask)
+scores = model_builder.score(context_features, example_features)
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`context_features`
+</td>
+<td>
+maps from context feature keys to [batch_size,
+feature_dims]-tensors of preprocessed context features.
+</td>
+</tr><tr>
+<td>
+`example_features`
+</td>
+<td>
+maps from example feature keys to [batch_size,
+list_size, feature_dims]-tensors of preprocessed example features.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+[batch_size, list_size]-tensor of mask for valid examples.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A [batch_size, list_size]-tensor of logits or a dict mapping task name to
+logits in the multi-task setting.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Preprocessor.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Preprocessor.md
index 4a5761d..673ca56 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Preprocessor.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Preprocessor.md
@@ -12,7 +12,7 @@ description: Interface for feature preprocessing.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L508-L554">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L539-L585">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -51,7 +51,7 @@ class SimplePreprocessor(Preprocessor):
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L536-L554">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L567-L585">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -66,7 +66,6 @@ source</a>
 Invokes the `Preprocessor` instance.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -96,7 +95,6 @@ maps example feature keys to `tf.keras.Input`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/PreprocessorWithSpec.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/PreprocessorWithSpec.md
index 199a4af..d26e031 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/PreprocessorWithSpec.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/PreprocessorWithSpec.md
@@ -13,7 +13,7 @@ description: Preprocessing inputs with provided spec.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L557-L621">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L588-L652">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -75,7 +75,7 @@ RaggedTensor to Tensor. Default to 0. if not specified.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L591-L621">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L622-L652">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Scorer.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Scorer.md
index 9199767..146381f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Scorer.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/Scorer.md
@@ -12,7 +12,7 @@ description: Interface for scorer.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L624-L666">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L655-L697">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -45,7 +45,7 @@ class SimpleScorer(Scorer):
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L646-L666">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L677-L697">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/TypeSpecInputCreator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/TypeSpecInputCreator.md
index 95b4c3d..5733539 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/TypeSpecInputCreator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/TypeSpecInputCreator.md
@@ -13,7 +13,7 @@ description: InputCreator with tensor type specs.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L457-L505">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L488-L536">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -43,7 +43,6 @@ input_creator=TypeSpecInputCreator(
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -77,7 +76,7 @@ A list of example feature keys.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L487-L505">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L518-L536">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/UnivariateScorer.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/UnivariateScorer.md
index 9e0d115..a75128f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/UnivariateScorer.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/UnivariateScorer.md
@@ -12,7 +12,7 @@ description: Interface for univariate scorer.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L669-L733">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L700-L764">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -47,7 +47,7 @@ class SimpleUnivariateScorer(UnivariateScorer):
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L711-L733">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L742-L764">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractDatasetBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractDatasetBuilder.md
index 3f3385a..1135a0e 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractDatasetBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractDatasetBuilder.md
@@ -85,7 +85,6 @@ signatures = dataset_builder.build_signatures()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -108,7 +107,6 @@ keyword arguments that might be used to build signatures.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -142,7 +140,6 @@ train_dataset = dataset_builder.build_train_dataset()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -165,7 +162,6 @@ keyword arguments that might be used to build training dataset.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -199,7 +195,6 @@ valid_dataset = dataset_builder.build_valid_dataset()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -223,7 +218,6 @@ dataset.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractPipeline.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractPipeline.md
index c83feb9..b79c775 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractPipeline.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/AbstractPipeline.md
@@ -102,7 +102,6 @@ loss = pipeline.build_loss()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -134,7 +133,6 @@ metrics = pipeline.build_metrics()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -166,7 +164,6 @@ weighted_metrics = pipeline.build_weighted_metrics()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -200,7 +197,6 @@ pipeline.train_and_validate()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -223,7 +219,6 @@ keyword arguments that might be used in the training pipeline.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/BaseDatasetBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/BaseDatasetBuilder.md
index a046efd..6ef6f6d 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/BaseDatasetBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/BaseDatasetBuilder.md
@@ -15,7 +15,7 @@ description: Builds datasets from feature specs.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L828-L973">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L824-L969">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -121,7 +121,7 @@ weights.
 
 <h3 id="build_signatures"><code>build_signatures</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L967-L973">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L963-L969">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -134,7 +134,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_train_dataset"><code>build_train_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L950-L956">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L946-L952">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -145,7 +145,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_valid_dataset"><code>build_valid_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L958-L965">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L954-L961">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/DatasetHparams.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/DatasetHparams.md
index ba28609..b86108f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/DatasetHparams.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/DatasetHparams.md
@@ -17,7 +17,7 @@ description: Hyperparameters used in BaseDatasetBuilder.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L311-L339">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L310-L338">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -28,6 +28,7 @@ Hyperparameters used in `BaseDatasetBuilder`.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.pipeline.DatasetHparams(
+    train_input_pattern: str,
     valid_input_pattern: str,
     train_batch_size: int,
     valid_batch_size: int,
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/ModelFitPipeline.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/ModelFitPipeline.md
index 0417c30..1b500cc 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/ModelFitPipeline.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/ModelFitPipeline.md
@@ -18,7 +18,7 @@ description: Pipeline using model.fit to train a ranking tf.keras.Model.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L342-L612">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L341-L608">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -110,7 +110,7 @@ A dict containing model hyperparameters.
 
 <h3 id="build_callbacks"><code>build_callbacks</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L444-L491">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L443-L490">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -130,7 +130,6 @@ callbacks = pipeline.build_callbacks()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -163,7 +162,6 @@ loss = pipeline.build_loss()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -195,7 +193,6 @@ metrics = pipeline.build_metrics()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -227,7 +224,6 @@ weighted_metrics = pipeline.build_weighted_metrics()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -241,7 +237,7 @@ A list or a dict of `tf.keras.metrics.Metric`s.
 
 <h3 id="export_saved_model"><code>export_saved_model</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L493-L518">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L492-L517">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -295,7 +291,7 @@ If given, export the model with weights from this checkpoint.
 
 <h3 id="train_and_validate"><code>train_and_validate</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L520-L612">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L519-L608">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -327,8 +323,8 @@ dataset_hparams = DatasetHparams(
 pipeline_hparams = pipeline.PipelineHparams(
     model_dir="model/",
     num_epochs=2,
-    num_train_steps=10,
-    num_valid_steps=2,
+    steps_per_epoch=5,
+    validation_steps=2,
     learning_rate=0.01,
     loss="softmax_loss")
 model_builder = SimpleModelBuilder(
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiLabelDatasetBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiLabelDatasetBuilder.md
index 7507e0b..4d379d3 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiLabelDatasetBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiLabelDatasetBuilder.md
@@ -15,7 +15,7 @@ description: Builds datasets for multi-task training.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L1086-L1187">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L1082-L1183">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -136,7 +136,7 @@ Feature spec for per-example weight.
 
 <h3 id="build_signatures"><code>build_signatures</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L967-L973">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L963-L969">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -149,7 +149,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_train_dataset"><code>build_train_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L950-L956">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L946-L952">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -160,7 +160,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_valid_dataset"><code>build_valid_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L958-L965">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L954-L961">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiTaskPipeline.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiTaskPipeline.md
index 04590ea..8e903af 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiTaskPipeline.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/MultiTaskPipeline.md
@@ -18,7 +18,7 @@ description: Pipeline for multi-task training.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L695-L786">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L691-L782">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -72,8 +72,8 @@ dataset_hparams = DatasetHparams(
 pipeline_hparams = PipelineHparams(
     model_dir="model/",
     num_epochs=2,
-    num_train_steps=10,
-    num_valid_steps=2,
+    steps_per_epoch=5,
+    validation_steps=2,
     learning_rate=0.01,
     loss={
         "task1": "softmax_loss",
@@ -130,7 +130,7 @@ A dict containing model hyperparameters.
 
 <h3 id="build_callbacks"><code>build_callbacks</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L444-L491">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L443-L490">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -150,7 +150,6 @@ callbacks = pipeline.build_callbacks()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -165,7 +164,7 @@ A list of `tf.keras.callbacks.Callback` or a
 
 <h3 id="build_loss"><code>build_loss</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L754-L763">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L750-L759">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -176,7 +175,7 @@ See `AbstractPipeline`.
 
 <h3 id="build_metrics"><code>build_metrics</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L765-L774">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L761-L770">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -187,7 +186,7 @@ See `AbstractPipeline`.
 
 <h3 id="build_weighted_metrics"><code>build_weighted_metrics</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L776-L786">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L772-L782">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -198,7 +197,7 @@ See `AbstractPipeline`.
 
 <h3 id="export_saved_model"><code>export_saved_model</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L493-L518">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L492-L517">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -252,7 +251,7 @@ If given, export the model with weights from this checkpoint.
 
 <h3 id="train_and_validate"><code>train_and_validate</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L520-L612">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L519-L608">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -284,8 +283,8 @@ dataset_hparams = DatasetHparams(
 pipeline_hparams = pipeline.PipelineHparams(
     model_dir="model/",
     num_epochs=2,
-    num_train_steps=10,
-    num_valid_steps=2,
+    steps_per_epoch=5,
+    validation_steps=2,
     learning_rate=0.01,
     loss="softmax_loss")
 model_builder = SimpleModelBuilder(
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/NullDatasetBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/NullDatasetBuilder.md
index ce05f84..0d20e02 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/NullDatasetBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/NullDatasetBuilder.md
@@ -15,7 +15,7 @@ description: A no-op wrapper of datasets and signatures.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L789-L825">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L785-L821">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -44,7 +44,6 @@ dataset_builder = NullDatasetBuilder(train_dataset, valid_dataset)
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -79,7 +78,7 @@ assigned.
 
 <h3 id="build_signatures"><code>build_signatures</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L823-L825">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L819-L821">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -92,7 +91,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_train_dataset"><code>build_train_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L815-L817">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L811-L813">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -105,7 +104,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_valid_dataset"><code>build_valid_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L819-L821">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L815-L817">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/PipelineHparams.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/PipelineHparams.md
index 4acf7c3..fbf8305 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/PipelineHparams.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/PipelineHparams.md
@@ -24,7 +24,7 @@ description: Hyperparameters used in ModelFitPipeline.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L249-L307">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L249-L306">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -35,9 +35,10 @@ Hyperparameters used in `ModelFitPipeline`.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.pipeline.PipelineHparams(
+    model_dir: str,
     num_epochs: int,
-    num_train_steps: int,
-    num_valid_steps: int,
+    steps_per_epoch: int,
+    validation_steps: int,
     learning_rate: float,
     loss: Union[str, Dict[str, str]],
     loss_reduction: str = &#x27;auto&#x27;,
@@ -79,21 +80,20 @@ An integer to specify the number of epochs of training.
 </td>
 </tr><tr>
 <td>
-`num_train_steps`
+`steps_per_epoch`
 </td>
 <td>
-An integer to specify the total number of training steps.
-Note that a mini-batch of data will be treated in each step and
-`num_train_steps` / `num_epochs` of steps will be taken in one epoch.
+An integer to specify the number of steps per epoch. When
+it is None, going over the training data once is counted as an epoch.
 </td>
 </tr><tr>
 <td>
-`num_valid_steps`
+`validation_steps`
 </td>
 <td>
 An integer to specify the number of validation steps in
 each epoch. Note that a mini-batch of data will be evaluated in each step
-and `num_valid_steps` of steps will be taken for validation in each epoch.
+and this is the number of steps taken for validation in each epoch.
 </td>
 </tr><tr>
 <td>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimpleDatasetBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimpleDatasetBuilder.md
index 03f6870..73226f1 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimpleDatasetBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimpleDatasetBuilder.md
@@ -15,7 +15,7 @@ description: Builds datasets from feature specs with a single label spec.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L992-L1083">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L988-L1079">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -129,7 +129,7 @@ Feature spec for per-example weight.
 
 <h3 id="build_signatures"><code>build_signatures</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L967-L973">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L963-L969">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -142,7 +142,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_train_dataset"><code>build_train_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L950-L956">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L946-L952">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -153,7 +153,7 @@ See `AbstractDatasetBuilder`.
 
 <h3 id="build_valid_dataset"><code>build_valid_dataset</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L958-L965">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L954-L961">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimplePipeline.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimplePipeline.md
index 3834598..eb99e37 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimplePipeline.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/pipeline/SimplePipeline.md
@@ -18,7 +18,7 @@ description: Pipleine for single-task training.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L621-L692">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L617-L688">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -68,8 +68,8 @@ dataset_hparams = DatasetHparams(
 pipeline_hparams = pipeline.PipelineHparams(
     model_dir="model/",
     num_epochs=2,
-    num_train_steps=10,
-    num_valid_steps=2,
+    steps_per_epoch=5,
+    validation_steps=2,
     learning_rate=0.01,
     loss="softmax_loss")
 model_builder = SimpleModelBuilder(
@@ -118,7 +118,7 @@ A dict containing model hyperparameters.
 
 <h3 id="build_callbacks"><code>build_callbacks</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L444-L491">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L443-L490">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -138,7 +138,6 @@ callbacks = pipeline.build_callbacks()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -153,7 +152,7 @@ A list of `tf.keras.callbacks.Callback` or a
 
 <h3 id="build_loss"><code>build_loss</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L669-L675">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L665-L671">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -164,7 +163,7 @@ See `AbstractPipeline`.
 
 <h3 id="build_metrics"><code>build_metrics</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L677-L683">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L673-L679">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -175,7 +174,7 @@ See `AbstractPipeline`.
 
 <h3 id="build_weighted_metrics"><code>build_weighted_metrics</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L685-L692">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L681-L688">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -186,7 +185,7 @@ See `AbstractPipeline`.
 
 <h3 id="export_saved_model"><code>export_saved_model</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L493-L518">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L492-L517">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -240,7 +239,7 @@ If given, export the model with weights from this checkpoint.
 
 <h3 id="train_and_validate"><code>train_and_validate</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L520-L612">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/pipeline.py#L519-L608">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -272,8 +271,8 @@ dataset_hparams = DatasetHparams(
 pipeline_hparams = pipeline.PipelineHparams(
     model_dir="model/",
     num_epochs=2,
-    num_train_steps=10,
-    num_valid_steps=2,
+    steps_per_epoch=5,
+    validation_steps=2,
     learning_rate=0.01,
     loss="softmax_loss")
 model_builder = SimpleModelBuilder(
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/recall.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/recall.md
index 4c488bf..dbf1627 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/recall.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/recall.md
@@ -27,8 +27,8 @@ Computes recall as weighted average of relevant examples.
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
-<!-- Tabular view -->
 
+<!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
@@ -75,7 +75,6 @@ A string used as the name for this metric.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
diff --git a/tensorflow_ranking/python/version.py b/tensorflow_ranking/python/version.py
index 8685c4f..61cd052 100644
--- a/tensorflow_ranking/python/version.py
+++ b/tensorflow_ranking/python/version.py
@@ -16,8 +16,8 @@
 
 # We follow Semantic Versioning (https://semver.org/)
 _MAJOR_VERSION = '0'
-_MINOR_VERSION = '3'
-_PATCH_VERSION = '3'
+_MINOR_VERSION = '4'
+_PATCH_VERSION = '0'
 
 # When building releases, we can update this value on the release branch to
 # reflect the current release candidate ('rc0', 'rc1') or, finally, the official
diff --git a/tensorflow_ranking/tools/pip_package/setup.py b/tensorflow_ranking/tools/pip_package/setup.py
index 784947e..bb57128 100644
--- a/tensorflow_ranking/tools/pip_package/setup.py
+++ b/tensorflow_ranking/tools/pip_package/setup.py
@@ -28,11 +28,11 @@
 # This version string is semver compatible, but incompatible with pip.
 # For pip, we will remove all '-' characters from this string, and use the
 # result for pip.
-_VERSION = '0.3.3'
+_VERSION = '0.4.0'
 
 REQUIRED_PACKAGES = [
     'absl-py >= 0.1.6', 'numpy >= 1.13.3', 'six >= 1.10.0',
-    'tensorflow-serving-api >= 2.0.0, < 3.0.0', 'tf-models-official >= 2.4.0'
+    'tensorflow-serving-api >= 2.0.0, < 3.0.0', 'tf-models-official >= 2.5.0'
 ]
 
 project_name = 'tensorflow_ranking'