Update version for release 0.5.2

PiperOrigin-RevId: 506166467

docs/api_docs/python/tfr.md

@@ -42,7 +42,7 @@ TensorFlow Ranking, the project to build ranking models on TensorFlow.
 **version**<a id="__version__"></a>
 </td>
 <td>
-`'0.5.0.dev'`
+`'0.5.2.dev'`
 </td>
 </tr>
 </table>

docs/api_docs/python/tfr/_redirects.yaml

@@ -25,6 +25,8 @@ redirects:
   to: /ranking/api_docs/python/tfr/keras/utils/GainFunction
 - from: /ranking/api_docs/python/tfr/keras/utils/TensorLike
   to: /ranking/api_docs/python/tfr/keras/model/TensorLike
+- from: /ranking/api_docs/python/tfr/utils/MetricFunction
+  to: /ranking/api_docs/python/tfr/utils/LossFunction
 - from: /ranking/api_docs/python/tfr/utils/TensorLike
   to: /ranking/api_docs/python/tfr/keras/model/TensorLike
 - from: /ranking/api_docs/python/tfr/utils/TransformationFunction

docs/api_docs/python/tfr/_toc.yaml

@@ -113,6 +113,8 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/losses/ApproxNDCGLoss
     - title: ClickEMLoss
       path: /ranking/api_docs/python/tfr/keras/losses/ClickEMLoss
+    - title: CoupledRankDistilLoss
+      path: /ranking/api_docs/python/tfr/keras/losses/CoupledRankDistilLoss
     - title: DCGLambdaWeight
       path: /ranking/api_docs/python/tfr/keras/losses/DCGLambdaWeight
     - title: GumbelApproxNDCGLoss
@@ -237,8 +239,6 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/pipeline/ModelFitPipeline
     - title: MultiLabelDatasetBuilder
       path: /ranking/api_docs/python/tfr/keras/pipeline/MultiLabelDatasetBuilder
-    - title: MultiObjectivePipeline
-      path: /ranking/api_docs/python/tfr/keras/pipeline/MultiObjectivePipeline
     - title: MultiTaskPipeline
       path: /ranking/api_docs/python/tfr/keras/pipeline/MultiTaskPipeline
     - title: NullDatasetBuilder
@@ -301,6 +301,8 @@ toc:
     path: /ranking/api_docs/python/tfr/utils/organize_valid_indices
   - title: padded_nd_indices
     path: /ranking/api_docs/python/tfr/utils/padded_nd_indices
+  - title: parse_keys_and_weights
+    path: /ranking/api_docs/python/tfr/utils/parse_keys_and_weights
   - title: ragged_to_dense
     path: /ranking/api_docs/python/tfr/utils/ragged_to_dense
   - title: reshape_first_ndims

docs/api_docs/python/tfr/all_symbols.md

@@ -61,6 +61,7 @@
 *   <a href="../tfr/keras/losses/ApproxMRRLoss.md"><code>tfr.keras.losses.ApproxMRRLoss</code></a>
 *   <a href="../tfr/keras/losses/ApproxNDCGLoss.md"><code>tfr.keras.losses.ApproxNDCGLoss</code></a>
 *   <a href="../tfr/keras/losses/ClickEMLoss.md"><code>tfr.keras.losses.ClickEMLoss</code></a>
+*   <a href="../tfr/keras/losses/CoupledRankDistilLoss.md"><code>tfr.keras.losses.CoupledRankDistilLoss</code></a>
 *   <a href="../tfr/keras/losses/DCGLambdaWeight.md"><code>tfr.keras.losses.DCGLambdaWeight</code></a>
 *   <a href="../tfr/keras/losses/GumbelApproxNDCGLoss.md"><code>tfr.keras.losses.GumbelApproxNDCGLoss</code></a>
 *   <a href="../tfr/keras/losses/LabelDiffLambdaWeight.md"><code>tfr.keras.losses.LabelDiffLambdaWeight</code></a>
@@ -120,7 +121,6 @@
 *   <a href="../tfr/keras/pipeline/DatasetHparams/dataset_reader.md"><code>tfr.keras.pipeline.DatasetHparams.dataset_reader</code></a>
 *   <a href="../tfr/keras/pipeline/ModelFitPipeline.md"><code>tfr.keras.pipeline.ModelFitPipeline</code></a>
 *   <a href="../tfr/keras/pipeline/MultiLabelDatasetBuilder.md"><code>tfr.keras.pipeline.MultiLabelDatasetBuilder</code></a>
-*   <a href="../tfr/keras/pipeline/MultiObjectivePipeline.md"><code>tfr.keras.pipeline.MultiObjectivePipeline</code></a>
 *   <a href="../tfr/keras/pipeline/MultiTaskPipeline.md"><code>tfr.keras.pipeline.MultiTaskPipeline</code></a>
 *   <a href="../tfr/keras/pipeline/NullDatasetBuilder.md"><code>tfr.keras.pipeline.NullDatasetBuilder</code></a>
 *   <a href="../tfr/keras/pipeline/PipelineHparams.md"><code>tfr.keras.pipeline.PipelineHparams</code></a>
@@ -146,13 +146,15 @@
 *   <a href="../tfr/keras/utils/symmetric_log1p.md"><code>tfr.keras.utils.symmetric_log1p</code></a>
 *   <a href="../tfr/utils.md"><code>tfr.utils</code></a>
 *   <a href="../tfr/utils/LossFunction.md"><code>tfr.utils.LossFunction</code></a>
+*   <a href="../tfr/utils/LossFunction.md"><code>tfr.utils.MetricFunction</code></a>
 *   <a href="../tfr/keras/model/TensorLike.md"><code>tfr.utils.TensorLike</code></a>
 *   <a href="../tfr/keras/utils/GainFunction.md"><code>tfr.utils.TransformationFunction</code></a>
 *   <a href="../tfr/utils/de_noise.md"><code>tfr.utils.de_noise</code></a>
 *   <a href="../tfr/utils/gather_per_row.md"><code>tfr.utils.gather_per_row</code></a>
 *   <a href="../tfr/utils/is_label_valid.md"><code>tfr.utils.is_label_valid</code></a>
 *   <a href="../tfr/utils/organize_valid_indices.md"><code>tfr.utils.organize_valid_indices</code></a>
 *   <a href="../tfr/utils/padded_nd_indices.md"><code>tfr.utils.padded_nd_indices</code></a>
+*   <a href="../tfr/utils/parse_keys_and_weights.md"><code>tfr.utils.parse_keys_and_weights</code></a>
 *   <a href="../tfr/utils/ragged_to_dense.md"><code>tfr.utils.ragged_to_dense</code></a>
 *   <a href="../tfr/utils/reshape_first_ndims.md"><code>tfr.utils.reshape_first_ndims</code></a>
 *   <a href="../tfr/utils/reshape_to_2d.md"><code>tfr.utils.reshape_to_2d</code></a>

docs/api_docs/python/tfr/extension/pipeline.md

@@ -3,9 +3,6 @@ description: Provides a RankingPipeline for running a TF-Ranking model.
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.extension.pipeline" />
 <meta itemprop="path" content="Stable" />
-<meta itemprop="property" content="absolute_import"/>
-<meta itemprop="property" content="division"/>
-<meta itemprop="property" content="print_function"/>
 </div>
 
 # Module: tfr.extension.pipeline
@@ -32,33 +29,3 @@ derive from this class and further tailor for their needs.
 
 [`class RankingPipeline`](../../tfr/extension/pipeline/RankingPipeline.md):
 Class to set up the input, train and eval processes for a TF Ranking model.
-
-<!-- Tabular view -->
-
- <table class="responsive fixed orange">
-<colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2"><h2 class="add-link">Other Members</h2></th></tr>
-
-<tr>
-<td>
-absolute_import<a id="absolute_import"></a>
-</td>
-<td>
-Instance of `__future__._Feature`
-</td>
-</tr><tr>
-<td>
-division<a id="division"></a>
-</td>
-<td>
-Instance of `__future__._Feature`
-</td>
-</tr><tr>
-<td>
-print_function<a id="print_function"></a>
-</td>
-<td>
-Instance of `__future__._Feature`
-</td>
-</tr>
-</table>

docs/api_docs/python/tfr/extension/pipeline/RankingPipeline.md

@@ -14,7 +14,7 @@ Ranking 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/extension/pipeline.py#L22-L412">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py#L18-L410">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -82,7 +82,6 @@ ranking_pipeline.train_and_eval()
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Note that you may</h2></th></tr>
@@ -147,7 +146,6 @@ ranking_pipeline.train_and_eval()
 ```
 
 <!-- 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>
@@ -235,7 +233,7 @@ generated.
 
 <h3 id="train_and_eval"><code>train_and_eval</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py#L405-L412">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py#L403-L410">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">

docs/api_docs/python/tfr/extension/premade.md

@@ -63,7 +63,6 @@ for tf-ranking BERT.
 [`TensorLike`](../../tfr/keras/model/TensorLike.md)
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Other Members</h2></th></tr>

docs/api_docs/python/tfr/extension/premade/TFRBertConfig.md

@@ -76,6 +76,34 @@ Inherits From:
 </code></pre>
 
 <!-- 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>
+
+<tr>
+<td>
+`default_params`<a id="default_params"></a>
+</td>
+<td>
+a Python dict or another ParamsDict object including the
+default parameters to initialize.
+</td>
+</tr><tr>
+<td>
+`restrictions`<a id="restrictions"></a>
+</td>
+<td>
+a list of strings, which define a list of restrictions to
+ensure the consistency of different parameters internally. Each
+restriction string is defined as a binary relation with a set of
+operators, including {'==', '!=',  '<', '<=', '>', '>='}.
+</td>
+</tr>
+</table>
+
 <!-- Tabular view -->
 
  <table class="responsive fixed orange">
@@ -244,7 +272,6 @@ Makes the ParamsDict immutable.
 Override the ParamsDict with a set of given params.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -301,7 +328,6 @@ bb2: 20 ccc: a1: 1 a3: 3` one can define two restrictions like this ['a.a1 ==
 b.ccc.a1', 'a.a2 <= b.bb.bb2']
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">What it enforces are</th></tr>
@@ -315,7 +341,6 @@ b.ccc.a1', 'a.a2 <= b.bb.bb2']
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Raises</th></tr>
@@ -358,8 +383,9 @@ Implements the membership test operator.
 )
 </code></pre>
 
-<!-- Tabular view -->
+Return self==value.
 
+<!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Class Variables</h2></th></tr>

docs/api_docs/python/tfr/extension/premade/TFRBertDataConfig.md

@@ -85,7 +85,7 @@ Inherits From:
     default_params: dataclasses.InitVar[Optional[Mapping[str, Any]]] = None,
     restrictions: dataclasses.InitVar[Optional[List[str]]] = None,
     input_path: Union[Sequence[str], str, base_config.Config] = '',
-    tfds_name: str = '',
+    tfds_name: Union[str, base_config.Config] = '',
     tfds_split: str = '',
     global_batch_size: int = 0,
     is_training: bool = True,
@@ -120,6 +120,34 @@ Inherits From:
 </code></pre>
 
 <!-- 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>
+
+<tr>
+<td>
+`default_params`<a id="default_params"></a>
+</td>
+<td>
+a Python dict or another ParamsDict object including the
+default parameters to initialize.
+</td>
+</tr><tr>
+<td>
+`restrictions`<a id="restrictions"></a>
+</td>
+<td>
+a list of strings, which define a list of restrictions to
+ensure the consistency of different parameters internally. Each
+restriction string is defined as a binary relation with a set of
+operators, including {'==', '!=',  '<', '<=', '>', '>='}.
+</td>
+</tr>
+</table>
+
 <!-- Tabular view -->
 
  <table class="responsive fixed orange">
@@ -442,7 +470,6 @@ Makes the ParamsDict immutable.
 Override the ParamsDict with a set of given params.
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -499,7 +526,6 @@ bb2: 20 ccc: a1: 1 a3: 3` one can define two restrictions like this ['a.a1 ==
 b.ccc.a1', 'a.a2 <= b.bb.bb2']
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">What it enforces are</th></tr>
@@ -513,7 +539,6 @@ b.ccc.a1', 'a.a2 <= b.bb.bb2']
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Raises</th></tr>
@@ -556,8 +581,9 @@ Implements the membership test operator.
 )
 </code></pre>
 
-<!-- Tabular view -->
+Return self==value.
 
+<!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Class Variables</h2></th></tr>

docs/api_docs/python/tfr/extension/premade/TFRBertModelBuilder.md

@@ -48,8 +48,8 @@ Inherits From: [`ModelBuilder`](../../../tfr/keras/model/ModelBuilder.md),
 </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>
@@ -119,7 +119,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>