ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/gemma1.py

@@ -0,0 +1,107 @@
+# 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 Gemma1 model."""
+
+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
+from torch import nn
+
+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_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="embedder",
+    final_norm="model.norm",
+    lm_head=None,
+)
+
+
+class Gemma1(model_builder.DecoderOnlyModel):
+  """A Gemma1 model built from the Edge Generative API layers."""
+  pass
+
+
+def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for a Gemma 2B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for a Gemma 2B 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,
+  )
+  embedding_dim = 2048
+  config = cfg.ModelConfig(
+      vocab_size=256000,
+      num_layers=18,
+      max_seq_len=8192,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**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_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  config = get_model_config_2b(kv_cache_max_len)
+  # Gemma 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_2b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_model_config_2b(**kwargs),
+      tensor_names=TENSOR_NAMES,
+      model_class=Gemma1,
+  )

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -0,0 +1,295 @@
+# 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 Gemma2 model."""
+
+from typing import List, Optional, Tuple
+
+from ai_edge_torch.generative.layers import attention
+from ai_edge_torch.generative.layers import builder
+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.layers.rotary_position_embedding as rotary_pos_emb
+from ai_edge_torch.generative.utilities import model_builder
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+
+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_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",
+    pre_ff_norm="model.layers.{}.pre_feedforward_layernorm",
+    post_ff_norm="model.layers.{}.post_feedforward_layernorm",
+    embedding="embedder",
+    final_norm="model.norm",
+    lm_head=None,
+)
+
+
+class Gemma2Block(attention.TransformerBlock):
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+      mask: Optional[torch.Tensor] = None,
+      input_pos: Optional[torch.Tensor] = None,
+      kv_cache: kv_utils.KVCacheEntry = None,
+  ) -> Tuple[torch.Tensor, Optional[kv_utils.KVCacheEntry]]:
+    """Forward function of the Gemma2Block.
+
+    Exactly the same as TransformerBlock but we call the post-attention norm
+    immediately after attention and not after the residual pointwise addition.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+      kv_cache (KVCacheEntry): the optional kv cache entry.
+
+    Returns:
+      output activation from this transformer block, and updated kv cache (if
+      passed in).
+    """
+
+    x_norm = self.pre_atten_norm(x)
+    attn_out, kv = self.atten_func(x_norm, rope, mask, input_pos, kv_cache)
+    attn_out_norm = self.post_atten_norm(attn_out)
+    x = x + attn_out_norm
+    output = x + self.ff(x)
+    return output, kv
+
+
+class Gemma2(nn.Module):
+  """A Gemma2 model built from the Edge Generative API layers."""
+
+  def __init__(self, config: cfg.ModelConfig):
+    super().__init__()
+
+    # Construct model layers.
+    self.tok_embedding = nn.Embedding(
+        config.vocab_size, config.embedding_dim, padding_idx=0
+    )
+    self.lm_head = nn.Linear(
+        config.embedding_dim,
+        config.vocab_size,
+        bias=config.lm_head_use_bias,
+    )
+    # Gemma2 re-uses the embedding as the head projection layer.
+    self.lm_head.weight.data = self.tok_embedding.weight.data
+    self.transformer_blocks = nn.ModuleList(
+        Gemma2Block(config.block_config(idx), config)
+        for idx in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.kv_cache_max,
+    )
+    # Gemma2 has same hyper parameters for each layer except for attention
+    # types. Use the first layer.
+    attn_config = config.block_config(0).attn_config
+    self.sliding_window_mask_cache = attn_utils.build_sliding_window_mask_cache(
+        size=config.kv_cache_max,
+        window_size=attn_config.sliding_window_size,
+    )
+    self.config = config
+
+  def get_attention_mask(
+      self, attn_type: cfg.AttentionType, input_pos: torch.Tensor
+  ) -> torch.Tensor:
+    if attn_type == cfg.AttentionType.LOCAL_SLIDING:
+      return self.sliding_window_mask_cache.index_select(2, input_pos)
+    return self.mask_cache.index_select(2, input_pos)
+
+  @torch.inference_mode
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    _, seq_len = tokens.size()
+    assert self.config.max_seq_len >= seq_len, (
+        f"Cannot forward sequence of length {seq_len}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
+
+    # token embeddings of shape (b, t, n_embd)
+    input_embeds = self.tok_embedding(tokens)
+    # RoPE parameters are the same for all blocks. Use the first layer.
+    attn_config = self.config.block_config(0).attn_config
+    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
+    rope = rotary_pos_emb.build_rope(input_pos, n_elem, attn_config.rotary_base)
+    mask = [
+        self.get_attention_mask(
+            self.config.block_config(i).attn_config.attn_type, input_pos
+        )
+        for i in range(self.config.num_layers)
+    ]
+
+    return self._forward_with_embeds(
+        input_embeds, rope, mask, input_pos, kv_cache, export_config
+    )
+
+  def _forward_with_embeds(
+      self,
+      input_embeds: torch.Tensor,
+      rope: Tuple[torch.Tensor, torch.Tensor],
+      mask: List[torch.Tensor],
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      export_config: Optional[model_builder.ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    """Forwards the model with input embeddings."""
+    assert len(self.transformer_blocks) == len(kv_cache.caches), (
+        "The number of transformer blocks and the number of KV cache entries"
+        " must be the same."
+    )
+
+    if self.config.embedding_scale is not None:
+      input_embeds = input_embeds * self.config.embedding_scale
+    x = input_embeds
+    updated_kv_entries = []
+    mask_input = mask is not None
+    for i, block in enumerate(self.transformer_blocks):
+      mask = (
+          mask
+          if mask_input
+          else self.get_attention_mask(
+              block.config.attn_config.attn_type, input_pos
+          )
+      )
+      kv_entry = kv_cache.caches[i] if kv_cache else None
+      x, kv_entry = block(x, rope, mask[i], input_pos, kv_entry)
+      if kv_entry:
+        updated_kv_entries.append(kv_entry)
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+
+    if export_config is not None:
+      if (
+          torch.numel(input_pos) > 1
+          and not export_config.output_logits_on_prefill
+      ):
+        return {"kv_cache": updated_kv_cache}
+
+    x = self.final_norm(x)
+    res = self.lm_head(x)  # (b, t, vocab_size)
+    if self.config.final_logit_softcap is not None:
+      res = res / self.config.final_logit_softcap
+      res = torch.tanh(res)
+      res = res * self.config.final_logit_softcap
+
+    return {"logits": res, "kv_cache": updated_kv_cache}
+
+
+def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
+  """Returns the model config for a Gemma2 2B model.
+
+  Args:
+    kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
+      is 1024.
+
+  Returns:
+    The model config for a Gemma 2B model.
+  """
+  norm_config = cfg.NormalizationConfig(
+      type=cfg.NormalizationType.RMS_NORM,
+      epsilon=1e-6,
+      zero_centered=True,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.GELU_TANH),
+      intermediate_size=9216,
+      pre_ff_norm_config=norm_config,
+      post_ff_norm_config=norm_config,
+  )
+
+  def get_block_config(idx: int) -> cfg.TransformerBlockConfig:
+    attn_config = cfg.AttentionConfig(
+        num_heads=8,
+        head_dim=256,
+        num_query_groups=4,
+        rotary_base=10000,
+        rotary_percentage=1.0,
+        qkv_transpose_before_split=True,
+        logit_softcap=50.0,
+        sliding_window_size=4096,
+        attn_type=(
+            cfg.AttentionType.GLOBAL
+            if idx % 2 == 0
+            else cfg.AttentionType.LOCAL_SLIDING
+        ),
+    )
+    return cfg.TransformerBlockConfig(
+        attn_config=attn_config,
+        ff_config=ff_config,
+        pre_attention_norm_config=norm_config,
+        post_attention_norm_config=norm_config,
+    )
+
+  num_layers = 26
+  embedding_dim = 2304
+  config = cfg.ModelConfig(
+      vocab_size=256000,
+      num_layers=num_layers,
+      max_seq_len=8192,
+      embedding_dim=embedding_dim,
+      embedding_scale=embedding_dim**0.5,
+      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,
+      lm_head_use_bias=False,
+      enable_hlfb=True,
+      final_logit_softcap=30.0,
+  )
+  return config
+
+
+def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
+  config = get_model_config_2b(kv_cache_max_len)
+  config.vocab_size = 128
+  config.num_layers = 2
+  config.max_seq_len = 2 * kv_cache_max_len
+  config.embedding_dim = 128
+  config.embedding_scale = config.embedding_dim**0.5
+  config.block_configs = config.block_configs[: config.num_layers]
+  for block_config in config.block_configs:
+    block_config.attn_config.num_heads = 4
+    block_config.attn_config.head_dim = 64
+    block_config.attn_config.sliding_window_size = 64
+    block_config.ff_config.intermediate_size = 128
+  return config
+
+
+def build_2b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+  return model_builder.build_decoder_only_model(
+      checkpoint_path=checkpoint_path,
+      config=get_model_config_2b(**kwargs),
+      tensor_names=TENSOR_NAMES,
+      model_class=Gemma2,
+  )

ai_edge_torch/generative/examples/gemma/verify_gemma1.py

@@ -0,0 +1,56 @@
+# 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."""
+
+import logging
+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
+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")
+
+  logging.info("Building the reauthored model from: %s", 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,43 @@
+# 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 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.",
+)
+
+
+def main(_):
+  checkpoint = kagglehub.model_download("google/gemma-2/pyTorch/gemma-2-2b-it")
+
+  verify_util.verify_gemma2(checkpoint, _PROMPTS.value, _MAX_NEW_TOKENS.value)
+
+
+if __name__ == "__main__":
+  app.run(main)

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -0,0 +1,157 @@
+# 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 logging
+import os
+from typing import List, Tuple
+
+from ai_edge_torch.generative.examples.gemma import gemma2
+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
+
+
+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_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) -> torch.Tensor:
+    """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 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 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 = "model.ckpt",
+    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.
+
+  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=verifier.ReauthoredModelWrapper(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_gemma2(
+    gemma2_model_path: str, prompts: List[str], max_new_tokens: int
+) -> bool:
+  """Verifies the reauthored Gemma2 model.
+
+  Return True if the verification passes, False otherwise.
+  """
+  logging.info("Building the reauthored model from: %s", gemma2_model_path)
+  reauthored_model = gemma2.build_2b_model(gemma2_model_path)
+
+  return verify_reauthored_gemma_model(
+      checkpoint=gemma2_model_path,
+      variant="2b-v2",
+      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,
+      atol=1e-04,
+  )

ai_edge_torch/generative/examples/llama/__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.
+# ==============================================================================