import tensorflow as tf

    if FLAGS.auto_recover:
      hooks.append(tf.data.experimental.CheckpointInputPipelineHook(estimator))

    estimator.train(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps, hooks=hooks)

  if FLAGS.do_eval:
    tf.logging.info("***** Running evaluation *****")
    tf.logging.info("  Batch size = %d", FLAGS.eval_batch_size)

    eval_input_fn = input_fn_builder(
        input_files=validation_input_files,
        max_seq_length=FLAGS.max_seq_length,
