def input_fn(self,
               mode,
               hparams,
               data_dir=None,
               params=None,
               config=None,
               force_repeat=False,
               prevent_repeat=False,
               dataset_kwargs=None):
    """Builds input pipeline for problem.

    Args:
      mode: tf.estimator.ModeKeys
      hparams: HParams, model hparams
      data_dir: str, data directory; if None, will use hparams.data_dir
      params: dict, may include "batch_size"
      config: RunConfig; should have the data_parallelism attribute if not using
        TPU
      force_repeat: bool, whether to repeat the data even if not training
      prevent_repeat: bool, whether to not repeat when in training mode.
        Overrides force_repeat.
      dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
        method when called

    Returns:
      (features_dict<str name, Tensor feature>, Tensor targets)
    """
    partition_id, num_partitions = self._dataset_partition(mode, config, params)
    is_training = mode == tf.estimator.ModeKeys.TRAIN
    if config and config.use_tpu:
      num_threads = 64
    else:
      num_threads = data_reader.cpu_count() if is_training else 1
    data_dir = data_dir or (hasattr(hparams, "data_dir") and hparams.data_dir)
    dataset_kwargs = dataset_kwargs or {}
    dataset_kwargs.update({
        "mode": mode,
        "data_dir": data_dir,
        "num_threads": num_threads,
        "hparams": hparams,
        "partition_id": partition_id,
        "num_partitions": num_partitions,
    })
    return data_reader.input_fn(
        self.dataset(**dataset_kwargs),
        self.filepattern(data_dir, mode),
        self.skip_random_fraction_when_training,
        self.batch_size_means_tokens,
        self.get_hparams().batch_size_multiplier,
        self.max_length(hparams),
        mode,
        hparams,
        data_dir=data_dir,
        params=params,
        config=config,
        force_repeat=force_repeat,
        prevent_repeat=prevent_repeat)