def yield_predictions_from_estimator(predictions, total, log_interval=10):
  """Yields predictions from Estimator.predict, with added error correction.

  This function handles the case of Estimator.predict occasionally restarting,
  causing results to be yielded out of order.

  Args:
    predictions: the return value of Estimator.predict. An iterable of dicts.
      Each dict MUST have a 'result_idx' attribute, used to track result order.
    total (int): total expected number of elements to yield from predictions.
    log_interval: log every this many seconds.

  Yields:
    the same dicts yielded from Estimator.predict, but in the right order. The
    result_idx element is removed from every dict.
  """
  predictions_iter = iter(predictions)
  total_yielded = 0

  start_time = time.time()
  last_log_timestamp = time.time()

  while total_yielded < total:
    try:
      result = next(predictions_iter)
    except StopIteration:
      raise ValueError(
          'Estimator.predict terminated before we got all results.')

    result_idx = result.pop('result_idx')

    if result_idx == total_yielded:
      # If results are always emitted from Estimator.predict in the same
      # order that they were fed into the Estimator, then we should always
      # expect result_idx to equal total_yielded. However, this does not always
      # happen, so we handle that in the `else` case below.
      yield result
      total_yielded += 1

      # Log progress.
      current_time = time.time()
      if current_time - last_log_timestamp > log_interval:
        total_time = current_time - start_time
        log_msg = 'Yielded {} results in {:.2f} secs.'.format(
            total_yielded, total_time)
        logging.info(log_msg)
        last_log_timestamp = current_time
    else:
      # If results start to arrive out of order, something has gone wrong.

      if result_idx < total_yielded:
        # This can happen if the TPU worker dies, causing Estimator.predict to
        # restart from the beginning. In this case, we just don't yield
        # anything on this step. Instead, we keep pulling things from the
        # iterator until we are back to where we were.
        if result_idx == 0:
          logging.warning('TPU worker seems to have restarted.')
      elif result_idx > total_yielded:
        # Something has gone really wrong.
        raise ValueError('Estimator.predict has somehow missed a result.')
Yields predictions from Estimator.predict, with added error correction.

  This function handles the case of Estimator.predict occasionally restarting,
  causing results to be yielded out of order.

  Args:
    predictions: the return value of Estimator.predict. An iterable of dicts.
      Each dict MUST have a 'result_idx' attribute, used to track result order.
    total (int): total expected number of elements to yield from predictions.
    log_interval: log every this many seconds.

  Yields:
    the same dicts yielded from Estimator.predict, but in the right order. The
    result_idx element is removed from every dict.
  
language/realm/retrieval.py