def embed_documents(
    shard_paths,
    shard_sizes,
    hub_module_spec,
    featurizer,
    batch_size,
    tpu_run_config = None):
  """Embeds documents either locally (CPU/GPU) or with a TPU worker.

  Note: TPUEstimator.predict currently requires the TPU worker to have a single
  "host" (a machine running TensorFlow that is physically connected to the TPU
  chips). This is not true for all TPU topologies -- some have multiple hosts.

  Args:
    shard_paths: a list of file paths, each specifying a GZIP'd TFRecord file
      containing documents stored as TF Examples. Doc embeddings will be
      concatenated in the order of shard_paths.
    shard_sizes: a list parallel to shard_paths, specifying the number of
      documents in each shard.
    hub_module_spec: path to the Hub module that will be used to embed the
      documents.
    featurizer: a Featurizer used to convert documents into Tensor features.
    batch_size: embed documents in batches of this size.
    tpu_run_config: configures the TPU worker. If None, run on CPU/GPU.

  Returns:
    a [total_docs, embed_dim] Numpy array.
  """
  embedder = DocumentEmbedder(
      hub_module_spec=hub_module_spec,
      featurizer=featurizer,
      use_tpu=(tpu_run_config is not None),
      run_config=tpu_run_config)

  def parse_tf_example(serialized):
    # FixedLenSequenceFeature requires allow_missing to be True, even though we
    # can't actually handle those cases.
    feature_spec = {
        'title':
            tf.FixedLenSequenceFeature([], tf.string, allow_missing=True),
        'text':
            tf.FixedLenSequenceFeature([], tf.string, allow_missing=True),
        'title_token_ids':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
        'body_token_ids':
            tf.FixedLenSequenceFeature([], tf.int64, allow_missing=True),
    }
    features = tf.parse_single_example(serialized, feature_spec)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(features.keys()):
      tensor = features[name]
      if tensor.dtype == tf.int64:
        tensor = tf.cast(tensor, tf.int32)
      features[name] = tensor

    return features

  def get_documents_dataset():
    # Note: num_parallel_reads should be None to guarantee that shard_paths
    # are visited sequentially, not in parallel.
    dataset = tf.data.TFRecordDataset(
        shard_paths,
        compression_type='GZIP',
        buffer_size=8 * 1024 * 1024,
        num_parallel_reads=None)

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

  embeds = embedder.embed(
      get_documents_dataset=get_documents_dataset,
      total_docs=sum(shard_sizes),
      batch_size=batch_size)
  # A list of [embed_dim] Numpy arrays.
  embeds_list = list(embeds)

  # A [total_docs, embed_dim] Numpy array.
  return np.stack(embeds_list, axis=0)
Embeds documents either locally (CPU/GPU) or with a TPU worker.

  Note: TPUEstimator.predict currently requires the TPU worker to have a single
  "host" (a machine running TensorFlow that is physically connected to the TPU
  chips). This is not true for all TPU topologies -- some have multiple hosts.

  Args:
    shard_paths: a list of file paths, each specifying a GZIP'd TFRecord file
      containing documents stored as TF Examples. Doc embeddings will be
      concatenated in the order of shard_paths.
    shard_sizes: a list parallel to shard_paths, specifying the number of
      documents in each shard.
    hub_module_spec: path to the Hub module that will be used to embed the
      documents.
    featurizer: a Featurizer used to convert documents into Tensor features.
    batch_size: embed documents in batches of this size.
    tpu_run_config: configures the TPU worker. If None, run on CPU/GPU.

  Returns:
    a [total_docs, embed_dim] Numpy array.
  
language/realm/retrieval.py