import tensorflow as tf

      with self._write_locks[index % rconst.NUM_FILE_SHARDS]:
        self._writers[index % rconst.NUM_FILE_SHARDS].write(example_bytes)

    else:
      if self._is_training:
        mask_start_index = data.pop(rconst.MASK_START_INDEX)
        batch_size = data[movielens.ITEM_COLUMN].shape[0]
        data[rconst.VALID_POINT_MASK] = np.less(np.arange(batch_size),
                                                mask_start_index)
        data = (data, data.pop("labels"))
      self._result_queue.put(data)

  def start_construction(self):
    if self._stream_files:
      tf.gfile.MakeDirs(self.current_data_root)
      template = os.path.join(self.current_data_root, rconst.SHARD_TEMPLATE)
      self._writers = [tf.io.TFRecordWriter(template.format(i))
                       for i in range(rconst.NUM_FILE_SHARDS)]

  def end_construction(self):
    if self._stream_files:
      [writer.close() for writer in self._writers]
      self._writers = []
      self._result_queue.put(self.current_data_root)

    self._epochs_completed += 1

  def data_generator(self, epochs_between_evals):
