Internal change

PiperOrigin-RevId: 465381174

official/nlp/serving/export_savedmodel.py

@@ -72,6 +72,12 @@ def define_flags():
   flags.DEFINE_bool("convert_tpu", False, "")
   flags.DEFINE_multi_integer("allowed_batch_size", None,
                              "Allowed batch sizes for batching ops.")
+  flags.DEFINE_integer("num_batch_threads", 4,
+                       "Number of threads to do TPU batching.")
+  flags.DEFINE_integer("batch_timeout_micros", 100000,
+                       "TPU batch function timeout in microseconds.")
+  flags.DEFINE_integer("max_enqueued_batches", 1000,
+                       "Max number of batches in queue for TPU batching.")
 
 
 def lookup_export_module(task: base_task.Task):
@@ -136,11 +142,11 @@ def main(_):
     options = converter_options_pb2.ConverterOptions()
     if FLAGS.allowed_batch_size is not None:
       allowed_batch_sizes = sorted(FLAGS.allowed_batch_size)
-      options.batch_options.num_batch_threads = 4
+      options.batch_options.num_batch_threads = FLAGS.num_batch_threads
       options.batch_options.max_batch_size = allowed_batch_sizes[-1]
-      options.batch_options.batch_timeout_micros = 100000
+      options.batch_options.batch_timeout_micros = FLAGS.batch_timeout_micros
       options.batch_options.allowed_batch_sizes[:] = allowed_batch_sizes
-      options.batch_options.max_enqueued_batches = 1000
+      options.batch_options.max_enqueued_batches = FLAGS.max_enqueued_batches
     converter_cli.ConvertSavedModel(
         export_dir, tpu_dir, function_alias="tpu_candidate", options=options,
         graph_rewrite_only=True)