def input_fn(params, is_train):
  """An input function satisfying the tf.estimator API."""
  dataset = ict_dataset.get_dataset(
      examples_path=params["examples_path"],
      mask_rate=params["mask_rate"],
      bert_hub_module_path=params["bert_hub_module_path"],
      query_seq_len=params["query_seq_len"],
      block_seq_len=params["block_seq_len"],
      num_block_records=params["num_block_records"],
      num_input_threads=params["num_input_threads"])
  batch_size = params["batch_size"] if is_train else params["eval_batch_size"]
  dataset = dataset.batch(batch_size, drop_remainder=True)
  return dataset
An input function satisfying the tf.estimator API.
language/orqa/models/ict_model.py