from tensorflow.contrib.learn.python.learn.io import data_feeder


def _get_input_fn(x, y, batch_size):
  # TODO(ipoloshukin): Remove this when refactor of data_feeder is done
  if hasattr(x, 'create_graph') and hasattr(y, 'create_graph'):
    def input_fn():
      return x.create_graph(), y.create_graph()
    return input_fn, None

  df = data_feeder.setup_train_data_feeder(x, y,
                                           n_classes=None,
                                           batch_size=batch_size)
  return df.input_builder, df.get_feed_dict_fn()


def _get_predict_input_fn(x, batch_size):
  # TODO(ipoloshukin): Remove this when refactor of data_feeder is done
