def embed_docs(docs_tfrecord_filename, estimator, ckpt_path, tok_index_to_keep,
               num_tokens, dimension):
  """Compute document features for given TFRecords."""
  tf.logging.info("***** Embedding documents using %s *****", ckpt_path)

  predict_input_fn = input_fn_builder(
      input_file=docs_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_features = np.empty((num_tokens, dimension), dtype=np.float32)
  doc_spans = []
  batch_start, num_features = 0, 0
  for result in estimator.predict(
      predict_input_fn, yield_single_examples=False, checkpoint_path=ckpt_path):
    if batch_start % 1000 == 0:
      tf.logging.info("Processing batch: %d", batch_start)
    for ii in range(result["doc_features"].shape[0]):
      doc_idx = batch_start + ii
      st = num_features
      for ti in tok_index_to_keep[doc_idx]:
        all_features[num_features, :] = result["doc_features"][ii, ti, :]
        num_features += 1
      doc_spans.append([st, num_features])
    batch_start += result["doc_features"].shape[0]

  return all_features, np.stack(doc_spans, axis=0)
Compute document features for given TFRecords.
language/labs/drkit/run_dualencoder_lsf.py