easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/input/parquet_input_v3.py

@@ -0,0 +1,203 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+
+from easy_rec.python.input.input import Input
+from easy_rec.python.utils.input_utils import get_type_defaults
+
+try:
+  from tensorflow.python.data.experimental.ops import parquet_dataset_ops
+  from tensorflow.python.data.experimental.ops import parquet_pybind
+  from tensorflow.python.data.experimental.ops import dataframe
+  from tensorflow.python.ops import gen_ragged_conversion_ops
+  from tensorflow.python.ops.work_queue import WorkQueue
+  _has_deep_rec = True
+except Exception:
+  _has_deep_rec = False
+  pass
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class ParquetInputV3(Input):
+
+  def __init__(self,
+               data_config,
+               feature_config,
+               input_path,
+               task_index=0,
+               task_num=1,
+               check_mode=False,
+               pipeline_config=None,
+               **kwargs):
+    if not _has_deep_rec:
+      raise RuntimeError('You should install DeepRec first.')
+    super(ParquetInputV3,
+          self).__init__(data_config, feature_config, input_path, task_index,
+                         task_num, check_mode, pipeline_config)
+
+    self._ignore_val_dict = {}
+    for f in data_config.input_fields:
+      if f.HasField('ignore_val'):
+        self._ignore_val_dict[f.input_name] = get_type_defaults(
+            f.input_type, f.ignore_val)
+
+    self._true_type_dict = {}
+    for fc in self._feature_configs:
+      if fc.feature_type in [fc.IdFeature, fc.TagFeature, fc.SequenceFeature]:
+        if fc.hash_bucket_size > 0 or len(
+            fc.vocab_list) > 0 or fc.HasField('vocab_file'):
+          self._true_type_dict[fc.input_names[0]] = tf.string
+        else:
+          self._true_type_dict[fc.input_names[0]] = tf.int64
+        if len(fc.input_names) > 1:
+          self._true_type_dict[fc.input_names[1]] = tf.float32
+      if fc.feature_type == fc.RawFeature:
+        self._true_type_dict[fc.input_names[0]] = tf.float32
+
+    self._reserve_fields = None
+    self._reserve_types = None
+    if 'reserve_fields' in kwargs and 'reserve_types' in kwargs:
+      self._reserve_fields = kwargs['reserve_fields']
+      self._reserve_types = kwargs['reserve_types']
+
+    # In ParquetDataset multi_value use input type
+    self._multi_value_types = {}
+
+  def _ignore_and_cast(self, name, value):
+    ignore_value = self._ignore_val_dict.get(name, None)
+    if ignore_value:
+      if isinstance(value, tf.SparseTensor):
+        indices = tf.where(tf.equal(value.values, ignore_value))
+        value = tf.SparseTensor(
+            tf.gather_nd(value.indices, indices),
+            tf.gather_nd(value.values, indices), value.dense_shape)
+      elif isinstance(value, tf.Tensor):
+        indices = tf.where(tf.not_equal(value, ignore_value), name='indices')
+        value = tf.SparseTensor(
+            indices=indices,
+            values=tf.gather_nd(value, indices),
+            dense_shape=tf.shape(value, out_type=tf.int64))
+    dtype = self._true_type_dict.get(name, None)
+    if dtype:
+      value = tf.cast(value, dtype)
+    return value
+
+  def _parse_dataframe_value(self, value):
+    if len(value.nested_row_splits) == 0:
+      return value.values
+    value.values.set_shape([None])
+    sparse_value = gen_ragged_conversion_ops.ragged_tensor_to_sparse(
+        value.nested_row_splits, value.values)
+    return tf.SparseTensor(sparse_value.sparse_indices,
+                           sparse_value.sparse_values,
+                           sparse_value.sparse_dense_shape)
+
+  def _parse_dataframe(self, df):
+    inputs = {}
+    for k, v in df.items():
+      if k in self._effective_fields:
+        if isinstance(v, dataframe.DataFrame.Value):
+          v = self._parse_dataframe_value(v)
+      elif k in self._label_fields:
+        if isinstance(v, dataframe.DataFrame.Value):
+          v = v.values
+      elif k in self._reserve_fields:
+        if isinstance(v, dataframe.DataFrame.Value):
+          v = v.values
+      else:
+        continue
+      inputs[k] = v
+    return inputs
+
+  def _build(self, mode, params):
+    input_files = []
+    for sub_path in self._input_path.strip().split(','):
+      input_files.extend(tf.gfile.Glob(sub_path))
+    file_num = len(input_files)
+    logging.info('[task_index=%d] total_file_num=%d task_num=%d' %
+                 (self._task_index, file_num, self._task_num))
+
+    task_index = self._task_index
+    task_num = self._task_num
+    if self._data_config.chief_redundant:
+      task_index = max(self._task_index - 1, 0)
+      task_num = max(self._task_num - 1, 1)
+
+    if self._data_config.pai_worker_queue and \
+        mode == tf.estimator.ModeKeys.TRAIN:
+      work_queue = WorkQueue(
+          input_files,
+          num_epochs=self.num_epochs,
+          shuffle=self._data_config.shuffle)
+      my_files = work_queue.input_dataset()
+    else:
+      my_files = []
+      for file_id in range(file_num):
+        if (file_id % task_num) == task_index:
+          my_files.append(input_files[file_id])
+
+    parquet_fields = parquet_pybind.parquet_fields(input_files[0])
+    parquet_input_fields = []
+    for f in parquet_fields:
+      if f.name in self._input_fields:
+        parquet_input_fields.append(f)
+
+    all_fields = set(self._effective_fields)
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      all_fields |= set(self._label_fields)
+    if self._reserve_fields:
+      all_fields |= set(self._reserve_fields)
+
+    selected_fields = []
+    for f in parquet_input_fields:
+      if f.name in all_fields:
+        selected_fields.append(f)
+
+    num_parallel_reads = min(self._data_config.num_parallel_calls,
+                             len(input_files) // task_num)
+    dataset = parquet_dataset_ops.ParquetDataset(
+        my_files,
+        batch_size=self._batch_size,
+        fields=selected_fields,
+        drop_remainder=self._data_config.drop_remainder,
+        num_parallel_reads=num_parallel_reads)
+    # partition_count=task_num,
+    # partition_index=task_index)
+
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      if self._data_config.shuffle:
+        dataset = dataset.shuffle(
+            self._data_config.shuffle_buffer_size,
+            seed=2020,
+            reshuffle_each_iteration=True)
+      dataset = dataset.repeat(self.num_epochs)
+    else:
+      dataset = dataset.repeat(1)
+
+    dataset = dataset.map(
+        self._parse_dataframe,
+        num_parallel_calls=self._data_config.num_parallel_calls)
+
+    # preprocess is necessary to transform data
+    # so that they could be feed into FeatureColumns
+    dataset = dataset.map(
+        map_func=self._preprocess,
+        num_parallel_calls=self._data_config.num_parallel_calls)
+
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      dataset = dataset.map(lambda x:
+                            (self._get_features(x), self._get_labels(x)))
+    else:
+      dataset = dataset.map(lambda x: (self._get_features(x)))
+    return dataset
+
+  def _preprocess(self, field_dict):
+    for k, v in field_dict.items():
+      field_dict[k] = self._ignore_and_cast(k, v)
+    return super(ParquetInputV3, self)._preprocess(field_dict)

easy_rec/python/input/rtp_input.py

@@ -0,0 +1,225 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+
+from easy_rec.python.input.input import Input
+from easy_rec.python.ops.gen_str_avx_op import str_split_by_chr
+from easy_rec.python.utils.check_utils import check_split
+from easy_rec.python.utils.check_utils import check_string_to_number
+from easy_rec.python.utils.input_utils import string_to_number
+from easy_rec.python.utils.tf_utils import get_tf_type
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class RTPInput(Input):
+  """RTPInput for parsing rtp fg new input format.
+
+  Our new format(csv in csv) of rtp output:
+     label0, item_id, ..., user_id, features
+  here the separator(,) could be specified by data_config.rtp_separator
+  For the feature column, features are separated by ,
+     multiple values of one feature are separated by , such as:
+     ...20beautysmartParis...
+  The features column and labels are specified by data_config.selected_cols,
+     columns are selected by indices as our csv file has no header,
+     such as: 0,1,4, means the 4th column is features, the 1st and 2nd
+     columns are labels
+  """
+
+  def __init__(self,
+               data_config,
+               feature_config,
+               input_path,
+               task_index=0,
+               task_num=1,
+               check_mode=False,
+               pipeline_config=None):
+    super(RTPInput,
+          self).__init__(data_config, feature_config, input_path, task_index,
+                         task_num, check_mode, pipeline_config)
+    logging.info('input_fields: %s label_fields: %s' %
+                 (','.join(self._input_fields), ','.join(self._label_fields)))
+    self._rtp_separator = self._data_config.rtp_separator
+    if not isinstance(self._rtp_separator, str):
+      self._rtp_separator = self._rtp_separator.encode('utf-8')
+    self._selected_cols = [
+        int(x) for x in self._data_config.selected_cols.split(',')
+    ]
+    self._num_cols = -1
+    self._feature_col_id = self._selected_cols[-1]
+    logging.info('rtp separator = %s' % self._rtp_separator)
+
+  def _parse_csv(self, line):
+    record_defaults = ['' for i in range(self._num_cols)]
+
+    # the actual features are in one single column
+    record_defaults[self._feature_col_id] = self._data_config.separator.join([
+        str(self.get_type_defaults(t, v))
+        for x, t, v in zip(self._input_fields, self._input_field_types,
+                           self._input_field_defaults)
+        if x not in self._label_fields
+    ])
+
+    check_list = [
+        tf.py_func(
+            check_split, [line, self._rtp_separator,
+                          len(record_defaults)],
+            Tout=tf.bool)
+    ] if self._check_mode else []
+    with tf.control_dependencies(check_list):
+      fields = tf.string_split(line, self._rtp_separator, skip_empty=False)
+
+    fields = tf.reshape(fields.values, [-1, len(record_defaults)])
+
+    labels = []
+    for idx, x in enumerate(self._selected_cols[:-1]):
+      field = fields[:, x]
+      fname = self._input_fields[idx]
+      ftype = self._input_field_types[idx]
+      tf_type = get_tf_type(ftype)
+      if field.dtype in [tf.string]:
+        check_list = [
+            tf.py_func(check_string_to_number, [field, fname], Tout=tf.bool)
+        ] if self._check_mode else []
+        with tf.control_dependencies(check_list):
+          field = tf.string_to_number(field, tf_type)
+      labels.append(field)
+
+    # only for features, labels excluded
+    record_types = [
+        t for x, t in zip(self._input_fields, self._input_field_types)
+        if x not in self._label_fields
+    ]
+    # assume that the last field is the generated feature column
+    print('field_delim = %s' % self._data_config.separator)
+    feature_str = fields[:, self._feature_col_id]
+    check_list = [
+        tf.py_func(
+            check_split,
+            [feature_str, self._data_config.separator,
+             len(record_types)],
+            Tout=tf.bool)
+    ] if self._check_mode else []
+    with tf.control_dependencies(check_list):
+      fields = str_split_by_chr(
+          feature_str, self._data_config.separator, skip_empty=False)
+    tmp_fields = tf.reshape(fields.values, [-1, len(record_types)])
+    rtp_record_defaults = [
+        str(self.get_type_defaults(t, v))
+        for x, t, v in zip(self._input_fields, self._input_field_types,
+                           self._input_field_defaults)
+        if x not in self._label_fields
+    ]
+    fields = []
+    for i in range(len(record_types)):
+      field = string_to_number(tmp_fields[:, i], record_types[i],
+                               rtp_record_defaults[i], i)
+      fields.append(field)
+
+    field_keys = [x for x in self._input_fields if x not in self._label_fields]
+    effective_fids = [field_keys.index(x) for x in self._effective_fields]
+    inputs = {field_keys[x]: fields[x] for x in effective_fids}
+
+    for x in range(len(self._label_fields)):
+      inputs[self._label_fields[x]] = labels[x]
+    return inputs
+
+  def _build(self, mode, params):
+    if type(self._input_path) != list:
+      self._input_path = self._input_path.split(',')
+    file_paths = []
+    for x in self._input_path:
+      file_paths.extend(tf.gfile.Glob(x))
+    assert len(file_paths) > 0, 'match no files with %s' % self._input_path
+
+    # try to figure out number of fields from one file
+    with tf.gfile.GFile(file_paths[0], 'r') as fin:
+      num_lines = 0
+      for line_str in fin:
+        line_tok = line_str.strip().split(self._rtp_separator)
+        if self._num_cols != -1:
+          assert self._num_cols == len(line_tok), \
+              'num selected cols is %d, not equal to %d, current line is: %s, please check rtp_separator and data.' % \
+              (self._num_cols, len(line_tok), line_str)
+        self._num_cols = len(line_tok)
+        num_lines += 1
+        if num_lines > 10:
+          break
+    logging.info('num selected cols = %d' % self._num_cols)
+
+    record_defaults = [
+        self.get_type_defaults(t, v)
+        for x, t, v in zip(self._input_fields, self._input_field_types,
+                           self._input_field_defaults)
+        if x in self._label_fields
+    ]
+
+    # the features are in one single column
+    record_defaults.append(
+        self._data_config.separator.join([
+            str(self.get_type_defaults(t, v))
+            for x, t, v in zip(self._input_fields, self._input_field_types,
+                               self._input_field_defaults)
+            if x not in self._label_fields
+        ]))
+
+    num_parallel_calls = self._data_config.num_parallel_calls
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      logging.info('train files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.Dataset.from_tensor_slices(file_paths)
+
+      if self._data_config.file_shard:
+        dataset = self._safe_shard(dataset)
+
+      if self._data_config.shuffle:
+        # shuffle input files
+        dataset = dataset.shuffle(len(file_paths))
+
+      # too many readers read the same file will cause performance issues
+      # as the same data will be read multiple times
+      parallel_num = min(num_parallel_calls, len(file_paths))
+      dataset = dataset.interleave(
+          tf.data.TextLineDataset,
+          cycle_length=parallel_num,
+          num_parallel_calls=parallel_num)
+
+      if not self._data_config.file_shard:
+        dataset = self._safe_shard(dataset)
+
+      if self._data_config.shuffle:
+        dataset = dataset.shuffle(
+            self._data_config.shuffle_buffer_size,
+            seed=2020,
+            reshuffle_each_iteration=True)
+      dataset = dataset.repeat(self.num_epochs)
+    else:
+      logging.info('eval files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.TextLineDataset(file_paths)
+      dataset = dataset.repeat(1)
+
+    dataset = dataset.batch(batch_size=self._data_config.batch_size)
+
+    dataset = dataset.map(
+        self._parse_csv,
+        num_parallel_calls=self._data_config.num_parallel_calls)
+
+    # preprocess is necessary to transform data
+    # so that they could be feed into FeatureColumns
+    dataset = dataset.map(
+        map_func=self._preprocess,
+        num_parallel_calls=self._data_config.num_parallel_calls)
+
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      dataset = dataset.map(lambda x:
+                            (self._get_features(x), self._get_labels(x)))
+    else:
+      dataset = dataset.map(lambda x: (self._get_features(x)))
+    return dataset

easy_rec/python/input/rtp_input_v2.py

@@ -0,0 +1,145 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+
+from easy_rec.python.input.input import Input
+from easy_rec.python.protos.dataset_pb2 import DatasetConfig
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class RTPInputV2(Input):
+  """RTPInput for parsing rtp fg input format.
+
+  the original rtp format, it is not efficient for training, the performance have to be tuned.
+  """
+
+  def __init__(self,
+               data_config,
+               feature_config,
+               input_path,
+               task_index=0,
+               task_num=1,
+               check_mode=False,
+               pipeline_config=None):
+    super(RTPInputV2,
+          self).__init__(data_config, feature_config, input_path, task_index,
+                         task_num, check_mode, pipeline_config)
+
+  def _parse_rtp(self, lines):
+    tf_types = [tf.string for x in self._input_field_types]
+
+    def _parse_one_line_tf(line):
+      line = tf.expand_dims(line, axis=0)
+      field_toks = tf.string_split(line, '\002').values
+      field_vals = tf.string_split(field_toks, '\003').values
+      field_vals = tf.reshape(field_vals, [-1, 2])
+      keys = field_vals[:, 0]
+      vals = field_vals[:, 1]
+      temp_vals = [
+          str(
+              self.get_type_defaults(self._input_field_types[i],
+                                     self._input_field_defaults[i]))
+          for i in range(len(self._input_fields))
+      ]
+      for i, key in enumerate(self._input_fields):
+        msk = tf.equal(key, keys)
+        val = tf.boolean_mask(vals, msk)
+        def_val = self.get_type_defaults(self._input_field_types[i],
+                                         self._input_field_defaults[i])
+        temp_vals[i] = tf.cond(
+            tf.reduce_any(msk), lambda: tf.reduce_join(val, separator=','),
+            lambda: tf.constant(str(def_val)))
+      return temp_vals
+
+    fields = tf.map_fn(
+        _parse_one_line_tf,
+        lines,
+        tf_types,
+        parallel_iterations=64,
+        name='parse_one_line_tf_map_fn')
+
+    def _convert(x, target_type, name):
+      if target_type in [DatasetConfig.FLOAT, DatasetConfig.DOUBLE]:
+        return tf.string_to_number(
+            x, tf.float32, name='convert_input_flt32/%s' % name)
+      elif target_type == DatasetConfig.INT32:
+        return tf.string_to_number(
+            x, tf.int32, name='convert_input_int32/%s' % name)
+      elif target_type == DatasetConfig.INT64:
+        return tf.string_to_number(
+            x, tf.int64, name='convert_input_int64/%s' % name)
+      return x
+
+    inputs = {
+        self._input_fields[x]: _convert(fields[x], self._input_field_types[x],
+                                        self._input_fields[x])
+        for x in self._effective_fids
+    }
+
+    for x in self._label_fids:
+      inputs[self._input_fields[x]] = fields[x]
+    return inputs
+
+  def _build(self, mode, params):
+    if type(self._input_path) != list:
+      self._input_path = self._input_path.split(',')
+    file_paths = []
+    for x in self._input_path:
+      file_paths.extend(tf.gfile.Glob(x))
+    assert len(file_paths) > 0, 'match no files with %s' % self._input_path
+
+    num_parallel_calls = self._data_config.num_parallel_calls
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      logging.info('train files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.Dataset.from_tensor_slices(file_paths)
+
+      if self._data_config.file_shard:
+        dataset = self._safe_shard(dataset)
+
+      if self._data_config.shuffle:
+        # shuffle input files
+        dataset = dataset.shuffle(len(file_paths))
+
+      # too many readers read the same file will cause performance issues
+      # as the same data will be read multiple times
+      parallel_num = min(num_parallel_calls, len(file_paths))
+      dataset = dataset.interleave(
+          tf.data.TextLineDataset,
+          cycle_length=parallel_num,
+          num_parallel_calls=parallel_num)
+
+      if not self._data_config.file_shard:
+        dataset = self._safe_shard(dataset)
+
+      if self._data_config.shuffle:
+        dataset = dataset.shuffle(
+            self._data_config.shuffle_buffer_size,
+            seed=2020,
+            reshuffle_each_iteration=True)
+      dataset = dataset.repeat(self.num_epochs)
+    else:
+      logging.info('eval files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.TextLineDataset(file_paths)
+      dataset = dataset.repeat(1)
+
+    dataset = dataset.batch(self._data_config.batch_size)
+    dataset = dataset.map(
+        self._parse_rtp, num_parallel_calls=num_parallel_calls)
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+    dataset = dataset.map(
+        map_func=self._preprocess, num_parallel_calls=num_parallel_calls)
+
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      dataset = dataset.map(lambda x:
+                            (self._get_features(x), self._get_labels(x)))
+    else:
+      dataset = dataset.map(lambda x: (self._get_features(x)))
+    return dataset

easy_rec/python/input/tfrecord_input.py

@@ -0,0 +1,100 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+
+from easy_rec.python.input.input import Input
+from easy_rec.python.utils.tf_utils import get_tf_type
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class TFRecordInput(Input):
+
+  def __init__(self,
+               data_config,
+               feature_config,
+               input_path,
+               task_index=0,
+               task_num=1,
+               check_mode=False,
+               pipeline_config=None):
+    super(TFRecordInput,
+          self).__init__(data_config, feature_config, input_path, task_index,
+                         task_num, check_mode, pipeline_config)
+
+    self.feature_desc = {}
+    for x, t, d, s in zip(self._input_fields, self._input_field_types,
+                          self._input_field_defaults, self._input_dims):
+      d = self.get_type_defaults(t, d)
+      t = get_tf_type(t)
+      if s == 1:
+        self.feature_desc[x] = tf.FixedLenFeature(
+            dtype=t, shape=[s], default_value=d)
+      else:
+        self.feature_desc[x] = tf.FixedLenFeature(
+            dtype=t, shape=[s], default_value=[d] * s)
+
+  def _parse_tfrecord(self, example):
+    try:
+      inputs = tf.parse_single_example(example, features=self.feature_desc)
+    except AttributeError:
+      inputs = tf.io.parse_single_example(example, features=self.feature_desc)
+    return inputs
+
+  def _build(self, mode, params):
+    if type(self._input_path) != list:
+      self._input_path = self._input_path.split(',')
+    file_paths = []
+    for x in self._input_path:
+      file_paths.extend(tf.gfile.Glob(x))
+    assert len(file_paths) > 0, 'match no files with %s' % self._input_path
+
+    num_parallel_calls = self._data_config.num_parallel_calls
+    data_compression_type = self._data_config.data_compression_type
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      logging.info('train files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.Dataset.from_tensor_slices(file_paths)
+      if self._data_config.shuffle:
+        # shuffle input files
+        dataset = dataset.shuffle(len(file_paths))
+      # too many readers read the same file will cause performance issues
+      # as the same data will be read multiple times
+      parallel_num = min(num_parallel_calls, len(file_paths))
+      dataset = dataset.interleave(
+          lambda x: tf.data.TFRecordDataset(
+              x, compression_type=data_compression_type),
+          cycle_length=parallel_num,
+          num_parallel_calls=parallel_num)
+      dataset = dataset.shard(self._task_num, self._task_index)
+      if self._data_config.shuffle:
+        dataset = dataset.shuffle(
+            self._data_config.shuffle_buffer_size,
+            seed=2020,
+            reshuffle_each_iteration=True)
+      dataset = dataset.repeat(self.num_epochs)
+    else:
+      logging.info('eval files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.TFRecordDataset(
+          file_paths, compression_type=data_compression_type)
+      dataset = dataset.repeat(1)
+
+    dataset = dataset.map(
+        self._parse_tfrecord, num_parallel_calls=num_parallel_calls)
+    dataset = dataset.batch(self._data_config.batch_size)
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+    dataset = dataset.map(
+        map_func=self._preprocess, num_parallel_calls=num_parallel_calls)
+
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      dataset = dataset.map(lambda x:
+                            (self._get_features(x), self._get_labels(x)))
+    else:
+      dataset = dataset.map(lambda x: (self._get_features(x)))
+    return dataset