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

ai_edge_torch/generative/examples/gemma3/verify_util.py

@@ -0,0 +1,247 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utility functions to verify the reauthored Gemma model."""
+
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from ai_edge_torch.generative.examples.gemma3 import gemma3
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
+from ai_edge_torch.generative.utilities.experimental import verifier
+from gemma import config as gemma_config
+from gemma import model as gemma_model
+import torch
+
+
+def _get_actual_input_len(tokens: torch.Tensor) -> int:
+  for i in range(tokens.shape[1]):
+    if tokens[0, i] == 0:
+      return i
+  return tokens.shape[1]
+
+
+class GemmaWrapper(verifier.ModelWrapper):
+  """Gemma model wrapper for verification.
+
+  Verifier calls model.forward() with maxium sequence length (1024) expecting
+  the output is logits while Gemma gets the input tokens with the actual length
+  and returns logits in a tuple.
+
+  Verifier runs tokenizer before model.generate() while Gemma runs the tokenizer
+  inside model.generate().
+  """
+
+  def _get_kv_caches(
+      self, max_seq_len: int
+  ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+    config = self.model.config
+    cache_size = (1, max_seq_len, config.num_key_value_heads, config.head_dim)
+    cache = torch.zeros(cache_size)
+    return [
+        (cache.clone(), cache.clone()) for _ in range(config.num_hidden_layers)
+    ]
+
+  def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+    """Forwards the model after reducing input tokens to the actual length."""
+    actual_input_len = _get_actual_input_len(tokens)
+    input_pos = torch.arange(0, actual_input_len, dtype=torch.long)
+    mask_cache = attn_utils.build_causal_mask_cache(tokens.shape[1])
+    local_mask_cache = attn_utils.build_sliding_window_mask_cache(
+        tokens.shape[1], self.model.config.sliding_window_size
+    )
+    _, logits = self.model.forward(
+        input_token_ids=tokens[0, :actual_input_len].unsqueeze(0),
+        input_positions=input_pos,
+        kv_write_indices=None,
+        kv_caches=self._get_kv_caches(tokens.shape[1]),
+        mask=mask_cache.index_select(2, input_pos),
+        output_positions=input_pos,
+        temperatures=None,
+        top_ps=torch.tensor([1.0], dtype=torch.float),
+        top_ks=torch.tensor([1], dtype=torch.long),
+        local_mask=local_mask_cache.index_select(2, input_pos),
+    )
+    return logits
+
+  def generate(
+      self, tokens: torch.Tensor, max_new_tokens: int
+  ) -> torch.IntTensor:
+    """Generates the response after decoding the tokens into a string."""
+    prompts = self.model.tokenizer.decode(tokens[0].tolist())
+    response = self.model.generate(
+        prompts, device="cpu", output_len=max_new_tokens, top_k=1
+    )
+    return torch.tensor([self.model.tokenizer.encode(prompts + response)])
+
+
+class UnifiedGemma3Wrapper(verifier.ReauthoredModelWrapper):
+  """Unified Gemma3 model wrapper for verification."""
+
+  def _init_kv_cache(self):
+    """Returns an initialized KV cache."""
+    return kv_utils.KVCacheTransposed.from_model_config(self.model.model.config)
+
+  def forward(
+      self, tokens: torch.Tensor, pixel_values: torch.Tensor = None
+  ) -> torch.Tensor:
+    """Forwards the model."""
+    mask = attn_utils.build_causal_mask_cache(
+        self.model.model.config.kv_cache_max_len
+    )
+    input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
+    mask = mask.index_select(2, input_pos)
+    output = self.model.model.forward(
+        tokens, input_pos, self._init_kv_cache(), mask=mask
+    )
+    return output["logits"]
+
+  def generate(
+      self,
+      prompts: torch.Tensor,
+      max_new_tokens: int,
+      pixel_values: torch.Tensor = None,
+      eos_token_id: Optional[int] = None,
+  ) -> torch.IntTensor:
+    """Generates the response."""
+    input_ids = prompts[0].int().tolist()
+    tokens = torch.tensor([input_ids])
+    input_pos = torch.arange(0, tokens.shape[1], dtype=torch.int)
+    kv_cache = self._init_kv_cache()
+    mask_cache = attn_utils.build_causal_mask_cache(
+        self.model.model.config.kv_cache_max_len
+    )
+    for _ in range(max_new_tokens):
+      mask = mask_cache.index_select(2, input_pos)
+      output = self.model.model.forward(
+          tokens, input_pos, kv_cache, mask=mask
+      )
+      logits, kv_cache = output["logits"], output["kv_cache"]
+      generated_token = logits[0][-1].argmax().item()
+      input_ids.append(generated_token)
+      if eos_token_id is not None and generated_token == eos_token_id:
+        break
+      tokens = torch.tensor([[generated_token]])
+      input_pos = torch.tensor([len(input_ids) - 1])
+    return torch.tensor([input_ids])
+
+
+class GemmaTokenizerWrapper(verifier.TokenizerWrapper):
+  """Tokenizer wrapper for verification.
+
+  Verifier expects the tokenizer to handle tokens in torch.Tensor while Gemma
+  tokenizer expects tokens in a list.
+  """
+
+  def encode(self, text: str, **_) -> torch.Tensor:
+    """Adds one more dimension to the output of the tokenizer."""
+    return torch.tensor([self.tokenizer.encode(text)])
+
+  def decode(self, tokens: torch.Tensor) -> str:
+    """Decodes the token sequence after converting to a list."""
+    return self.tokenizer.decode(tokens.tolist())
+
+
+def verify_reauthored_gemma_model(
+    checkpoint: str,
+    variant: str,
+    reauthored_model: torch.nn.Module,
+    generate_prompts: List[str],
+    forward_input_ids: List[List[int]],
+    weight_filename: str,
+    tokenizer_filename: str = "tokenizer.model",
+    max_new_tokens: int = 20,
+    rtol: float = 1e-05,
+    atol: float = 1e-05,
+) -> bool:
+  """Verifies the reauthored Gemma model against the original model.
+
+  Args:
+      checkpoint: Path to the Gemma checkpoint.
+      variant: Gemma model variant.
+      reauthored_model: The reauthored model to verify.
+      generate_prompts: List of prompts for generation.
+      forward_input_ids: List of input ids for forward pass.
+      weight_filename: Name of the weight file.
+      tokenizer_filename: Name of the tokenizer file.
+      max_new_tokens: Maximum number of new tokens to generate.
+      rtol: Relative tolerance for comparison.
+      atol: Absolute tolerance for comparison.
+
+  Returns:
+      True if the verification passes, False otherwise.
+  """
+  config = gemma_config.get_model_config(variant)
+  config.tokenizer = os.path.join(checkpoint, tokenizer_filename)
+  # Use float32 to be compatible with the reauthored model.
+  config.dtype = torch.float32
+
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_model = gemma_model.GemmaForCausalLM(config).eval()
+  original_model.load_weights(os.path.join(checkpoint, weight_filename))
+
+  return verifier.verify_reauthored_model(
+      original_model=GemmaWrapper(original_model),
+      reauthored_model=UnifiedGemma3Wrapper(reauthored_model),
+      tokenizer=GemmaTokenizerWrapper(original_model.tokenizer),
+      generate_prompts=generate_prompts,
+      max_new_tokens=max_new_tokens,
+      forward_input_ids=forward_input_ids,
+      rtol=rtol,
+      atol=atol,
+  )
+
+
+def verify_gemma3(
+    checkpoint: str,
+    prompts: List[str],
+    max_new_tokens: int,
+    variant: str,
+    weight_filename: str,
+) -> bool:
+  """Verifies the reauthored Gemma3 model.
+
+  Args:
+      checkpoint: Path to the Gemma checkpoint.
+      prompts: List of prompts for generation.
+      max_new_tokens: Maximum number of new tokens to generate.
+      variant: Gemma model variant.
+      weight_filename: Name of the weight file.
+
+  Returns:
+      True if the verification passes, False otherwise.
+  """
+  gemma3_model_path = os.path.join(checkpoint, weight_filename)
+  logging.info("Building the reauthored model from: %s", gemma3_model_path)
+
+  if variant == "1b":
+    reauthored_model = UnifiedGemma3Wrapper(
+        gemma3.build_model_1b(gemma3_model_path)
+    )
+  else:
+    raise ValueError(f"Unsupported Gemma3 variant: {variant}")
+
+  return verify_reauthored_gemma_model(
+      checkpoint=checkpoint,
+      variant=variant,
+      reauthored_model=reauthored_model,
+      generate_prompts=prompts,
+      forward_input_ids=[[2, 651, 9456, 576, 573, 3520, 3858, 603, 235248]],
+      max_new_tokens=max_new_tokens,
+      weight_filename=weight_filename,
+      atol=1e-04,
+  )

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(),
   )
 

ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py

@@ -16,62 +16,25 @@
 """Example of converting a Phi-4 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 phi4
 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/phi4'),
-    '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',
-    'phi4',
-    '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("phi4")
 
 def main(_):
   pytorch_model = phi4.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/phi/convert_to_tflite.py

@@ -24,54 +24,19 @@ from ai_edge_torch.generative.examples.phi import phi2
 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/phi2'),
-    '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',
-    'phi2',
-    '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("phi2")
 
 def main(_):
   pytorch_model = phi2.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(),
   )