Multiple beam output candidates not handled by serving_utils

[walmsley]

Working with serving multiple beam search outputs, and bumped into an issue. I've actually solved the issue, but need to figure out how to submit the change without breaking for other people.

The code at: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/serving/serving_utils.py#L119

outputs = [
      (_decode(prediction["outputs"], output_decoder),
       prediction["scores"])
      for prediction in predictions
  ]

assumes that beam search returns only a single candidate; but I've needed to modify it to the following to handle multiple outputs:

outputs = [
      [(_decode(output, output_decoder),
       prediction["scores"][i]) for i,output in enumerate(prediction["outputs"])]
      for prediction in predictions
  ]

The question is: how to design this to gracefully handle the multi-beam-output case while also not changing the output for single-output cases?

[walmsley]

If this issue remains unaddressed, people who only need to serve their model with multiple beam search candidates as output can just use the code pasted above.

[lkluo]

I do not quite follow this. You can pick up the first beam search result from multiple beams once you need only one translation.