Clean up export configs.

PiperOrigin-RevId: 744102468

Author: haozha111

ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting AMD-Llama-135m model to multi-signature tflite model."""
 
-import os
 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
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("amd-llama-135m")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = amd_llama_135m.build_model(

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -15,16 +15,13 @@
 
 """Example of converting DeepSeek R1 distilled models to tflite model."""
 
-import os
-import pathlib
-
 from absl import app
-from absl import flags
 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
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("deepseek")
+ExportConfig = export_config.ExportConfig
 
 def main(_):
   pytorch_model = deepseek.build_model(

ai_edge_torch/generative/examples/experimental/gemma/convert_gemma2_gpu_to_tflite.py

@@ -15,18 +15,15 @@
 
 """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.experimental.gemma import gemma2_gpu
 from ai_edge_torch.generative.layers.experimental import kv_cache
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 import torch
 
 flags = converter.define_conversion_flags('gemma2-2b')
+ExportConfig = export_config.ExportConfig
 
 def _create_mask(mask_len, kv_cache_max_len):
   mask = torch.full(

ai_edge_torch/generative/examples/experimental/gemma/gemma2_gpu.py

@@ -30,6 +30,7 @@
 from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -152,7 +153,7 @@ def forward(
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCacheBase,
       mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCacheBase]:
     _, seq_len = tokens.size()
     assert self.config.max_seq_len >= seq_len, (
@@ -185,7 +186,7 @@ def _forward_with_embeds(
       mask: torch.Tensor | List[torch.Tensor],
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCacheBase,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCacheBase]:
     """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (

ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py

@@ -15,13 +15,14 @@
 
 """Example of converting a Gemma1 model to multi-signature tflite model."""
 
-import os
 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.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("gemma-2b")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = gemma1.build_2b_model(

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -15,14 +15,14 @@
 
 """Example of converting a Gemma2 model to multi-signature tflite model."""
 
-import os
 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
+from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags("gemma2-2b")
+ExportConfig = export_config.ExportConfig
+
 
 def main(_):
   pytorch_model = gemma2.build_2b_model(

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -23,6 +23,7 @@
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -151,7 +152,7 @@ def forward(
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     _, seq_len = tokens.size()
     assert self.config.max_seq_len >= seq_len, (
@@ -184,7 +185,7 @@ def _forward_with_embeds(
       mask: torch.Tensor | List[torch.Tensor],
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -15,16 +15,16 @@
 
 """Example of converting a Gemma3 model to multi-signature tflite model."""
 
-import os
 from absl import app
-from absl import flags
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.layers.experimental import kv_cache
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+from ai_edge_torch.generative.utilities import export_config
 import torch
 
 flags = converter.define_conversion_flags('gemma3-1b')
+ExportConfig = export_config.ExportConfig
+
 
 _MODEL_SIZE = flags.DEFINE_string(
     'model_size',

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -23,6 +23,7 @@
 from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -244,7 +245,7 @@ def forward(
       input_embeds: Optional[torch.Tensor] = None,
       mask: Optional[torch.Tensor] = None,
       image_indices: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCacheBase]:
 
     pixel_mask = None
@@ -288,7 +289,7 @@ def _forward_with_embeds(
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCacheBase,
       pixel_mask: Optional[torch.Tensor] = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCacheBase]:
     """Forwards the model with input embeddings."""
     assert len(self.transformer_blocks) == len(kv_cache.caches), (

ai_edge_torch/generative/examples/gemma3/gemma3.py

@@ -24,7 +24,7 @@
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
-from ai_edge_torch.generative.utilities import model_builder
+from ai_edge_torch.generative.utilities import export_config as export_cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -83,7 +83,7 @@ def forward(
       image_indices: Optional[torch.Tensor] = None,
       image_feat_indices: Optional[torch.Tensor] = None,
       pixel_values: torch.Tensor = None,
-      export_config: Optional[model_builder.ExportConfig] = None,
+      export_config: Optional[export_cfg.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     _, seq_len = tokens.size()
     assert self.config.decoder_config.max_seq_len >= seq_len, (
@@ -150,6 +150,7 @@ def forward(
         export_config=export_config,
     )
 
+
 def get_fake_model_config(**kwargs) -> Gemma3MMConfig:
   return Gemma3MMConfig(
       image_encoder_config=image_encoder.get_fake_image_encoder_config(),