Replace deprecated tf.contrib.data.map_and_batch(...) with `tf.data…

….experimental.map_and_batch(...)`. PiperOrigin-RevId: 523833192
google-research · Apr 12, 2023 · b772393 · b772393
1 parent 68286fb
commit b772393
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/classifier_utils.py b/classifier_utils.py
@@ -26,7 +26,6 @@
 from albert import tokenization
 import tensorflow.compat.v1 as tf
 from tensorflow.compat.v1 import estimator as tf_estimator
-from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import metrics as contrib_metrics
 from tensorflow.contrib import tpu as contrib_tpu
 
@@ -739,7 +738,7 @@ def input_fn(params):
       d = d.shuffle(buffer_size=100)
 
     d = d.apply(
-        contrib_data.map_and_batch(
+        tf.data.experimental.map_and_batch(
             lambda record: _decode_record(record, name_to_features),
             batch_size=batch_size,
             drop_remainder=drop_remainder))

diff --git a/squad_utils.py b/squad_utils.py
@@ -33,7 +33,6 @@
 from six.moves import range
 import tensorflow.compat.v1 as tf
 from tensorflow.compat.v1 import estimator as tf_estimator
-from tensorflow.contrib import data as contrib_data
 from tensorflow.contrib import layers as contrib_layers
 from tensorflow.contrib import tpu as contrib_tpu
 
@@ -696,7 +695,7 @@ def input_fn(params):
       d = d.shuffle(buffer_size=100)
 
     d = d.apply(
-        contrib_data.map_and_batch(
+        tf.data.experimental.map_and_batch(
             lambda record: _decode_record(record, name_to_features),
             batch_size=batch_size,
             drop_remainder=drop_remainder))