def embed(self, get_documents_dataset, total_docs, batch_size):
    """Embeds a Dataset of documents using Estimator.

    Args:
      get_documents_dataset: a function that returns a TF Dataset, where each
        element is a dict with attributes described below.
      total_docs: total number of documents in the Dataset.
      batch_size: number of documents to embed in each batch.  Each element in
        the Dataset returned by get_documents_dataset should be a dict with the
        attributes described below.

    get_documents_dataset should return a Dataset over dicts, each containing at
    least the following attributes:
    - title_token_ids: a 1-D int Tensor.
    - body_token_ids: a 1-D int Tensor.

    Yields:
      a [embed_dim] Numpy array, one for each document.
    """
    if total_docs < 1:
      raise ValueError('Must embed at least 1 document.')

    # These hyperparams are passed to Estimator.
    params = {
        'vocab_path':
            self._featurizer.tokenizer.vocab_path,
        'do_lower_case':
            self._featurizer.tokenizer.do_lower_case,
        'query_seq_len':
            self._featurizer.query_seq_len,
        'candidate_seq_len':
            self._featurizer.candidate_seq_len,
        'num_candidates':
            self._featurizer.num_candidates,
        'max_masks':
            self._featurizer.max_masks,
        'separate_candidate_segments':
            self._featurizer.separate_candidate_segments,
    }

    def input_fn(params):
      """Constructs the dataset fed to Estimator."""
      # We cannot access self._featurizer via closure, because this function is
      # passed to another device. Hence, we need to reconstruct the featurizer
      # from its hyerparameters (passed through `params`).

      tokenizer = featurization.Tokenizer(
          vocab_path=params['vocab_path'],
          do_lower_case=params['do_lower_case'])

      featurizer = featurization.Featurizer(
          query_seq_len=params['query_seq_len'],
          candidate_seq_len=params['candidate_seq_len'],
          num_candidates=params['num_candidates'],
          max_masks=params['max_masks'],
          tokenizer=tokenizer,
          separate_candidate_segments=params['separate_candidate_segments'])

      dataset = get_documents_dataset()

      def featurize(doc_dict):
        return featurizer.featurize_document_tf(doc_dict['title_token_ids'],
                                                doc_dict['body_token_ids'])

      dataset = dataset.map(
          featurize, num_parallel_calls=tf.data.experimental.AUTOTUNE)

      # Add a document index variable.
      dataset = dataset.enumerate()

      def _enumerate_to_dict(result_idx, tensor_dict):
        return dict(tensor_dict, result_idx=result_idx)

      dataset = dataset.map(
          _enumerate_to_dict, num_parallel_calls=tf.data.experimental.AUTOTUNE)

      # Pad the end of the dataset with one full extra batch.
      # This ensures that we don't drop the remainder.
      if total_docs % batch_size != 0:
        # Pad using the first value of the dataset, repeated batch_size times.
        pad_vals = dataset.take(1).repeat(batch_size)
        dataset = dataset.concatenate(pad_vals)

      # Batch the dataset.
      dataset = dataset.batch(batch_size, drop_remainder=True)
      dataset = dataset.prefetch(2)  # Prefetch for efficiency.
      return dataset

    def model_fn(features, labels, mode, params):
      """Constructs the model used by Estimator."""
      del labels, params
      embedder_module = hub.Module(
          spec=self._hub_module_spec, name='embedder', trainable=False)

      # Remove the result_idx before feeding features to the module.
      result_idx = features.pop('result_idx')

      # [batch_size, embed_dim]
      embeds = embedder_module(inputs=features, signature='projected')

      return tf_estimator.tpu.TPUEstimatorSpec(
          mode=mode, predictions={
              'embeds': embeds,
              'result_idx': result_idx
          })

    estimator = tf_estimator.tpu.TPUEstimator(
        use_tpu=self._use_tpu,
        model_fn=model_fn,
        model_dir=None,  # Don't persist model.
        config=self._run_config,
        params=params,
        train_batch_size=batch_size,
        predict_batch_size=batch_size)

    logging.info('Embedding %d documents total.', total_docs)
    predictions = estimator.predict(
        input_fn=input_fn, yield_single_examples=True)
    for result in yield_predictions_from_estimator(
        predictions, total=total_docs, log_interval=self._log_interval):
      yield result['embeds']
Embeds a Dataset of documents using Estimator.

    Args:
      get_documents_dataset: a function that returns a TF Dataset, where each
        element is a dict with attributes described below.
      total_docs: total number of documents in the Dataset.
      batch_size: number of documents to embed in each batch.  Each element in
        the Dataset returned by get_documents_dataset should be a dict with the
        attributes described below.

    get_documents_dataset should return a Dataset over dicts, each containing at
    least the following attributes:
    - title_token_ids: a 1-D int Tensor.
    - body_token_ids: a 1-D int Tensor.

    Yields:
      a [embed_dim] Numpy array, one for each document.
    
language/realm/retrieval.py