Set default of mask_as_input and transport_kv_cache true for gemma2 & 3

PiperOrigin-RevId: 756532578

Author: protobird-git

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -20,7 +20,9 @@
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
-flags = converter.define_conversion_flags("gemma2-2b")
+flags = converter.define_conversion_flags(
+    "gemma2-2b", default_mask_as_input=True, default_transpose_kv_cache=True
+)
 
 
 def main(_):

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -20,7 +20,9 @@
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
-flags = converter.define_conversion_flags('gemma3-1b')
+flags = converter.define_conversion_flags(
+    'gemma3-1b', default_mask_as_input=True, default_transpose_kv_cache=True
+)
 
 _MODEL_SIZE = flags.DEFINE_string(
     'model_size',

ai_edge_torch/generative/utilities/converter.py

@@ -42,7 +42,11 @@ def forward(self, *export_args, **export_kwargs):
     return self.module(*export_args, **full_kwargs)
 
 
-def define_conversion_flags(model_name: str):
+def define_conversion_flags(
+    model_name: str,
+    default_mask_as_input: bool = False,
+    default_transpose_kv_cache: bool = False,
+):
   """Defines common flags used for model conversion."""
 
   flags.DEFINE_string(
@@ -83,13 +87,13 @@ def define_conversion_flags(model_name: str):
   )
   flags.DEFINE_bool(
       'mask_as_input',
-      False,
+      default_mask_as_input,
       'If true, the mask will be passed in as input. Otherwise, mask will be '
       'built by the model internally.',
   )
   flags.DEFINE_bool(
       'transpose_kv_cache',
-      False,
+      default_transpose_kv_cache,
       'If true, the model will be converted with transposed KV cache.',
   )
   return flags