ai-edge-torch-nightly 0.3.0.dev20241108__py3-none-any.whl → 0.3.0.dev20241109__py3-none-any.whl

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -93,6 +93,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=[get_block_config(i) for i in range(num_layers)],
       final_norm_config=norm_config,
+      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/examples/paligemma/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 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.
+# ==============================================================================

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -0,0 +1,103 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+"""Example of building a decoder of PaliGemma 3B model which is Gemma1."""
+
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import model_builder
+import ai_edge_torch.generative.utilities.loader as loading_utils
+
+TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
+    ff_up_proj="language_model.model.layers.{}.mlp.up_proj",
+    ff_down_proj="language_model.model.layers.{}.mlp.down_proj",
+    ff_gate_proj="language_model.model.layers.{}.mlp.gate_proj",
+    attn_query_proj="language_model.model.layers.{}.self_attn.q_proj",
+    attn_key_proj="language_model.model.layers.{}.self_attn.k_proj",
+    attn_value_proj="language_model.model.layers.{}.self_attn.v_proj",
+    attn_output_proj="language_model.model.layers.{}.self_attn.o_proj",
+    pre_attn_norm="language_model.model.layers.{}.input_layernorm",
+    post_attn_norm="language_model.model.layers.{}.post_attention_layernorm",
+    embedding="language_model.model.embed_tokens",
+    final_norm="language_model.model.norm",
+    lm_head=None,
+)
+
+
+def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for the decoder of a PaliGemma 3B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for the decoder of a PaliGemma 3B model.
+  """
+  attn_config = cfg.AttentionConfig(
+      num_heads=8,
+      head_dim=256,
+      num_query_groups=1,
+      rotary_base=10000,
+      rotary_percentage=1.0,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=16384,
+  )
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+      zero_centered=True,
+  )
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=257216,
+      num_layers=18,
+      max_seq_len=8192,
+      embedding_dim=2048,
+      embedding_scale=2048**0.5,
+      kv_cache_max_len=kv_cache_max_len,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_fake_decoder_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  config = get_decoder_config(kv_cache_max_len)
+  # PaliGemma decoder has only one block config.
+  config.block_config(0).ff_config.intermediate_size = 128
+  config.vocab_size = 128
+  config.num_layers = 2
+  config.max_seq_len = 2 * kv_cache_max_len
+  return config
+
+
+def build_decoder(
+    checkpoint_path: str, **kwargs
+) -> model_builder.DecoderOnlyModel:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_decoder_config(**kwargs),
+      tensor_names=TENSOR_NAMES,
+  )

ai_edge_torch/generative/examples/paligemma/verify_decoder.py

@@ -0,0 +1,75 @@
+# Copyright 2024 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 decoder of PaliGemma 3B model."""
+
+import logging
+import pathlib
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.paligemma import decoder
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+_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.",
+)
+
+
+def main(_):
+  checkpoint = "google/paligemma-3b-mix-224"
+  logging.info("Loading the original model from: %s", checkpoint)
+  original_full_model = (
+      transformers.PaliGemmaForConditionalGeneration.from_pretrained(checkpoint)
+  )
+  original_language_model = original_full_model.eval().language_model
+
+  # Locate the cached dir.
+  cached_config_file = transformers.utils.cached_file(
+      checkpoint, transformers.utils.CONFIG_NAME
+  )
+  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = decoder.build_decoder(reauthored_checkpoint)
+
+  logging.info("Loading the tokenizer from: %s", checkpoint)
+  # It works only when GemmaTokenizerFast is available. In some environments,
+  # use_fast=False doeesn't work either if the tokenizer cannot load the
+  # sentencepiece model file properly.
+  processor = transformers.AutoProcessor.from_pretrained(checkpoint)
+
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_language_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(processor.tokenizer),
+      generate_prompts=_PROMPTS.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/utilities/verifier.py

@@ -143,7 +143,7 @@ def verify_with_input_ids(
     kv_cache_max_len: int = 1024,
     rtol: float = 1e-05,
     atol: float = 1e-05,
-) -> bool:
+):
   """Verifies if the model reauthored generates the same output of the oringal.
 
   It compares only one outputs from the original and the reauthored model.
@@ -157,8 +157,9 @@ def verify_with_input_ids(
     rtol (float): The relative tolerance for the comparison.
     atol (float): The absolute tolerance for the comparison.
 
-  Returns:
-    True if the model reauthored generates the same output of the original.
+  Raises:
+    AssertError if the model reauthored fails to generate the same output of the
+      original.
   """
   tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
   tokens[0, : len(input_ids)] = torch.tensor([input_ids]).int()
@@ -173,7 +174,7 @@ def verify_with_input_ids(
   logits_reauthored = outputs_reauthored[0, len(input_ids) - 1, :]
   logging.info("logits_reauthored: %s", logits_reauthored)
 
-  return torch.allclose(
+  assert torch.allclose(
       logits_original, logits_reauthored, rtol=rtol, atol=atol
   )
 
@@ -184,7 +185,7 @@ def verify_model_with_prompts(
     tokenizer: TokenizerWrapper,
     prompts: str,
     max_new_tokens: int,
-) -> bool:
+):
   """Verifies if the model reauthored generates the same answer of the oringal.
 
   It compares an answer, i.e. multiple continuous outputs generated by the
@@ -198,8 +199,9 @@ def verify_model_with_prompts(
     prompts (str): The input prompts to generate answers.
     max_new_tokens (int): The maximum number of new tokens to generate.
 
-  Returns:
-    True if the model reauthored generates the same answer of the original.
+  Raises:
+    AssertError if the model reauthored fails to generate the same answer of the
+      original.
   """
   prompt_tokens = tokenizer.encode(prompts)
 
@@ -213,7 +215,7 @@ def verify_model_with_prompts(
   response_reauthored = tokenizer.decode(outputs_reauthored[0])
   logging.info("outputs from reauthored model: [[%s]]", response_reauthored)
 
-  return response_original == response_reauthored
+  assert response_original == response_reauthored
 
 
 def verify_reauthored_model(
@@ -225,6 +227,7 @@ def verify_reauthored_model(
     forward_input_ids: List[List[int]] = [[1, 2, 3, 4]],
     rtol: float = 1e-05,
     atol: float = 1e-05,
+    continue_on_failure: bool = False,
 ):
   """Verifies the reauthored model against the original model.
 
@@ -247,21 +250,31 @@ def verify_reauthored_model(
       forward with.
     rtol (float): The relative tolerance for the comparison.
     atol (float): The absolute tolerance for the comparison.
+    continue_on_failure (bool): If True, it continues to verify the next prompt
+      or input IDs even if a previous one fails.
   """
   for input_ids in forward_input_ids:
     logging.info("Verifying the reauthored model with input IDs: %s", input_ids)
-    if verify_with_input_ids(
-        original_model, reauthored_model, input_ids, rtol=rtol, atol=atol
-    ):
-      logging.info("PASS")
+    try:
+      verify_with_input_ids(
+          original_model, reauthored_model, input_ids, rtol=rtol, atol=atol
+      )
+    except AssertionError as e:
+      logging.error("*** FAILED *** verify with input IDs: %s", input_ids)
+      if not continue_on_failure:
+        raise e
     else:
-      logging.error("FAILED")
+      logging.info("*** PASSED *** verify with input IDs: %s", input_ids)
 
   for prompts in generate_prompts:
-    logging.info("Verifying the reauthored model with prompts:%s", prompts)
-    if verify_model_with_prompts(
-        original_model, reauthored_model, tokenizer, prompts, max_new_tokens
-    ):
-      logging.info("PASS")
+    logging.info("Verifying the reauthored model with prompts: %s", prompts)
+    try:
+      verify_model_with_prompts(
+          original_model, reauthored_model, tokenizer, prompts, max_new_tokens
+      )
+    except AssertionError as e:
+      logging.error("*** FAILED *** verify with prompts: %s", prompts)
+      if not continue_on_failure:
+        raise e
     else:
-      logging.error("FAILED")
+      logging.info("*** PASSED *** verify with prompts: %s", prompts)

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241108.dist-info → ai_edge_torch_nightly-0.3.0.dev20241109.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241108
+Version: 0.3.0.dev20241109
 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.dev20241108.dist-info → ai_edge_torch_nightly-0.3.0.dev20241109.dist-info}/RECORD

@@ -3,7 +3,7 @@ ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=sBOl0mgVPJtokiP8qTbTtY0R_qIaF0KNiALh7P3AJEk,706
+ai_edge_torch/version.py,sha256=gXIkg8ND_yLshQAgj7V9B03CLweQwexHau-idN_Q-SQ,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -58,8 +58,11 @@ ai_edge_torch/generative/examples/llama/llama.py,sha256=AMcCbuDBxEfbO-l3KiEXbUaX
 ai_edge_torch/generative/examples/llama/verify.py,sha256=X7oKQi85M789ugBrOlMvzk8eSRR3Kf1Mprfl-U-WIpo,2842
 ai_edge_torch/generative/examples/openelm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=85FVEt6cKFP2UzCLC78tAkbwGlGhAArtG7Wa75NxJik,2185
-ai_edge_torch/generative/examples/openelm/openelm.py,sha256=JsrtuUY4q1Rovxsht2cGCuANUj1sUKnah6bAoSe8AoU,4387
+ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sFakstoPDcOHSak0IGFEEq_HQMBBSMcx-WVCDZqcVDo,4411
 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/decoder.py,sha256=JSb9h3gcIh5oYrbLU6rI8OU8FzfWeTCFJT5XRWu4btE,3675
+ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
 ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py,sha256=rkbTtMaqSVG48cm-NTxR_LDgZmXAEBqayTm9O49oMXc,2171
 ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=3go690yX6PFeXMdpY7y4JZorAwxX0HT_b_pKZieauvk,2169
@@ -139,7 +142,7 @@ ai_edge_torch/generative/utilities/model_builder.py,sha256=89jt80UUfDzYBi-x077HB
 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
 ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
-ai_edge_torch/generative/utilities/verifier.py,sha256=wQ4EtIED_a6FRsaOXeoQVZiHNx07esOYCQYbDVLgZ2o,9520
+ai_edge_torch/generative/utilities/verifier.py,sha256=5C2cm54d9kwL7nGRX-YfnBIJny1ICNhiU-LB3IqJq2E,10075
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k63ylp7sdZmtZ8o,4790
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
@@ -186,8 +189,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.dev20241108.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/METADATA,sha256=gp2VN_X4YPdK8axZYIhqafgiJhCwfiN_tOWT-yL3lW0,1897
-ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241109.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241109.dist-info/METADATA,sha256=Bx7bAskkKsT43fx4RZoccotaUwx2QOPeSWUfnarpsU4,1897
+ai_edge_torch_nightly-0.3.0.dev20241109.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
+ai_edge_torch_nightly-0.3.0.dev20241109.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241109.dist-info/RECORD,,

{ai_edge_torch_nightly-0.3.0.dev20241108.dist-info → ai_edge_torch_nightly-0.3.0.dev20241109.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.44.0)
+Generator: bdist_wheel (0.45.0)
 Root-Is-Purelib: true
 Tag: py3-none-any