ai-edge-torch-nightly 0.4.0.dev20250328__py3-none-any.whl → 0.4.0.dev20250330__py3-none-any.whl

ai_edge_torch/generative/examples/gemma3/verify_gemma3.py

@@ -0,0 +1,90 @@
+# 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.
+# ==============================================================================
+
+"""Verifies the reauthored Gemma3 model."""
+
+import glob
+import logging
+import os
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.gemma3 import verify_util
+import kagglehub
+
+
+_PROMPTS = flags.DEFINE_multi_string(
+    "prompts",
+    "What is the meaning of life?",
+    "The input prompts to generate answers.",
+)
+_MAX_NEW_TOKENS = flags.DEFINE_integer(
+    "max_new_tokens",
+    30,
+    "The maximum size of the generated tokens.",
+)
+_CHECKPOINT = flags.DEFINE_string(
+    "checkpoint",
+    "",
+    "The checkpoint to verify.",
+)
+_VARIANT = flags.DEFINE_string(
+    "variant",
+    "1b",
+    "The variant of the model to verify.",
+)
+_WEIGHT_FILENAME = flags.DEFINE_string(
+    "weight_filename",
+    None,
+    "The weightfilename of the model to verify.",
+)
+
+
+def find_first_ckpt(folder):
+  """Finds the first .ckpt file in a folder."""
+  ckpt_files = sorted(glob.glob(os.path.join(folder, "*.ckpt")))
+  return os.path.basename(ckpt_files[0]) if ckpt_files else None
+
+
+def main(_):
+  if _CHECKPOINT.value:
+    checkpoint = _CHECKPOINT.value
+  else:
+    checkpoint = kagglehub.model_download(
+        "google/gemma-3/pyTorch/gemma-3-1b-it"
+    )
+
+  # If the weight filename is not specified, use the first checkpoint.
+  if _WEIGHT_FILENAME.value is None:
+    weight_filename = find_first_ckpt(checkpoint)
+    logging.info(
+        "NOTE: using the first weight file `%s` from `%s`",
+        weight_filename,
+        checkpoint,
+    )
+  else:
+    weight_filename = _WEIGHT_FILENAME.value
+
+  # Verify the reauthored model by comparing the outputs with the original one.
+  verify_util.verify_gemma3(
+      checkpoint,
+      _PROMPTS.value,
+      _MAX_NEW_TOKENS.value,
+      _VARIANT.value,
+      weight_filename,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

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/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250328"
+__version__ = "0.4.0.dev20250330"

{ai_edge_torch_nightly-0.4.0.dev20250328.dist-info → ai_edge_torch_nightly-0.4.0.dev20250330.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250328
+Version: 0.4.0.dev20250330
 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.4.0.dev20250328.dist-info → ai_edge_torch_nightly-0.4.0.dev20250330.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=IyhMWqN-g3wNhaYTXhegaL93NTmZkKXsXS6yx4E2kko,706
+ai_edge_torch/version.py,sha256=vjzDDfl72IBJApHHN-UvaZwIZohudx1WyOQK2VwUXa4,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -69,6 +69,8 @@ ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=xAjM
 ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=4Vf1zA94qLyNzj9iLU0jrd3kzFFZXft4uiItoIBjKyM,15632
 ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=NQzqZ55cmC8tGlZ1SKkDeD0Su8mZ79KiazCS8X08xUY,6473
 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=u30qiZu3HJCTt5noWqtf9PgGLKQ87ke4Zpa4cpG6-As,8883
 ai_edge_torch/generative/examples/llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/llama/convert_to_tflite.py,sha256=tMSsqg7LU3LR-PHtKvlWtLCqlk71mfcO9hANU4vnvDM,2734
 ai_edge_torch/generative/examples/llama/llama.py,sha256=UKvMO85_5z1vEY5MVu6QBW_vpQYA8LWHbJI4Yx6BrCc,6592
@@ -240,8 +242,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 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.4.0.dev20250328.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250328.dist-info/METADATA,sha256=j22JbcB95xcu4aqu-G4NRbh1NxwLi9GpPjJjvpLsaSE,1966
-ai_edge_torch_nightly-0.4.0.dev20250328.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250328.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250328.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250330.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250330.dist-info/METADATA,sha256=WJm9GPjx-Tg4b_ABbbP7hVrvLR5cgcR350UBcuI2j9A,1966
+ai_edge_torch_nightly-0.4.0.dev20250330.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250330.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250330.dist-info/RECORD,,