ai-edge-torch-nightly 0.5.0.dev20250513__py3-none-any.whl → 0.5.0.dev20250514__py3-none-any.whl

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -15,7 +15,7 @@
 
 """Example of building a Decoder for Gemma3 model."""
 
-from typing import List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
@@ -410,7 +410,11 @@ def get_fake_decoder_config_1b(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model_1b(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_model_1b(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs,
+) -> nn.Module:
   # TODO(b/403644647): Better error handling for loading checkpoints with
   # different tensor names.
   for tensor_names in TENSOR_NAMES_DICT.values():
@@ -420,6 +424,7 @@ def build_model_1b(checkpoint_path: str, **kwargs) -> nn.Module:
           config=get_decoder_config_1b(**kwargs),
           tensor_names=tensor_names,
           model_class=Decoder,
+          custom_loader=custom_loader,
       )
     except KeyError as ke:
       continue

ai_edge_torch/generative/examples/gemma3/gemma3.py

@@ -16,8 +16,7 @@
 """Example of building a Gemma3 gpu model."""
 
 from dataclasses import dataclass
-from typing import List, Optional, Tuple
-import xmlrpc
+from typing import List, Optional, Tuple, Callable, Dict
 
 from ai_edge_torch.generative.examples.gemma3 import decoder
 from ai_edge_torch.generative.examples.gemma3 import image_encoder
@@ -166,9 +165,14 @@ def get_fake_model_config(**kwargs) -> Gemma3MMConfig:
       mm_extra_tokens=32,
   )
 
-def build_model_1b(checkpoint_path: str, **kwargs) -> decoder.Decoder:
+
+def build_model_1b(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs,
+) -> decoder.Decoder:
   if checkpoint_path:
-    model = decoder.build_model_1b(checkpoint_path, **kwargs)
+    model = decoder.build_model_1b(checkpoint_path, custom_loader, **kwargs)
   else:
     config = decoder.get_decoder_config_1b(**kwargs)
     model = decoder.Decoder(config)

ai_edge_torch/generative/examples/gemma3/verify_util.py

@@ -17,7 +17,7 @@
 
 import logging
 import os
-from typing import List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
@@ -167,6 +167,7 @@ def verify_reauthored_gemma_model(
     generate_prompts: List[str],
     forward_input_ids: List[List[int]],
     weight_filename: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
     tokenizer_filename: str = "tokenizer.model",
     max_new_tokens: int = 20,
     rtol: float = 1e-05,
@@ -196,7 +197,14 @@ def verify_reauthored_gemma_model(
 
   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))
+  checkpoint_path = os.path.join(checkpoint, weight_filename)
+  if custom_loader is None:
+    original_model.load_weights(checkpoint_path)
+  else:
+    original_model.load_state_dict(
+        custom_loader(checkpoint_path)["model_state_dict"],
+        strict=False,
+    )
 
   return verifier.verify_reauthored_model(
       original_model=GemmaWrapper(original_model),
@@ -216,6 +224,7 @@ def verify_gemma3(
     max_new_tokens: int,
     variant: str,
     weight_filename: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
 ) -> bool:
   """Verifies the reauthored Gemma3 model.
 
@@ -225,6 +234,7 @@ def verify_gemma3(
       max_new_tokens: Maximum number of new tokens to generate.
       variant: Gemma model variant.
       weight_filename: Name of the weight file.
+      custom_loader: A custom loader to load the weights.
 
   Returns:
       True if the verification passes, False otherwise.
@@ -234,7 +244,7 @@ def verify_gemma3(
 
   if variant == "1b":
     reauthored_model = UnifiedGemma3Wrapper(
-        gemma3.build_model_1b(gemma3_model_path)
+        gemma3.build_model_1b(gemma3_model_path, custom_loader)
     )
   else:
     raise ValueError(f"Unsupported Gemma3 variant: {variant}")
@@ -247,5 +257,6 @@ def verify_gemma3(
       forward_input_ids=[[2, 651, 9456, 576, 573, 3520, 3858, 603, 235248]],
       max_new_tokens=max_new_tokens,
       weight_filename=weight_filename,
+      custom_loader=custom_loader,
       atol=1e-04,
   )

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -22,8 +22,6 @@ from ai_edge_torch.generative.utilities import export_config
 import torch
 
 flags = converter.define_conversion_flags('paligemma2-3b-224')
-ExportConfig = export_config.ExportConfig
-
 
 _VERSION = flags.DEFINE_enum(
     'version',
@@ -32,6 +30,7 @@ _VERSION = flags.DEFINE_enum(
     'The version of PaliGemma model to verify.',
 )
 
+
 def main(_):
   pytorch_model = paligemma.build_model(
       flags.FLAGS.checkpoint_path,
@@ -51,7 +50,7 @@ def main(_):
       pixel_seq_len=(config.image_size // config.patch_size) ** 2,
       quantize=flags.FLAGS.quantize,
       config=pytorch_model.config.decoder_config,
-      export_config=ExportConfig(),
+      export_config=export_config.get_from_flags(),
   )
 
 

ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py

@@ -21,8 +21,6 @@ from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
 flags = converter.define_conversion_flags('qwen_vl')
-ExportConfig = export_config.ExportConfig
-
 
 _IMAGE_HEIGHT = flags.DEFINE_integer(
     'image_height',
@@ -35,6 +33,7 @@ _IMAGE_WIDTH = flags.DEFINE_integer(
     'The width of image.',
 )
 
+
 def main(_):
   pytorch_model = qwen_vl.build_model(
       flags.FLAGS.checkpoint_path,
@@ -60,7 +59,7 @@ def main(_):
       ),
       quantize=flags.FLAGS.quantize,
       config=pytorch_model.config.decoder_config,
-      export_config=ExportConfig(),
+      export_config=export_config.get_from_flags(),
   )
 
 

ai_edge_torch/generative/utilities/converter.py

@@ -280,6 +280,15 @@ def convert_to_tflite(
       '' if not lora_ranks else f'_lora{",".join(map(str, lora_ranks))}'
   )
 
+  if pixel_values_size is not None:
+    assert pixel_seq_len > 0, 'pixel_seq_len must be greater than 0'
+    max_prefill_seq_len = max(prefill_seq_lens)
+    assert kv_size > max_prefill_seq_len + pixel_seq_len, (
+        f'The KV cache size ({kv_size}) must be greater than the maximum '
+        f'prefill sequence length ({max_prefill_seq_len}) + pixel sequence '
+        f'length ({pixel_seq_len})'
+    )
+
   if export_config is not None:
     if export_config.decode_batch_size > 1:
       output_name_prefix += f'_dbs{export_config.decode_batch_size}'

ai_edge_torch/generative/utilities/loader.py

@@ -117,7 +117,12 @@ class ModelLoader:
     final_norm: str = None
     lm_head: str = None
 
-  def __init__(self, file_name: str, names: TensorNames) -> None:
+  def __init__(
+      self,
+      file_name: str,
+      names: TensorNames,
+      custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+  ) -> None:
     """ModelLoader constructor.
 
     Can be used to load multiple models of the same type.
@@ -126,10 +131,15 @@ class ModelLoader:
         file_name (str): Path to the checkpoint. Can be a directory or an exact
           file.
         names (TensorNames): An instance of `TensorNames` to determine mappings.
+        custom_loader (Callable[[str], Dict[str, torch.Tensor]]): A custom
+          loader to be used. If not provided, the class will determine a proper
+          loader.
     """
     self._file_name = file_name
     self._names = names
-    self._loader = self._get_loader()
+    self._loader = (
+        custom_loader if custom_loader is not None else self._get_loader()
+    )
 
   def get_state(self) -> Dict[str, torch.Tensor]:
     return self._loader(self._file_name)

ai_edge_torch/generative/utilities/model_builder.py

@@ -16,6 +16,7 @@
 """Utilities to be used for re-authoring transformer models."""
 
 import copy
+from typing import Callable, Dict
 from typing import Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
@@ -160,9 +161,12 @@ def build_decoder_only_model(
     config: cfg.ModelConfig,
     tensor_names: loading_utils.ModelLoader.TensorNames,
     model_class: type[nn.Module] = DecoderOnlyModel,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
 ) -> nn.Module:
   transformer = model_class(config)
-  loader = loading_utils.ModelLoader(checkpoint_path, tensor_names)
+  loader = loading_utils.ModelLoader(
+      checkpoint_path, tensor_names, custom_loader
+  )
   loader.load(
       transformer, strict=not config.lm_head_share_weight_with_embedding
   )

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250513"
+__version__ = "0.5.0.dev20250514"

{ai_edge_torch_nightly-0.5.0.dev20250513.dist-info → ai_edge_torch_nightly-0.5.0.dev20250514.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250513
+Version: 0.5.0.dev20250514
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.5.0.dev20250513.dist-info → ai_edge_torch_nightly-0.5.0.dev20250514.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
-ai_edge_torch/version.py,sha256=Q2u2GS0KjqxWhznlOZBgkCi4NAQcdpjJzkUYdcGYQ5o,706
+ai_edge_torch/version.py,sha256=ZvSDZpKkUslpMEN4pPp4xI6n8g3mHZMdfIcYeWth5Dg,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=QVugYVfbyaeBgSKKbhFzHG5oXA7t3M-40JcpcdSu6W8,5436
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -69,11 +69,11 @@ ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=jhiyinOqPt5ZZjEa
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=n7f2nF6Lin_tDvPs0JVldsuaBzo7pAwi5YAHAhlIxQg,6139
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
 ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=JLXXn2mFEBs4DlHH_O6hpEG9KInJqsCdWy3DrgUjT1c,1827
-ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=v0ZcKrIAvERQLb1wK1Vc_ewWWVZgJFUdRTyoVY0Lfus,14955
-ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=GACDBI_MsFowR8A3wAWrpzradPYe-AUgB9ZjXaVBG-s,6485
+ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=shdgLzKDUi0vyNOAsrIVAEFb3Adltsri6Rx1-wxzVf4,15089
+ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=ZorRtnbElWsctcA0nEbfwjx0C578voF7fjFEvWSR5Ck,6582
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=uRoLoBWzFtQz5wFZfPCxbkvZsgPAqSkUUsV3977GbYc,5184
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
-ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=KnE9ME3mrpQkAxFlBOJLsqcQkjsdDL1ClNhJahX5K5I,8960
+ai_edge_torch/generative/examples/gemma3/verify_util.py,sha256=1vfAtayH_I_qTpqhzu6n9xnCuvhgTzhS8IzZviW2dJQ,9418
 ai_edge_torch/generative/examples/hammer/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/hammer/convert_to_tflite.py,sha256=9r8LXyaoBXYIIhhe1WQgEIjaxALQPE1dO2N6qopyWCk,1753
 ai_edge_torch/generative/examples/hammer/hammer.py,sha256=76INcjffvaNCQ02fzXcxJUW_6EKHs4sg3q1nDBbEpHE,3431
@@ -90,7 +90,7 @@ ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=wRdT7bWbCX
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=hPcXYHj-nBP56TOeQQejB3HRzv6yHSftHOx0OEPP5M8,4574
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=4W26ZtPF5Cb9mpHYuRM4b2QB_4W76zf4WV36KzexVjs,2446
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=7HHXkC-IIu7ieBvBI4RlXs_oITz7R8a6YVYQskAs_Uk,2023
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=fkP-mWrih1s-vgJ41fLt8v5JE-UOs8Zrngh6ElQ6PMw,1997
 ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=-EYUZp55dfRY1E-N0Pr3b9i5c7Tt1XvYxvsRixguVS8,5527
 ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=WB8r-e_Crog1ItBq3Zse_nUG-foFyBcJsuEG26r_Ji8,6076
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=SvuR97sjkBtfkerH7Hu1UXB8kCFLpEATNbPfCbNAyfo,5614
@@ -114,7 +114,7 @@ ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=eOpv3scJr4mVs
 ai_edge_torch/generative/examples/qwen/qwen.py,sha256=m8APYzo9N0SXsdvCxC8HtCcbN3W7gLKkRBL-Tg0BWXU,4223
 ai_edge_torch/generative/examples/qwen/verify.py,sha256=9_AyEJTeUfvhhID64Rto2bflFPyXMFokdQLsseLUMiI,2775
 ai_edge_torch/generative/examples/qwen_vl/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=yVebRatt2SLCsGvrYTBXOM-0S2REhkpikHTyy5MCjUw,2222
+ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=4Gntv6LBIxd0CaKkb-koLzGTdBEOGgVf3ob99lAuvuY,2196
 ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=7RFM25tDj_b0FkpSv8RUWir8K8v9p2jMtwZmP4VAUhw,4474
 ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=nHzBe_YSPnUe1d5i09v4bePQomVifzJNeUjRfprmxC0,14878
 ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=mfLFrT8NPEPh9CqlJYHwh-I2y6ST7hH_vEmbZYartHQ,7764
@@ -192,10 +192,10 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=vQWmpzMkJ2hP
 ai_edge_torch/generative/test/test_quantize.py,sha256=kKJ01wscTC2t_Ylr7huO5gNKES01gm3dT1gx52z15PA,7356
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=OMBy_nos9mEGMQOAD8o0on-gAkRk-kliodFSTthD5BE,14612
+ai_edge_torch/generative/utilities/converter.py,sha256=4zcDlhgCQQyLylH8NLgVjnelou2pW6HWJHBFYsFyHuw,15020
 ai_edge_torch/generative/utilities/export_config.py,sha256=5IvR3grlMd4mWO5c_Y4x9Fk1b1xa57MzlYNE8XUaN28,2049
-ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
-ai_edge_torch/generative/utilities/model_builder.py,sha256=IG-88o7nWI9XrNDnwnQ-MoilsuqJ7KwrnbP3bn2EY9U,6334
+ai_edge_torch/generative/utilities/loader.py,sha256=nw2REQ9sGWDwphShfRqNFICYmwIjqLp6bDcwVmsNTtg,14067
+ai_edge_torch/generative/utilities/model_builder.py,sha256=tBfOcsI_NcneggHqkCSydYN3ZgmkzPc6nW0AJrA81wI,6461
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
@@ -251,8 +251,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/METADATA,sha256=qh5r3x7C0ksa3D2WriWd0yePFgxK8urh9aSsJCC_gjY,2074
-ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250513.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/METADATA,sha256=4_d1LvNhvXOHKlqYZDcBYSLdYDmoGvWMgCK5PJasNiU,2074
+ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250514.dist-info/RECORD,,