ai-edge-torch-nightly 0.3.0.dev20240919__py3-none-any.whl → 0.3.0.dev20240921__py3-none-any.whl

ai_edge_torch/_convert/test/test_convert.py

@@ -23,12 +23,12 @@ from ai_edge_torch import config
 from ai_edge_torch._convert import conversion_utils
 from ai_edge_torch.testing import model_coverage
 import numpy as np
-import tensorflow as tf
 import torch
 from torch import nn
 import torchvision
 
 from absl.testing import absltest as googletest
+from ai_edge_litert import interpreter as tfl_interpreter  # pylint: disable=g-direct-tensorflow-import
 
 
 @dataclasses.dataclass
@@ -466,7 +466,9 @@ class TestConvert(googletest.TestCase):
     np.testing.assert_almost_equal(edge_output["y_data_2_0"], args[1])
     np.testing.assert_almost_equal(edge_output["y_data_2_1"], args[2])
 
-    interpreter = tf.lite.Interpreter(model_content=edge_model._tflite_model)
+    interpreter = tfl_interpreter.Interpreter(
+        model_content=edge_model._tflite_model
+    )
     runner = interpreter.get_signature_runner("serving_default")
     output_details = runner.get_output_details()
     self.assertIn("x", output_details.keys())
@@ -477,7 +479,9 @@ class TestConvert(googletest.TestCase):
   def _compare_tflite_torch_args_kwargs(self, model, args, kwargs, flat_inputs):
     model.eval()
     edge_model = ai_edge_torch.convert(model, args, kwargs)
-    interpreter = tf.lite.Interpreter(model_content=edge_model._tflite_model)
+    interpreter = tfl_interpreter.Interpreter(
+        model_content=edge_model._tflite_model
+    )
     runner = interpreter.get_signature_runner("serving_default")
     input_details = runner.get_input_details()
     self.assertEqual(input_details.keys(), flat_inputs.keys())

ai_edge_torch/generative/examples/gemma/{convert_to_tflite.py → convert_gemma1_to_tflite.py}

@@ -13,14 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Example of converting a Gemma model to multi-signature tflite model."""
+"""Example of converting a Gemma1 model to multi-signature tflite model."""
 
 import os
 import pathlib
 
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.utilities import converter
 
 _CHECKPOINT_PATH = flags.DEFINE_string(
@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/gemma_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -51,12 +51,14 @@ _QUANTIZE = flags.DEFINE_bool(
 
 
 def main(_):
-  pytorch_model = gemma.build_2b_model(
+  pytorch_model = gemma1.build_2b_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'gemma_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/gemma2_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = gemma2.build_2b_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'gemma2_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/gemma/{gemma.py → gemma1.py}

@@ -13,10 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Example of building a Gemma model."""
-
-import os
-import pathlib
+"""Example of building a Gemma1 model."""
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
@@ -24,7 +21,6 @@ from ai_edge_torch.generative.layers import kv_cache as kv_utils
 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.utilities.loader as loading_utils
-import numpy as np
 import torch
 from torch import nn
 
@@ -32,13 +28,11 @@ TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
     ff_down_proj="model.layers.{}.mlp.down_proj",
     ff_gate_proj="model.layers.{}.mlp.gate_proj",
-    attn_query_proj="model.layers.{}.self_attn.q_proj",
-    attn_key_proj="model.layers.{}.self_attn.k_proj",
-    attn_value_proj="model.layers.{}.self_attn.v_proj",
+    attn_fused_qkv_proj="model.layers.{}.self_attn.qkv_proj",
     attn_output_proj="model.layers.{}.self_attn.o_proj",
     pre_attn_norm="model.layers.{}.input_layernorm",
     post_attn_norm="model.layers.{}.post_attention_layernorm",
-    embedding="model.embed_tokens",
+    embedding="embedder",
     final_norm="model.norm",
     lm_head=None,
 )
@@ -192,30 +186,3 @@ def build_2b_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model, strict=False)
   model.eval()
   return model
-
-
-def define_and_run_2b(checkpoint_path: str) -> None:
-  """Instantiates and runs a Gemma 2B model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  gemma_goldens = torch.load(current_dir / "gemma_lm_logits.pt")
-
-  kv_cache_max_len = 1024
-  model = build_2b_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :4] = idx
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  output = model.forward(tokens, input_pos, kv)
-  print("comparing with goldens..")
-  assert torch.allclose(
-      gemma_goldens, output["logits"][0, idx.shape[1] - 1, :], atol=1e-02
-  )
-
-
-if __name__ == "__main__":
-  input_checkpoint_path = os.path.join(
-      pathlib.Path.home(), "Downloads/llm_data/gemma-2b"
-  )
-  define_and_run_2b(input_checkpoint_path)

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -267,29 +267,3 @@ def build_2b_model(checkpoint_path: str, **kwargs) -> nn.Module:
   loader.load(model, strict=False)
   model.eval()
   return model
-
-
-def define_and_run_2b(checkpoint_path: str) -> None:
-  """Instantiates and runs a Gemma2 2B model."""
-
-  current_dir = pathlib.Path(__file__).parent.resolve()
-  gemma2_goldens = torch.load(current_dir / "gemma2it_2b_golden.pt")
-  print("Running GEMMA 2")
-  kv_cache_max_len = 1024
-  model = build_2b_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
-  toks = torch.from_numpy(
-      np.array([2, 651, 9456, 576, 573, 3520, 3858, 603, 235248])
-  )
-  tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  tokens[0, :9] = toks
-  input_pos = torch.arange(0, kv_cache_max_len, dtype=torch.int)
-  kv = kv_utils.KVCache.from_model_config(model.config)
-  out = model.forward(tokens, input_pos, kv)
-  out_final = out["logits"][0, 8, :]
-  assert torch.allclose(gemma2_goldens, out_final, atol=1e-04)
-
-
-if __name__ == "__main__":
-  torch.set_printoptions(sci_mode=True)
-  path = os.path.join(pathlib.Path.home(), "Downloads/llm_data/gemma2-2b")
-  define_and_run_2b(path)

ai_edge_torch/generative/examples/gemma/verify_gemma1.py

@@ -0,0 +1,55 @@
+# 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 Gemma1 model."""
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.gemma import gemma1
+from ai_edge_torch.generative.examples.gemma import verify_util
+from ai_edge_torch.generative.utilities import verifier
+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.",
+)
+
+
+def main(_):
+  checkpoint = kagglehub.model_download("google/gemma/pyTorch/2b-it")
+
+  verifier.log_msg("Building the reauthored model from", checkpoint)
+  reauthored_model = gemma1.build_2b_model(checkpoint)
+
+  verify_util.verify_reauthored_gemma_model(
+      checkpoint=checkpoint,
+      variant="2b",
+      reauthored_model=reauthored_model,
+      weight_filename="gemma-2b-it.ckpt",
+      generate_prompts=_PROMPTS.value,
+      forward_input_ids=[[1, 2, 3, 4]],
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/gemma/verify_gemma2.py

@@ -0,0 +1,55 @@
+# 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 Gemma2 model."""
+
+from absl import app
+from absl import flags
+from ai_edge_torch.generative.examples.gemma import gemma2
+from ai_edge_torch.generative.examples.gemma import verify_util
+from ai_edge_torch.generative.utilities import verifier
+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.",
+)
+
+
+def main(_):
+  checkpoint = kagglehub.model_download("google/gemma-2/pyTorch/gemma-2-2b-it")
+
+  verifier.log_msg("Building the reauthored model from", checkpoint)
+  reauthored_model = gemma2.build_2b_model(checkpoint)
+
+  verify_util.verify_reauthored_gemma_model(
+      checkpoint=checkpoint,
+      variant="2b-v2",
+      reauthored_model=reauthored_model,
+      generate_prompts=_PROMPTS.value,
+      forward_input_ids=[[2, 651, 9456, 576, 573, 3520, 3858, 603, 235248]],
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -0,0 +1,142 @@
+# 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.
+# ==============================================================================
+
+"""Utility functions to verify the reauthored Gemma model."""
+
+import dataclasses
+import os
+from typing import List, Tuple
+
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+from ai_edge_torch.generative.utilities import verifier
+from gemma import config as gemma_config
+from gemma import model as gemma_model
+import torch
+
+
+@dataclasses.dataclass
+class _Output:
+  logits: torch.Tensor
+
+
+class GemmaWrapper(verifier.ModelWrapper):
+  """Gemma model wrapper for verification.
+
+  Verifier calls model.forward() with maxium sequence length (1024) expecting
+  the output has 'logits' field 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 __init__(self, model: torch.nn.Module, max_new_tokens: int):
+    super().__init__(model)
+    self.max_new_tokens = max_new_tokens
+
+  def _get_actual_input_len(self, tokens: torch.Tensor) -> int:
+    for i in range(tokens.shape[1]):
+      if tokens[0, i] == 0:
+        return i
+    return tokens.shape[1]
+
+  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) -> _Output:
+    """Forwards the model after reducing input tokens to the actual length."""
+    actual_input_len = self._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])
+    _, 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),
+    )
+    return _Output(logits.float())
+
+  def generate(self, tokens: torch.Tensor) -> torch.Tensor:
+    """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=self.max_new_tokens, top_k=1
+    )
+    return torch.tensor([self.model.tokenizer.encode(prompts + response)])
+
+
+class TokenizerWrapper(torch.nn.Module):
+  """Tokenizer wrapper for verification.
+
+  Verifier expects the tokenizer to handle tokens in torch.Tensor while Gemma
+  tokenizer expects tokens in a list.
+  """
+
+  def __init__(self, tokenizer: torch.nn.Module):
+    super().__init__()
+    self.tokenizer = tokenizer
+
+  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 = "model.ckpt",
+    tokenizer_filename: str = "tokenizer.model",
+    max_new_tokens: int = 20,
+    rtol: float = 1e-05,
+    atol: float = 1e-05,
+):
+  """Verifies the reauthored Gemma model against the original model."""
+  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
+
+  verifier.log_msg("Loading the original model from", checkpoint)
+  original_model = gemma_model.GemmaForCausalLM(config).eval()
+  original_model.load_weights(os.path.join(checkpoint, weight_filename))
+
+  verifier.verify_reauthored_model(
+      original_model=GemmaWrapper(original_model, max_new_tokens),
+      reauthored_model=reauthored_model,
+      tokenizer=TokenizerWrapper(original_model.tokenizer),
+      generate_prompts=generate_prompts,
+      forward_input_ids=forward_input_ids,
+      rtol=rtol,
+      atol=atol,
+  )

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/openelm_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = openelm.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'openelm_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/openelm/verify.py

@@ -33,8 +33,10 @@ _PROMPTS = flags.DEFINE_multi_string(
 def main(_):
   checkpoint = "apple/OpenELM-3B"
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(
-      checkpoint, trust_remote_code=True
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(
+          checkpoint, trust_remote_code=True
+      ),
   )
 
   # Locate the cached dir.
@@ -50,10 +52,10 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
   )
 
 

ai_edge_torch/generative/examples/phi/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/phi2_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = phi2.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'phi2_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/phi/verify.py

@@ -24,15 +24,25 @@ import transformers
 
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
-    "What is the meaning of life?",
+    "Instruct: Write an email about the weather Output:",
     "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 = kagglehub.model_download("Microsoft/phi/transformers/2")
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
+  generation_config = transformers.GenerationConfig.from_pretrained(checkpoint)
+  generation_config.max_new_tokens = _MAX_NEW_TOKENS.value
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(checkpoint),
+      hf_generation_config=generation_config,
+  )
 
   verifier.log_msg("Building the reauthored model from", checkpoint)
   reauthored_model = phi2.build_model(checkpoint)
@@ -41,10 +51,10 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
       atol=1e-03,
   )
 

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/smollm_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = smollm.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'smollm_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/smollm/verify.py

@@ -33,8 +33,9 @@ _PROMPTS = flags.DEFINE_multi_string(
 def main(_):
   checkpoint = "HuggingFaceTB/SmolLM-135M"
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
-
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(checkpoint),
+  )
   # Locate the cached dir.
   cached_config_file = transformers.utils.cached_file(
       checkpoint, transformers.utils.CONFIG_NAME
@@ -47,10 +48,10 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
       atol=1e-04,
   )
 

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -30,17 +30,17 @@ _CHECKPOINT_PATH = flags.DEFINE_string(
 )
 _TFLITE_PATH = flags.DEFINE_string(
     'tflite_path',
-    '/tmp/tiny_llama_q8_seq512_ekv1024.tflite',
+    '/tmp/',
     'The tflite file path to export.',
 )
 _PREFILL_SEQ_LEN = flags.DEFINE_integer(
     'prefill_seq_len',
-    512,
+    1024,
     'The maximum size of prefill input tensor.',
 )
 _KV_CACHE_MAX_LEN = flags.DEFINE_integer(
     'kv_cache_max_len',
-    1024,
+    1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
 _QUANTIZE = flags.DEFINE_bool(
@@ -54,9 +54,11 @@ def main(_):
   pytorch_model = tiny_llama.build_model(
       _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
   )
+  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
+  output_filename = f'tinyllama_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
   converter.convert_to_tflite(
       pytorch_model,
-      tflite_path=_TFLITE_PATH.value,
+      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
       quantize=_QUANTIZE.value,
   )

ai_edge_torch/generative/examples/tiny_llama/verify.py

@@ -33,10 +33,11 @@ _PROMPTS = flags.DEFINE_multi_string(
 def main(_):
   checkpoint = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
   verifier.log_msg("Loading the original model from", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(
-      checkpoint, trust_remote_code=True
+  wrapper_model = verifier.ModelWrapper(
+      model=transformers.AutoModelForCausalLM.from_pretrained(
+          checkpoint, trust_remote_code=True
+      ),
   )
-
   # Locate the cached dir.
   cached_config_file = transformers.utils.cached_file(
       checkpoint, transformers.utils.CONFIG_NAME
@@ -49,10 +50,10 @@ def main(_):
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(
-      original_model=original_model,
+      original_model=wrapper_model,
       reauthored_model=reauthored_model,
       tokenizer=tokenizer,
-      prompts=_PROMPTS.value,
+      generate_prompts=_PROMPTS.value,
       atol=1e-04,
   )
 

ai_edge_torch/generative/layers/feed_forward.py

@@ -18,7 +18,6 @@ from typing import Callable, Optional
 
 import torch
 from torch import nn
-import torch.nn.functional as F
 
 
 class SequentialFeedForward(nn.Module):

ai_edge_torch/generative/quantize/example.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 import ai_edge_torch
-from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.quantize import quant_recipes
 import numpy as np
 import torch
@@ -22,8 +22,8 @@ import torch
 
 def main():
   # Build a PyTorch model as usual
-  config = gemma.get_fake_model_config()
-  model = gemma.Gemma(config)
+  config = gemma1.get_fake_model_config()
+  model = gemma1.Gemma(config)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
   tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx

ai_edge_torch/generative/test/test_model_conversion.py

@@ -25,7 +25,7 @@ import numpy as np
 import torch
 
 from absl.testing import absltest as googletest
-from tensorflow.lite.python import interpreter
+from ai_edge_litert import interpreter
 
 
 class TestModelConversion(googletest.TestCase):

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -17,7 +17,7 @@
 
 import ai_edge_torch
 from ai_edge_torch import config as ai_edge_config
-from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.examples.openelm import openelm
 from ai_edge_torch.generative.examples.phi import phi2
@@ -28,7 +28,7 @@ import numpy as np
 import torch
 
 from absl.testing import absltest as googletest
-from tensorflow.lite.python import interpreter
+from ai_edge_litert import interpreter
 
 
 class TestModelConversion(googletest.TestCase):
@@ -82,9 +82,9 @@ class TestModelConversion(googletest.TestCase):
       ai_edge_config.Config.use_torch_xla,
       reason="tests with custom ops are not supported on oss",
   )
-  def test_gemma(self):
-    config = gemma.get_fake_model_config()
-    pytorch_model = gemma.Gemma(config).eval()
+  def test_gemma1(self):
+    config = gemma1.get_fake_model_config()
+    pytorch_model = gemma1.Gemma(config).eval()
     self._test_model(
         config, pytorch_model, "serving_default", atol=1e-2, rtol=1e-5
     )

ai_edge_torch/generative/utilities/verifier.py

@@ -16,17 +16,65 @@
 """Common utility functions to verify the reauthored models."""
 
 import datetime
-from typing import List
+from typing import List, Optional, Union
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-import numpy as np
 import torch
+import transformers
 
 
 def log_msg(*args):
   print("[%s]" % datetime.datetime.now(), *args)
 
 
+class ModelWrapper(torch.nn.Module):
+  """A wrapper for the model to be verified, this could be a HuggingFace model
+
+  or a regular PyTorch model.
+  """
+
+  def __init__(
+      self,
+      model: torch.nn.Module,
+      model_format: str = "huggingface",
+      hf_generation_config: Optional[transformers.GenerationConfig] = None,
+  ):
+    """Initializes the wrapper.
+
+    Args:
+      model (torch.nn.Module): The original model. This could be a model built
+        from HuggingFace transformers, or a regular PyTorch model.
+      model_format (str): The format of the model. It should be either
+        "huggingface" or "pytorch".
+      hf_generation_config (transformers.GenerationConfig): The HuggingFace
+        generation config. This config will only be used if the underlying model
+        is built from HuggingFace transformers.
+    """
+    super().__init__()
+    self.model = model
+    self.model_format = model_format
+    self.hf_generation_config = hf_generation_config
+
+  def generate(
+      self, inputs: torch.Tensor
+  ) -> Union[transformers.utils.ModelOutput, torch.LongTensor]:
+    if self.model_format == "huggingface":
+      return self.model.generate(
+          inputs=inputs, generation_config=self.hf_generation_config
+      )
+    else:
+      raise NotImplementedError(
+          "generate() is not implemented for model format: %s"
+          % self.model_format
+      )
+
+  def forward(
+      self,
+      inputs: torch.Tensor,
+  ):
+    return self.model.forward(inputs)
+
+
 def forward(
     model: torch.nn.Module,
     tokens: torch.Tensor,
@@ -75,9 +123,9 @@ def generate(
 
 
 def verify_with_input_ids(
-    original_model: torch.nn.Module,
+    original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
-    input_ids: torch.Tensor = torch.from_numpy(np.array([[1, 2, 3, 4]])).int(),
+    input_ids: List[int],
     kv_cache_max_len: int = 1024,
     rtol: float = 1e-05,
     atol: float = 1e-05,
@@ -87,10 +135,10 @@ def verify_with_input_ids(
   It compares only one outputs from the original and the reauthored model.
 
   Args:
-    original_model (torch.nn.Module): The original model.
+    original_model (ModelWrapper): The original model.
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
-    input_ids (torch.Tensor): The input token IDs to forward.
+    input_ids (List[int]): The input token IDs to forward with.
     kv_cache_max_len (int): The maximum sequence length of the KV cache.
     rtol (float): The relative tolerance for the comparison.
     atol (float): The absolute tolerance for the comparison.
@@ -99,18 +147,17 @@ def verify_with_input_ids(
     True if the model reauthored generates the same output of the original.
   """
   tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.int, device="cpu")
-  input_ids_len = input_ids.shape[1]
-  tokens[0, :input_ids_len] = input_ids
+  tokens[0, : len(input_ids)] = torch.tensor([input_ids]).int()
 
   log_msg("Forwarding the original model...")
   outputs_original = original_model.forward(tokens)
-  logits_original = outputs_original.logits[0, input_ids_len - 1, :]
+  logits_original = outputs_original.logits[0, len(input_ids) - 1, :]
   log_msg("logits_original: ", logits_original)
 
   log_msg("Forwarding the reauthored model...")
   kv_cache = kv_utils.KVCache.from_model_config(reauthored_model.config)
   outputs_reauthored = forward(reauthored_model, tokens, kv_cache)
-  logits_reauthored = outputs_reauthored[0][0, input_ids_len - 1, :]
+  logits_reauthored = outputs_reauthored[0][0, len(input_ids) - 1, :]
   log_msg("logits_reauthored:", logits_reauthored)
 
   return torch.allclose(
@@ -119,7 +166,7 @@ def verify_with_input_ids(
 
 
 def verify_model_with_prompts(
-    original_model: torch.nn.Module,
+    original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
     tokenizer: torch.nn.Module,
     prompts: str,
@@ -130,7 +177,7 @@ def verify_model_with_prompts(
   original and the reauthored model.
 
   Args:
-    original_model (torch.nn.Module): The original model.
+    original_model (ModelWrapper): The original model.
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
     tokenizer (torch.nn.Module): The tokenizer.
@@ -156,10 +203,11 @@ def verify_model_with_prompts(
 
 
 def verify_reauthored_model(
-    original_model: torch.nn.Module,
+    original_model: ModelWrapper,
     reauthored_model: torch.nn.Module,
     tokenizer: torch.nn.Module,
-    prompts: List[str],
+    generate_prompts: List[str],
+    forward_input_ids: List[List[int]] = [[1, 2, 3, 4]],
     rtol: float = 1e-05,
     atol: float = 1e-05,
 ):
@@ -174,26 +222,29 @@ def verify_reauthored_model(
   It prints out "PASS" or "FAILED" to the console.
 
   Args:
-    original_model (torch.nn.Module): The original model.
+    original_model (ModelWrapper): The original model.
     reauthored_model (torch.nn.Module): The model reauthored with ai_edge_torch
       Generative API.
     tokenizer (torch.nn.Module): The tokenizer.
-    prompts (List[str]): List of the input prompts to generate answers.
+    generate_prompts (List[str]): List of the input prompts to generate answers.
+    forward_input_ids (List[torch.Tensor]): List if ihe input token IDs to
+      forward with.
     rtol (float): The relative tolerance for the comparison.
     atol (float): The absolute tolerance for the comparison.
   """
-  log_msg("Verifying the reauthored model with an arbitrary input...")
-  if verify_with_input_ids(
-      original_model, reauthored_model, rtol=rtol, atol=atol
-  ):
-    log_msg("PASS")
-  else:
-    log_msg("FAILED")
+  for input_ids in forward_input_ids:
+    log_msg("Verifying the reauthored model with input IDs:", input_ids)
+    if verify_with_input_ids(
+        original_model, reauthored_model, input_ids, rtol=rtol, atol=atol
+    ):
+      log_msg("PASS")
+    else:
+      log_msg("FAILED")
 
-  for p in prompts:
-    log_msg("Verifying the reauthored model with prompts:", p)
+  for prompts in generate_prompts:
+    log_msg("Verifying the reauthored model with prompts:", prompts)
     if verify_model_with_prompts(
-        original_model, reauthored_model, tokenizer, p
+        original_model, reauthored_model, tokenizer, prompts
     ):
       log_msg("PASS")
     else:

ai_edge_torch/model.py

@@ -27,6 +27,8 @@ from typing import Callable
 import numpy.typing as npt
 import tensorflow as tf
 
+from ai_edge_litert import interpreter as tfl_interpreter  # pylint: disable=g-direct-tensorflow-import
+
 DEFAULT_SIGNATURE_NAME = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
 
 
@@ -65,7 +67,7 @@ class TfLiteModel(Model):
       tflite_model: A TFlite serialized object.
     """
     self._tflite_model = tflite_model
-    self._interpreter_builder = lambda: tf.lite.Interpreter(
+    self._interpreter_builder = lambda: tfl_interpreter.Interpreter(
         model_content=self._tflite_model,
         experimental_default_delegate_latest_features=True,
     )
@@ -75,12 +77,13 @@ class TfLiteModel(Model):
     return self._tflite_model
 
   def set_interpreter_builder(
-      self, builder: Callable[[], tf.lite.Interpreter]
+      self, builder: Callable[[], tfl_interpreter.Interpreter]
   ) -> None:
     """Sets a custom interpreter builder.
 
     Args:
-      builder: A function that returns a `tf.lite.Interpreter` or its subclass.
+      builder: A function that returns a `tfl_interpreter.Interpreter` or its
+        subclass.
     """
     self._interpreter_builder = builder
 
@@ -166,7 +169,7 @@ class TfLiteModel(Model):
 
     # Check if this is indeed a tflite model:
     try:
-      interpreter = tf.lite.Interpreter(model_content=model_content)
+      interpreter = tfl_interpreter.Interpreter(model_content=model_content)
       interpreter.get_signature_list()
     except:
       return None

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20240919.dist-info → ai_edge_torch_nightly-0.3.0.dev20240921.dist-info}/METADATA

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

@@ -2,8 +2,8 @@ ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,116
 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=D86Gw3pIRcpnTebUPKlnPbPGJae1S6Fw4DZZ3ZkD0zw,3730
-ai_edge_torch/model.py,sha256=NYV6Mkaje_ditIEI_s_7nLP_-8i4kbGM8nRzieVkbUI,5397
-ai_edge_torch/version.py,sha256=N5hYc9s2RU44J1_oe0UfJhTFo0d4JvMlKvxNlYtK0GI,706
+ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
+ai_edge_torch/version.py,sha256=t9zajdsiowClI2fG0RkKVonPF-SUx9UBuUDOEZFU9y4,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=5uPwHhmc6kwiIz-CqaiHDejf2SOWMHrb-rYEHm69wKc,3801
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -25,7 +25,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=FSufFZEeTLBpUnzE1Iy-LvNN0mhDynWMNg7Mei8RpLQ,14973
+ai_edge_torch/_convert/test/test_convert.py,sha256=40QRxQFNeSRr4dLXJkzG-wKUlvJtsfv62cdvRrmBv5w,15097
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -39,22 +39,25 @@ ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=09VbyWErOMP9BXGwZpwvqzN5RaOqRigsELfxNRVeWns,2024
-ai_edge_torch/generative/examples/gemma/convert_to_tflite.py,sha256=qJKQu6lKuSVhn8JR7KUeInq0u6yqgxEi7hfKCrZrIqY,2019
-ai_edge_torch/generative/examples/gemma/gemma.py,sha256=hjpSPzEjPHuxwRJ-vHHtCCf2PSTnm30Mp0ajYYtDivo,7489
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=gCLOti-4xHunjphNBbx9St6faRteSakm8Oex6R1Xek0,10272
+ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=evmUj_4yygQthSRU-ke-Xn1qFNDCZKbegqINWfruKwU,2184
+ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=RZDs6oY-NLYrPNtfuJDweIHzGUL2kzpIc3AW_1p8gGg,2186
+ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=cahMzvJNJfShIw4uqoBRX5iBZrI3rvsha6wpNHzeYJ0,6369
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=KsM6VlzluTqbodG24IFr3biPxBrLay0z0gmnG0bcU2U,9277
+ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=B14IR4mIw6qBVUbiIRdfdUzHMCIJCJ0RFPsYOxA46qc,1776
+ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=kSzn1ITJXqrtNQax1m4WTAnC3gMxBpcvGA7-xTO1Iuw,1802
+ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=HBK2d8FcWFoxVDF5zk9sLSbKZEtwZQhX-K_zm4AvQtQ,5160
 ai_edge_torch/generative/examples/openelm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=HnqP3te1Qvy4SKaaqPrsG05eojiKDJShp4H3jPC9tYg,2023
+ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=85FVEt6cKFP2UzCLC78tAkbwGlGhAArtG7Wa75NxJik,2185
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=gGkHELNrt4xqnu11fCh3sJbZ7OsPyvoiF1J1aKCs5r8,7532
-ai_edge_torch/generative/examples/openelm/verify.py,sha256=2qFdyLfcefdA3s1KQ-ZGWo4XReMXkEQAvpUEyJE5iqM,2057
+ai_edge_torch/generative/examples/openelm/verify.py,sha256=QdFKymQSCYFJcYVvA63u5uIsn1YxJ0JZD5UqN6gxraI,2112
 ai_edge_torch/generative/examples/phi/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=viIkbAgknE3zxavTZtib87cMIG2_-jJXtxJPcmB2pGQ,2007
+ai_edge_torch/generative/examples/phi/convert_to_tflite.py,sha256=3go690yX6PFeXMdpY7y4JZorAwxX0HT_b_pKZieauvk,2169
 ai_edge_torch/generative/examples/phi/phi2.py,sha256=YwAszA53aOjvaMJ5wua2-5rP79N21Un_Y5yBCfFSYNU,6189
-ai_edge_torch/generative/examples/phi/verify.py,sha256=R9BjOArnn-3svoIApmP1NwO47n8KIFikOF0_MEgTOa4,1770
+ai_edge_torch/generative/examples/phi/verify.py,sha256=QPYX6weEZGMEXt_Vb2hNARPAECQBKzx-KCivd4dzOrw,2145
 ai_edge_torch/generative/examples/smollm/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=86hvBleyFXWmwy3Ke5J7x7WcCtG20D2kiBNrodE0R4w,2017
+ai_edge_torch/generative/examples/smollm/convert_to_tflite.py,sha256=zPrDTDeRVWFi9DS32uNi-RLpzOStFOk5MhNla4ixeew,2179
 ai_edge_torch/generative/examples/smollm/smollm.py,sha256=hyhMk-b5762Q2xmjdD47g85dcbBSNJXNPIsifm1DRto,3239
-ai_edge_torch/generative/examples/smollm/verify.py,sha256=JzidfVMMFDXzDdwn7ToDPuMo6eaoENNZGpEzX3f61Jk,1976
+ai_edge_torch/generative/examples/smollm/verify.py,sha256=G2dAcl-VhAbx1E1PEqM6hpzPF24HqFZaz7UBEpJSQ3w,2022
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=tL6w2dr6VP66IXjSKo9StDNP-wl0RO3fh6dIliiYlFA,4656
@@ -78,16 +81,16 @@ ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=QyLeCqDnk71WvvFH68g9UeF-HytonSk1ItGF9dc7Zj8,5854
 ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=e_Kqm5dStSrNE9_aIYC-vYJRsqLn-hJVkmR4QjYqZI0,5913
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=Yg5G1LePoryeTib35lqICqaDW6foLUzSRgwJ2FlklIw,2040
+ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=ekxd8efjMgEvauUu3PidWOC-DszPHn5sqU753F7sJIM,2201
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=tlWpa7Aun3u3w5b-9EBtW7olhmSf8W-tn5bKUIwC-ys,6044
-ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=jld5PlGOQXMIWc1WoDYL_1nnsoVzRfrg-WgnsxRgaEU,2041
+ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=LUChL5tA7FHL_DlTg5QKvGInmH9AwVVw9a-omcndiz8,2095
 ai_edge_torch/generative/fx_passes/__init__.py,sha256=jrzCB3ZyY_t5jJM1e2Czdt3DjAIL43R0_a-T-I7wOzw,1155
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=hhxSQvkDMv0isZJhmuLiod66ZODaJ8uSPSVTJVHBabQ,1931
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=Z0Y_G8IG0LmvLX2u9D8__Fkr22szB-az6wMNnZpzhkA,13233
 ai_edge_torch/generative/layers/attention_utils.py,sha256=68GXGR2HSWBFViTxX7cHifzVG-kcLS2IL2tQJPIpupg,7344
 ai_edge_torch/generative/layers/builder.py,sha256=toT9Tl1x9o5KbG-eGOEViUr4fd_4f-XLZdMQT0Ae5_8,5130
-ai_edge_torch/generative/layers/feed_forward.py,sha256=dfS1psdmomgs4EbwzkYyV_xx1xl3P1lU-3GoS8m0Avw,4221
+ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
 ai_edge_torch/generative/layers/kv_cache.py,sha256=2El7kZYnQRCRcVc63xgiAdBh9oVOksDu35p9XggvaGE,6148
 ai_edge_torch/generative/layers/model_config.py,sha256=d0Y-EFb4Rr7iLZ4Bsdf1i92KuhY1BXRqyeUN2kuu510,6923
 ai_edge_torch/generative/layers/normalization.py,sha256=l_36uFdruJwqqyubnBTM0M-iGiJfeFafyXKPPK8KHVo,6713
@@ -98,7 +101,7 @@ ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=c8rtlfDaeKmUfiiTKPmQhNW
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
 ai_edge_torch/generative/layers/unet/model_config.py,sha256=8ze9kVWMuyZVQcgK7hWYw9TM1W9lXD-2j0iMHlxoGX4,9267
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/quantize/example.py,sha256=n_YFFP3dpKjeNKYZicDGL5LqtjqwhYEIaDrC6-Ci2vE,1539
+ai_edge_torch/generative/quantize/example.py,sha256=tlACaRsz6lqOxakzpXVFJZYfFKOiFqetcYVJqWVRdPE,1542
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
 ai_edge_torch/generative/quantize/quant_recipe.py,sha256=tKnuJq6hPD23JPCB9nPAlE1UHAwdbChkgPShiVaz4CE,5156
 ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=4fgmP_GgeiFUOkIaC9ZZXC12eO3DQZdrWDXRz5YXiwU,2270
@@ -107,8 +110,8 @@ ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVu
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=W6Bh0gYDzmwb0j9HdD5_D7Z7FPToP2HSyFrmwIXuFqo,3793
 ai_edge_torch/generative/test/test_loader.py,sha256=8y74ChO3CZCfEi1eCf3-w47kRgAI4qPYCXpi8rTQXMA,3378
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=DBlqxW2IT-dZYzEfOMAp86Wtqiu6kgSWZ9BKZR1Clrw,5467
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=dUYFarOldejqbMpa0j0vIDvXlWPAancuI8di3XkGxm8,4498
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=s-EVLOQGjIeVtgNI8Ggs37pkRdErAliT6NhrrFigPOE,5459
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=mAK8Pm4mgGyilDSBtFazCRDetoqYKKB0sGC83MPKE0M,4494
 ai_edge_torch/generative/test/test_quantize.py,sha256=8geJhKwYBU20m0mdGPD1BUFwQ0lZKNtCB04SOLO18y4,5980
 ai_edge_torch/generative/test/utils.py,sha256=YvEhO2HIj1LkBs5du1UxY-cGRW9HMyAYsOUhgsTrTpA,1796
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
@@ -116,7 +119,7 @@ ai_edge_torch/generative/utilities/converter.py,sha256=MQUg2ZLmfk_2csWmQWKD_II0b
 ai_edge_torch/generative/utilities/loader.py,sha256=b9iotIhVDX-Zc9XjIDUaLxnV395AyBnkQe3dV5YA7Co,13297
 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/verifier.py,sha256=QAv1uJdI5o1yfphr_DpzxhZswKa4VG3JZUpqbCCWKMk,7114
+ai_edge_torch/generative/utilities/verifier.py,sha256=vU9KgmFS7I9jNS_3H2SWROx-rbNqtMKgQC2MRhdqQ4g,8803
 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
@@ -163,8 +166,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.dev20240919.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/METADATA,sha256=NkHYIOMz-5DNKJuSQ8wE-3Nz1R6a9YZ59M-Nq8sAnJg,1859
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240919.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/METADATA,sha256=SWy7BhOQDe0_SBF17deNndzt1bEYy7iXUxy0KznIPYM,1859
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240921.dist-info/RECORD,,