e93afea875a544bd4b887b9f093985dd4d23ec1e,official/nlp/tasks/tagging.py,,predict,#Any#Any#Any#,219

Before Change

  dataset = orbit.utils.make_distributed_dataset(tf.distribute.get_strategy(),
                                                 task.build_inputs, params)
  // Set `num_steps` to -1 to exhaust the dataset.
  predict_ids, sentence_ids = loop_fn(
      iter(dataset), num_steps=-1, state=([], []), reduce_fn=reduce_fn)
  return predict_ids, sentence_ids

After Change

  dataset = orbit.utils.make_distributed_dataset(tf.distribute.get_strategy(),
                                                 task.build_inputs, params)
  outputs = utils.predict(predict_step, aggregate_fn, dataset)
  return outputs["predict_ids"], outputs["sentence_ids"]

Instances