ai-edge-torch-nightly 0.3.0.dev20250204__py3-none-any.whl → 0.3.0.dev20250205__py3-none-any.whl

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -13,11 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Example of converting a PaliGemma model to multi-signature tflite model.
-
-DISCLAIMER: It works only with ODML Torch conversion backend. Refer to
-https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/pytorch_converter/README.md#use-odml-torch-conversion-backend-experimental.
-"""
+"""Example of converting a PaliGemma model to multi-signature tflite model."""
 
 import os
 import pathlib

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -55,7 +55,6 @@ class Decoder(model_builder.DecoderOnlyModel):
       input_embeds: torch.Tensor = None,
       mask: Optional[torch.Tensor] = None,
       export_config: Optional[model_builder.ExportConfig] = None,
-      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(
@@ -64,11 +63,11 @@ class Decoder(model_builder.DecoderOnlyModel):
 
     assert input_embeds is not None
 
-    repo_pos = input_pos + 1  # PaliGemma position is 1-based.
+    rope_pos = input_pos + 1  # PaliGemma position is 1-based.
     # ROPE parameters for all attn_configs are the same. Take the first one.
     attn_config = self.config.block_config(0).attn_config
     n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
-    rope = self.config.build_rope(repo_pos, n_elem, attn_config.rotary_base)
+    rope = self.config.build_rope(rope_pos, n_elem, attn_config.rotary_base)
 
     # The first part of input_embeds are image embeddings. Diagonal causal mask
     # doesn't work here.

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -58,34 +58,23 @@ class Decoder2(gemma2.Gemma2):
       input_embeds: torch.Tensor = None,
       mask: Optional[torch.Tensor] = None,
       export_config: Optional[model_builder.ExportConfig] = None,
-      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(tokens, input_pos, kv_cache, mask, export_config)
 
     assert input_embeds is not None
 
-    repo_pos = input_pos + 1  # PaliGemma2 position is 1-based.
+    rope_pos = input_pos + 1  # PaliGemma2 position is 1-based.
     # ROPE parameters for all attn_configs are the same. Take the first one.
     attn_config = self.config.block_config(0).attn_config
     n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
-    rope = self.config.build_rope(repo_pos, n_elem, attn_config.rotary_base)
+    rope = self.config.build_rope(rope_pos, n_elem, attn_config.rotary_base)
 
     if mask is None:
-      if called_by_generate:
-        # PaliGemma2 generate() uses a diagonal causal mask even with image
-        # embeds.
-        mask = [
-            self.get_attention_mask(
-                self.config.block_config(i).attn_config.attn_type, input_pos
-            )
-            for i in range(self.config.num_layers)
-        ]
-      else:
-        # By default, don't mask image embeds with a diagonal causal mask.
-        embeds_len = input_embeds.shape[1]
-        mask = torch.zeros(embeds_len, self.config.kv_cache_max)
-        mask[:, embeds_len:] = float("-inf")
+      # By default, don't mask image embeds with a diagonal causal mask.
+      embeds_len = input_embeds.shape[1]
+      mask = torch.zeros(embeds_len, self.config.kv_cache_max)
+      mask[:, embeds_len:] = float("-inf")
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -15,7 +15,7 @@
 
 """Example of building a full-stack of PaliGemma model."""
 
-from dataclasses import dataclass
+import dataclasses
 from typing import Optional
 
 from ai_edge_torch.generative.examples.paligemma import decoder
@@ -31,7 +31,7 @@ from torch import nn
 PROJECTION_TENSOR_NAME = "multi_modal_projector.linear"
 
 
-@dataclass
+@dataclasses.dataclass
 class PaliGemmaConfig:
   """PaliGemma model configurations."""
 
@@ -39,7 +39,6 @@ class PaliGemmaConfig:
   decoder_config: cfg.ModelConfig
 
   image_token_id: int
-  image_projection_scale: float
   image_projection_use_bias: bool = False
 
 
@@ -73,7 +72,6 @@ class PaliGemma(nn.Module):
       mask: Optional[torch.Tensor] = None,
       pixel_values: torch.Tensor = None,
       export_config: Optional[model_builder.ExportConfig] = None,
-      called_by_generate: bool = True,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if pixel_values is None:
       return self.decoder(
@@ -83,14 +81,13 @@ class PaliGemma(nn.Module):
           mask=mask,
           input_embeds=None,
           export_config=export_config,
-          called_by_generate=called_by_generate,
       )
 
     input_embeds = self.decoder.tok_embedding(tokens)
 
     image_encoded = self.image_encoder(pixel_values=pixel_values)
     image_embeds = self.image_projection(image_encoded)
-    image_embeds = image_embeds / self.config.image_projection_scale
+    image_embeds = image_embeds / self.config.decoder_config.embedding_scale
 
     # Merging image_embeds into text_embeds as PaliGemmaForConditionalGeneration
     # can be done like:
@@ -116,7 +113,6 @@ class PaliGemma(nn.Module):
         mask=mask,
         input_embeds=input_embeds,
         export_config=export_config,
-        called_by_generate=called_by_generate,
     )
 
 
@@ -130,7 +126,6 @@ def get_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
       image_encoder_config=image_encoder.get_image_encoder_config(),
       decoder_config=get_decoder_config(**kwargs),
       image_token_id=257152,
-      image_projection_scale=2048**0.5,
       image_projection_use_bias=True,
   )
 
@@ -140,7 +135,6 @@ def get_fake_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
       image_encoder_config=image_encoder.get_fake_image_encoder_config(),
       decoder_config=get_decoder_config(**kwargs),
       image_token_id=127,
-      image_projection_scale=128**0.5,
       image_projection_use_bias=True,
   )
 

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -41,7 +41,7 @@ _IMAGE_URL = flags.DEFINE_string(
 )
 _PROMPTS = flags.DEFINE_string(
     "prompts",
-    "describe en",
+    "<image><bos>describe en",
     "The input prompts to generate answers.",
 )
 _MAX_NEW_TOKENS = flags.DEFINE_integer(
@@ -59,16 +59,9 @@ _CHECKPOINT = {
 class ReauthoredPaliGemmaWrapper(verifier.ReauthoredModelWrapper):
   """Reauthored PaliGemma model wrapper."""
 
-  def __init__(self, model: torch.nn.Module):
-    super().__init__(model)
-    self.forward_called_by_generate = False
-
   def _init_kv_cache(self):
     return kv_cache.KVCache.from_model_config(self.model.config.decoder_config)
 
-  def _get_extra_args_for_forward(self):
-    return {"called_by_generate": self.forward_called_by_generate}
-
 
 def main(_):
   if _VERSION.value == "1":
@@ -137,7 +130,6 @@ def main(_):
   logging.info("outputs_from_original_model: [[%s]]", response_original)
 
   logging.info("Generating answer with the reauthored model...")
-  wrapped_reauthored_model.forward_called_by_generate = True
   outputs_reauthored = wrapped_reauthored_model.generate(
       prompts=inputs["input_ids"],
       pixel_values=inputs["pixel_values"],

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20250204"
+__version__ = "0.3.0.dev20250205"

{ai_edge_torch_nightly-0.3.0.dev20250204.dist-info → ai_edge_torch_nightly-0.3.0.dev20250205.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250204
+Version: 0.3.0.dev20250205
 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.3.0.dev20250204.dist-info → ai_edge_torch_nightly-0.3.0.dev20250205.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=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=4XOGz1x6yfOnkOtBndF7qE1L3Ma12ZMJNwQ7wIWkyEs,706
+ai_edge_torch/version.py,sha256=3qCqU6b85lrBJn0A7eFSW9dGx1TkEsCXhffIwwFwUv4,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -73,12 +73,12 @@ ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=pyxRGgMxrn
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=VkigoqhAr8ew95neb3TifYv-SLOSheaWKv2AH0iKDrc,2441
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=scLsguzzuHfKYDWUd2uZkKYVRzdAbQHLd-kPam8QwvM,3004
-ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=S_W-0ojRu2Vd5SLNPs1kC-70xHB8AdSWslm-yPxyezk,5478
-ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=W009ky-yobueTzdaybSCqBAvNyArLXW3jDyp5MarzZU,6376
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=a6ISb96xhEJc1TtaFGCUiA4msKedPTAeMvkWrfIklx4,2792
+ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=z658dW_D0Iqvo6xnh4vG7_o17-Fufndyis8Rq5yafJY,5439
+ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=GZa0Ou_DvOijB2nTL_jRvGbn0_dvJPosQAPf47yqicw,5988
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=7K1xl64UvoHaYmqWjIbahwXHfppwTQ8sN7JrpGKX1XQ,5771
-ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=CEMG9gh51ev1KXPew927a6nfampiXX9bL6m-25tNYN8,6340
-ai_edge_torch/generative/examples/paligemma/verify.py,sha256=KT3Ruy40tSESxQuy-Sw01NAI3zId1BZr6Bp7FZj1wZk,5622
+ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=x1mgRtVLxkCTvlkPow3y7ADoGTjUh5uc5pF46mxatLw,6099
+ai_edge_torch/generative/examples/paligemma/verify.py,sha256=HLcu1fWMtFFFONAqVW94rOBqq4XvFHtatX3JFGOsfZw,5345
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
 ai_edge_torch/generative/examples/paligemma/verify_decoder2.py,sha256=tm-UfLr0YeBRVcQsWLBOMWI9JUzHmtPEbYK2vpITpqY,2534
 ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py,sha256=vNm-wTT8BD6zbX6GocfP1QrVoHl0zSvuVxoXN36eeiU,3540
@@ -227,8 +227,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 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.3.0.dev20250204.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/METADATA,sha256=Rf4w5EMQlNWOoFIuVlXUZPU9vmXlOJW7oB4yPrtgK0c,1966
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250204.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/METADATA,sha256=F9YG6dtQw7Vh9T4m0C2z4JAiddvpobcdY-Rxjmh4WX4,1966
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250205.dist-info/RECORD,,