def embed_qrys(qrys_tfrecord_filename, estimator, ckpt_path):
  """Compute query features for given TFRecords."""
  tf.logging.info("***** Embedding queries using %s *****", ckpt_path)

  predict_input_fn = input_fn_builder(
      input_file=qrys_tfrecord_filename,
      seq_length=FLAGS.max_seq_length,
      qry_length=FLAGS.max_query_length,
      ent_length=FLAGS.max_entity_length,
      is_training=False,
      drop_remainder=False)

  # If running eval on the TPU, you will need to specify the number of
  # steps.
  all_starts, all_ends, all_bows = [], [], []
  for result in estimator.predict(
      predict_input_fn, yield_single_examples=False, checkpoint_path=ckpt_path):
    if len(all_starts) % 1000 == 0:
      tf.logging.info("Processing example: %d", len(all_starts))
    all_starts += [result["qry_st_features"]]
    all_ends += [result["qry_en_features"]]
    all_bows += [result["qry_bow_features"]]

  return (np.concatenate(all_starts, axis=0), np.concatenate(all_ends, axis=0),
          np.concatenate(all_bows, axis=0))
Compute query features for given TFRecords.
language/labs/drkit/run_dualencoder_lsf.py