def predict(self, feature_batch):
    """Runs a prediction on a set of features.

    Calling multiple times does *not* regenerate the graph which makes predict
    much faster.

    Args:
      feature_batch: A list of list of features. IMPORTANT: If you're only
        classifying 1 thing, you still need to make it a batch of 1 by wrapping
        it in a list (i.e. predict([my_feature]), not predict(my_feature).

    Returns:
      results: A list of results from estimator.predict.
    """
    self.next_features = feature_batch
    if self.first_run:
      self.batch_size = len(feature_batch)
      self.predictions = self.estimator.predict(
          input_fn=self.input_fn(self._create_generator))
      self.first_run = False
    elif self.batch_size != len(feature_batch):
      raise ValueError("All batches must be of the same size. First-batch:" +
                       str(self.batch_size) + " This-batch:" +
                       str(len(feature_batch)))

    results = []
    for _ in range(self.batch_size):
      results.append(next(self.predictions))
    return results
Runs a prediction on a set of features.

    Calling multiple times does *not* regenerate the graph which makes predict
    much faster.

    Args:
      feature_batch: A list of list of features. IMPORTANT: If you're only
        classifying 1 thing, you still need to make it a batch of 1 by wrapping
        it in a list (i.e. predict([my_feature]), not predict(my_feature).

    Returns:
      results: A list of results from estimator.predict.
    
language/labs/drkit/bert_utils.py