def dataset_from_preprocessors(preprocessing_servers,
                               rpcs_per_tf_op = 5):
  """Get dataset of serialized examples from a preprocessing job.

  Args:
    preprocessing_servers: BNS addresses for all preprocessing tasks.
    rpcs_per_tf_op: Number of RPCs to issue at once to a single server in a
      single TF op. Each TF op has some overhead, so usually a value like 20
      works well.

  Returns:
    Tensorflow dataset which yields single serialized examples.
  """

  def _ragged_where(cond, x):
    assert len(cond.shape) == 1
    assert len(x.shape) == 1
    return tf.gather(x, tf.squeeze(tf.where(cond), axis=1))

  def _pop_and_filter_valid(addresses):
    assert len(addresses.shape) == 1, "Expected replicated address tensor."
    minutes_in_ms = 60000
    result = _make_rpc_op(addresses, 10 * minutes_in_ms)
    is_ok = tf.math.equal(result.status_code, tf.errors.OK)
    return _ragged_where(is_ok, result.response)

  def _get_single_server_dataset(server):
    d = tf.data.Dataset.from_tensors(
        tf.tile(tf.expand_dims(server, 0), [rpcs_per_tf_op]))
    d = d.repeat()
    return d.map(_pop_and_filter_valid)

  dataset = tf.data.Dataset.from_tensor_slices(preprocessing_servers)
  dataset = dataset.apply(
      tf.data.experimental.parallel_interleave(
          _get_single_server_dataset,
          sloppy=True,
          cycle_length=len(preprocessing_servers),
      ))
  dataset = dataset.apply(tf.data.experimental.unbatch())
  return dataset
Get dataset of serialized examples from a preprocessing job.

  Args:
    preprocessing_servers: BNS addresses for all preprocessing tasks.
    rpcs_per_tf_op: Number of RPCs to issue at once to a single server in a
      single TF op. Each TF op has some overhead, so usually a value like 20
      works well.

  Returns:
    Tensorflow dataset which yields single serialized examples.
  
language/realm/preprocessing.py