Modify serving utils (#1495)

* Add decode logic for model with return_beams=True.

* Add print logic for model with return_beams=True.

Author: Drunkar

tensor2tensor/serving/query.py

@@ -90,14 +90,26 @@ def main(_):
     outputs = serving_utils.predict([inputs], problem, request_fn)
     outputs, = outputs
     output, score = outputs
-    print_str = """
+    if len(score.shape) > 0:
+      print_str = """
+Input:
+{inputs}
+
+Output (Scores [{score}]):
+{output}
+        """
+      score_text = ",".join(["{:.3f}".format(s) for s in score])
+      print(print_str.format(inputs=inputs, output=output, score=score_text))
+    else:
+      print_str = """
 Input:
 {inputs}
 
 Output (Score {score:.3f}):
 {output}
-    """
-    print(print_str.format(inputs=inputs, output=output, score=score))
+        """
+      print(print_str.format(inputs=inputs, output=output, score=score))
+
     if FLAGS.inputs_once:
       break
 

tensor2tensor/serving/serving_utils.py

@@ -94,7 +94,10 @@ def _encode(inputs, encoder, add_eos=True):
 
 
 def _decode(output_ids, output_decoder):
-  return output_decoder.decode(output_ids, strip_extraneous=True)
+  if len(output_ids.shape) > 1:
+    return [output_decoder.decode(o, strip_extraneous=True) for o in output_ids]
+  else:
+    return output_decoder.decode(output_ids, strip_extraneous=True)