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
          })
Constructs the model used by Estimator.
language/realm/retrieval.py