Export models with magic numbers if gpu_dynamic_shapes=true

- Remove duplicated code to convert pytorch model to tflite
- Promote decode_batch_size as a common flag, not only for smollm
- Exporting test signatures will be in a following CL

PiperOrigin-RevId: 811525306

Author: protobird-git

ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("amd-llama-135m")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = amd_llama_135m.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(amd_llama_135m.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -18,33 +18,14 @@
 from absl import app
 from ai_edge_torch.generative.examples.deepseek import deepseek
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags(
     'deepseek', default_mask_as_input=True, default_transpose_kv_cache=True
 )
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = deepseek.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(deepseek.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("gemma-2b")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = gemma1.build_2b_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(gemma1.build_2b_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -18,33 +18,14 @@
 from absl import app
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags(
     "gemma2-2b", default_mask_as_input=True, default_transpose_kv_cache=True
 )
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = gemma2.build_2b_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(gemma2.build_2b_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags(
     'gemma3-1b', default_mask_as_input=True, default_transpose_kv_cache=True
@@ -33,36 +31,14 @@
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
   if _MODEL_SIZE.value == '1b':
-    pytorch_model = gemma3.build_model_1b(
-        checkpoint_path,
-        custom_loader=loader.maybe_get_custom_loader(
-            checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-        ),
-        mask_cache_size=converter.get_mask_cache_size_from_flags(),
-    )
+    model_builder = gemma3.build_model_1b
   elif _MODEL_SIZE.value == '270m':
-    pytorch_model = gemma3.build_model_270m(
-        checkpoint_path,
-        custom_loader=loader.maybe_get_custom_loader(
-            checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-        ),
-        mask_cache_size=converter.get_mask_cache_size_from_flags(),
-    )
+    model_builder = gemma3.build_model_270m
   else:
     raise ValueError(f'Unsupported model size: {_MODEL_SIZE.value}')
 
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(model_builder)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/hammer/convert_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.hammer import hammer
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags('hammer')
 
@@ -37,24 +35,7 @@
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(_BUILDER[_MODEL_SIZE.value])
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/llama/convert_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.llama import llama
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags('llama')
 
@@ -37,24 +35,7 @@
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(_BUILDER[_MODEL_SIZE.value])
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.openelm import openelm
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("openelm")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = openelm.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(openelm.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.phi import phi3
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("phi3")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = phi3.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(phi3.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.phi import phi4
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
-from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("phi4")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = phi4.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(phi4.build_model)
 
 
 if __name__ == '__main__':