Revert usage of Datasets API

PiperOrigin-RevId: 163421122

tensor2tensor/utils/data_reader.py

@@ -20,7 +20,6 @@
 
 import math
 import os
-import random
 
 # Dependency imports
 
@@ -114,17 +113,18 @@ def decode_record(record):
     return dict(zip(decode_items, decoded))
 
   with tf.name_scope("examples_in"):
+    # Read serialized examples using slim parallel_reader.
     data_files = tf.contrib.slim.parallel_reader.get_data_files(data_sources)
-    if training:
-      random.shuffle(data_files)
-    dataset = tf.contrib.data.TFRecordDataset(data_files)
     num_readers = min(4 if training else 1, len(data_files))
-    dataset = dataset.map(decode_record, num_threads=num_readers)
-    if training:
-      dataset = dataset.shuffle(capacity)
-    dataset = dataset.repeat(None if training else 1)
-    it = dataset.make_one_shot_iterator()
-    return it.get_next()
+    _, example_serialized = tf.contrib.slim.parallel_reader.parallel_read(
+        data_sources,
+        tf.TFRecordReader,
+        num_epochs=None if training else 1,
+        shuffle=training,
+        capacity=2 * capacity,
+        min_after_dequeue=capacity,
+        num_readers=num_readers)
+    return decode_record(example_serialized)
 
 
 def preprocessing(examples, data_file_pattern, mode):