ai-edge-torch-nightly 0.4.0.dev20250330__py3-none-any.whl → 0.4.0.dev20250331__py3-none-any.whl

ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py

@@ -16,61 +16,25 @@
 """Example of converting AMD-Llama-135m model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
 from absl import flags
 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.model_builder import ExportConfig
 
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/amd-llama-135m'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'amd_llama',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
+flags = converter.define_conversion_flags("amd-llama-135m")
 
 def main(_):
   pytorch_model = amd_llama_135m.build_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
-      lora_ranks=_LORA_RANKS.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=ExportConfig(),
   )
 

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -24,54 +24,19 @@ from ai_edge_torch.generative.examples.deepseek import deepseek
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/deepseek'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'deepseek',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
-
+flags = converter.define_conversion_flags("deepseek")
 
 def main(_):
   pytorch_model = deepseek.build_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
-      lora_ranks=_LORA_RANKS.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=ExportConfig(),
   )
 

ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py

@@ -16,62 +16,24 @@
 """Example of converting a Gemma1 model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/gemma-2b'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'gemma',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
-
+flags = converter.define_conversion_flags("gemma-2b")
 
 def main(_):
   pytorch_model = gemma1.build_2b_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
-      lora_ranks=_LORA_RANKS.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=ExportConfig(),
   )
 

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -16,62 +16,25 @@
 """Example of converting a Gemma2 model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/gemma2-2b'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'gemma2',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
-
+flags = converter.define_conversion_flags("gemma2-2b")
 
 def main(_):
   pytorch_model = gemma2.build_2b_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
-      lora_ranks=_LORA_RANKS.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=ExportConfig(),
   )
 

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -16,8 +16,6 @@
 """Example of converting a Gemma3 model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.gemma3 import gemma3
@@ -26,48 +24,14 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 import torch
 
+flags = converter.define_conversion_flags('gemma3-1b')
+
 _MODEL_SIZE = flags.DEFINE_string(
     'model_size',
     '1b',
     'The size of the model to convert.',
 )
 
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/gemma3-1b'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'gemma3',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (32, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    2048,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    False,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
-
 
 def _create_mask(mask_len, kv_cache_max_len):
   mask = torch.full(
@@ -101,21 +65,22 @@ def _create_export_config(
 def main(_):
   if _MODEL_SIZE.value == '1b':
     pytorch_model = gemma3.build_model_1b(
-        _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+        flags.FLAGS.checkpoint_path,
+        kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
     )
     config = pytorch_model.config
   else:
     raise ValueError(f'Unsupported model size: {_MODEL_SIZE.value}')
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
       config=config,
-      lora_ranks=_LORA_RANKS.value,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=_create_export_config(
-          _PREFILL_SEQ_LENS.value, _KV_CACHE_MAX_LEN.value
+          flags.FLAGS.prefill_seq_lens, flags.FLAGS.kv_cache_max_len
       ),
   )
 

ai_edge_torch/generative/examples/llama/convert_to_tflite.py

@@ -16,55 +16,21 @@
 """Example of converting Llama 3.2 1B model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.llama import llama
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 
+
+flags = converter.define_conversion_flags('llama')
+
 _MODEL_SIZE = flags.DEFINE_enum(
     'model_size',
     '1b',
     ['1b', '3b'],
     'The size of the model to verify.',
 )
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/llama'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'llama',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
 
 _BUILDER = {
     '1b': llama.build_1b_model,
@@ -74,15 +40,15 @@ _BUILDER = {
 
 def main(_):
   pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
-      lora_ranks=_LORA_RANKS.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=ExportConfig(),
   )
 

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -16,62 +16,25 @@
 """Example of converting OpenELM model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.openelm import openelm
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/openelm'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'openelm',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
-
+flags = converter.define_conversion_flags("openelm")
 
 def main(_):
   pytorch_model = openelm.build_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
-      lora_ranks=_LORA_RANKS.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=ExportConfig(),
   )
 

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -16,8 +16,6 @@
 """Example of converting a PaliGemma model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.paligemma import paligemma
@@ -25,61 +23,32 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 import torch
 
+flags = converter.define_conversion_flags('paligemma2-3b-224')
+
 _VERSION = flags.DEFINE_enum(
     'version',
     '2',
     ['1', '2'],
     'The version of PaliGemma model to verify.',
 )
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/paligemma2-3b-224'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'paligemma',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LEN = flags.DEFINE_integer(
-    'prefill_seq_len',
-    1024,
-    'The maximum size of prefill input tensor.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-
 
 def main(_):
   pytorch_model = paligemma.build_model(
-      _CHECKPOINT_PATH.value,
+      flags.FLAGS.checkpoint_path,
       version=int(_VERSION.value),
-      kv_cache_max_len=_KV_CACHE_MAX_LEN.value,
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
 
   config = pytorch_model.image_encoder.config.image_embedding
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=f'{_OUTPUT_NAME_PREFIX.value}_{_VERSION.value}',
-      prefill_seq_len=_PREFILL_SEQ_LEN.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=f'{flags.FLAGS.output_name_prefix}_{_VERSION.value}',
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       pixel_values_size=torch.Size(
           [1, config.channels, config.image_size, config.image_size]
       ),
-      quantize=_QUANTIZE.value,
+      quantize=flags.FLAGS.quantize,
       config=pytorch_model.config.decoder_config,
       export_config=ExportConfig(),
   )

ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py

@@ -16,62 +16,25 @@
 """Example of converting a Phi-3.5 model to multi-signature tflite model."""
 
 import os
-import pathlib
-
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.phi import phi3
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities.model_builder import ExportConfig
 
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/phi3'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_OUTPUT_PATH = flags.DEFINE_string(
-    'output_path',
-    '/tmp/',
-    'The path to export the tflite model.',
-)
-_OUTPUT_NAME_PREFIX = flags.DEFINE_string(
-    'output_name_prefix',
-    'phi3',
-    'The prefix of the output tflite model name.',
-)
-_PREFILL_SEQ_LENS = flags.DEFINE_multi_integer(
-    'prefill_seq_lens',
-    (8, 64, 128, 256, 512, 1024),
-    'List of the maximum sizes of prefill input tensors.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-_LORA_RANKS = flags.DEFINE_multi_integer(
-    'lora_ranks',
-    None,
-    'If set, the model will be converted with the provided list of LoRA ranks.',
-)
-
+flags = converter.define_conversion_flags("phi3")
 
 def main(_):
   pytorch_model = phi3.build_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
+      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
   )
   converter.convert_to_tflite(
       pytorch_model,
-      output_path=_OUTPUT_PATH.value,
-      output_name_prefix=_OUTPUT_NAME_PREFIX.value,
-      prefill_seq_len=_PREFILL_SEQ_LENS.value,
-      quantize=_QUANTIZE.value,
-      lora_ranks=_LORA_RANKS.value,
+      output_path=flags.FLAGS.output_path,
+      output_name_prefix=flags.FLAGS.output_name_prefix,
+      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
+      quantize=flags.FLAGS.quantize,
+      lora_ranks=flags.FLAGS.lora_ranks,
       export_config=ExportConfig(),
   )