Enable multi-file inference

PiperOrigin-RevId: 475338276

Author: Mesh TensorFlow Team

mesh_tensorflow/transformer/utils.py

@@ -1500,6 +1500,35 @@ def input_fn(params):
   write_lines_to_file(decodes, output_filename)
 
 
+@gin.configurable
+def decode_from_files(
+    estimator,
+    vocabulary,
+    model_type,
+    batch_size,
+    sequence_length,
+    checkpoint_path=None,
+    input_filenames=gin.REQUIRED,
+    output_filenames=gin.REQUIRED,
+    eos_id=1,
+    repeats=1,
+):
+  """Decodes from multiple files and writes to output_filenames."""
+  if len(input_filenames) != len(output_filenames):
+    raise ValueError("Input and output filename lists must have equal length.")
+  for input_filename, output_filename in zip(input_filenames, output_filenames):
+    decode_from_file(estimator=estimator,
+                     vocabulary=vocabulary,
+                     model_type=model_type,
+                     batch_size=batch_size,
+                     sequence_length=sequence_length,
+                     checkpoint_path=checkpoint_path,
+                     input_filename=input_filename,
+                     output_filename=output_filename,
+                     eos_id=eos_id,
+                     repeats=repeats)
+
+
 @gin.configurable
 def decode_from_dataset(estimator,
                         vocabulary,