ai-edge-torch-nightly 0.5.0.dev20250516__py3-none-any.whl → 0.5.0.dev20250518__py3-none-any.whl

ai_edge_torch/generative/examples/hammer/verify.py

@@ -15,15 +15,9 @@
 
 """Verifies the reauthored Hammer 2.1 0.5B and 1.5B models."""
 
-import logging
-import pathlib
-
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.hammer import hammer
-from ai_edge_torch.generative.utilities import transformers_verifier
-from ai_edge_torch.generative.utilities import verifier
-import transformers
+from ai_edge_torch.generative.examples.hammer import verify_util
 
 
 _MODEL_SIZE = flags.DEFINE_enum(
@@ -48,37 +42,13 @@ _CHECKPOINT = {
     "1.5b": "MadeAgents/Hammer2.1-1.5b",
 }
 
-_BUILDER = {
-    "0.5b": hammer.build_0_5b_model,
-    "1.5b": hammer.build_1_5b_model,
-}
-
 
 def main(_):
-  checkpoint = _CHECKPOINT[_MODEL_SIZE.value]
-  logging.info("Loading the original model from: %s", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
-
-  # 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 = _BUILDER[_MODEL_SIZE.value](reauthored_checkpoint)
-
-  logging.info("Loading the tokenizer from: %s", checkpoint)
-  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
-
-  verifier.verify_reauthored_model(
-      original_model=transformers_verifier.TransformersModelWrapper(
-          original_model
-      ),
-      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
-      tokenizer=verifier.TokenizerWrapper(tokenizer),
-      generate_prompts=_PROMPTS.value,
+  verify_util.verify_hammer(
+      model_size=_MODEL_SIZE.value,
+      checkpoint_dir=_CHECKPOINT[_MODEL_SIZE.value],
       max_new_tokens=_MAX_NEW_TOKENS.value,
-      atol=1e-04,
+      prompts=_PROMPTS.value,
   )
 
 

ai_edge_torch/generative/examples/hammer/verify_util.py

@@ -0,0 +1,82 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils for verifying the Hammer 2.1 model."""
+import logging
+import os
+import pathlib
+
+from ai_edge_torch.generative.examples.hammer import hammer
+from ai_edge_torch.generative.utilities import loader
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+_BUILDER = {
+    "0.5b": hammer.build_0_5b_model,
+    "1.5b": hammer.build_1_5b_model,
+}
+
+DEFAULT_PROMPTS = ["What is the meaning of life?"]
+
+
+def verify_hammer(
+    model_size: str,
+    checkpoint_dir: str,
+    weight_filename: str = "model.safetensors",
+    max_new_tokens: int = 30,
+    initialize_from_local: bool = True,
+    prompts: list[str] | None = None,
+) -> bool:
+  """Verifies the reauthored Hammer 2.1 model with a custom loader."""
+  logging.info("Loading the original model from: %s", checkpoint_dir)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(
+      checkpoint_dir
+  )
+
+  logging.info("Building the reauthored model from: %s", checkpoint_dir)
+  custom_loader = (
+      None
+      if initialize_from_local
+      else loader.get_custom_loader("", "safetensors")
+  )
+
+  if initialize_from_local:
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint_dir, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    reauthored_checkpoint = os.path.join(checkpoint_dir, weight_filename)
+
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = _BUILDER[model_size](
+      checkpoint_path=reauthored_checkpoint,
+      custom_loader=custom_loader,
+  )
+
+  logging.info("Loading the tokenizer from: %s", checkpoint_dir)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint_dir)
+  return verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=DEFAULT_PROMPTS if prompts is None else prompts,
+      max_new_tokens=max_new_tokens,
+      atol=1e-04,
+  )

ai_edge_torch/generative/examples/llama/llama.py

@@ -121,9 +121,7 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True,
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -152,7 +150,6 @@ def get_1b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       kv_cache_max_len=kv_cache_max_len,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config

ai_edge_torch/generative/examples/llama/verify.py

@@ -15,15 +15,9 @@
 
 """Verifies the reauthored Llama 3.2-1B model."""
 
-import logging
-import pathlib
-
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.llama import llama
-from ai_edge_torch.generative.utilities import transformers_verifier
-from ai_edge_torch.generative.utilities import verifier
-import transformers
+from ai_edge_torch.generative.examples.llama import verify_util
 
 _MODEL_SIZE = flags.DEFINE_enum(
     "model_size",
@@ -47,40 +41,13 @@ _CHECKPOINT = {
     "3b": "meta-llama/Llama-3.2-3B-Instruct",
 }
 
-_BUILDER = {
-    "1b": llama.build_1b_model,
-    "3b": llama.build_3b_model,
-}
-
 
 def main(_):
-  checkpoint = _CHECKPOINT[_MODEL_SIZE.value]
-  logging.info("Loading the original model from: %s", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
-
-  # 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 = _BUILDER[_MODEL_SIZE.value](reauthored_checkpoint)
-
-  logging.info("Loading the tokenizer from: %s", checkpoint)
-  # Llama tokenizer_config.json sets a fast tokenizer class explicitly,
-  # "PreTrainedTokenizerFast". It works only when the fast tokenizer is
-  # available.
-  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
-
-  verifier.verify_reauthored_model(
-      original_model=transformers_verifier.TransformersModelWrapper(
-          original_model
-      ),
-      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
-      tokenizer=verifier.TokenizerWrapper(tokenizer),
-      generate_prompts=_PROMPTS.value,
+  verify_util.verify_llama_3_2(
+      model_size=_MODEL_SIZE.value,
+      checkpoint_dir=_CHECKPOINT[_MODEL_SIZE.value],
       max_new_tokens=_MAX_NEW_TOKENS.value,
-      atol=1e-04,
+      prompts=_PROMPTS.value,
   )
 
 

ai_edge_torch/generative/examples/llama/verify_util.py

@@ -0,0 +1,81 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils for verifying the Llama 3.2-1B model."""
+import logging
+import os
+import pathlib
+
+from ai_edge_torch.generative.examples.llama import llama
+from ai_edge_torch.generative.utilities import loader
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+_BUILDER = {
+    "1b": llama.build_1b_model,
+    "3b": llama.build_3b_model,
+}
+
+DEFAULT_PROMPTS = ["What is the meaning of life?"]
+
+
+def verify_llama_3_2(
+    model_size: str,
+    checkpoint_dir: str,
+    weight_filename: str = "model.safetensors",
+    max_new_tokens: int = 30,
+    initialize_from_local: bool = True,
+    prompts: list[str] | None = None,
+) -> bool:
+  """Verifies the reauthored Llama 3.2 model with a custom loader."""
+  logging.info("Loading the original model from: %s", checkpoint_dir)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(
+      checkpoint_dir
+  )
+
+  logging.info("Building the reauthored model from: %s", checkpoint_dir)
+  custom_loader = (
+      None
+      if initialize_from_local
+      else loader.get_custom_loader("", "safetensors")
+  )
+
+  if initialize_from_local:
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint_dir, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    reauthored_checkpoint = os.path.join(checkpoint_dir, weight_filename)
+
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = _BUILDER[model_size](
+      checkpoint_path=reauthored_checkpoint,
+      custom_loader=custom_loader,
+  )
+
+  logging.info("Loading the tokenizer from: %s", checkpoint_dir)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint_dir)
+  return verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=DEFAULT_PROMPTS if prompts is None else prompts,
+      max_new_tokens=max_new_tokens,
+      atol=1e-04,
+  )

ai_edge_torch/generative/examples/openelm/openelm.py

@@ -53,7 +53,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for an OpenELM model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, enable_hlfb=True
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6
   )
   num_heads = [12] * 4 + [16] * 14 + [20] * 12 + [24] * 6
   num_query_groups = [3] * 4 + [4] * 14 + [5] * 12 + [6] * 6
@@ -101,7 +101,6 @@ 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/openelm/verify.py

@@ -15,14 +15,9 @@
 
 """Verifies the reauthored OpenELM-3B model."""
 
-import logging
-import pathlib
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.openelm import openelm
-from ai_edge_torch.generative.utilities import transformers_verifier
-from ai_edge_torch.generative.utilities import verifier
-import transformers
+from ai_edge_torch.generative.examples.openelm import verify_util
 
 
 _PROMPTS = flags.DEFINE_multi_string(
@@ -38,32 +33,10 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
 
 
 def main(_):
-  checkpoint = "apple/OpenELM-3B"
-  logging.info("Loading the original model from: %s", checkpoint)
-  original_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
-  )
-  reauthored_checkpoint = pathlib.Path(cached_config_file).parent
-  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
-  reauthored_model = openelm.build_model(str(reauthored_checkpoint))
-
-  tokenizer_checkpoint = "meta-llama/Llama-2-7b-hf"
-  logging.info("Loading the tokenizer from: %s", tokenizer_checkpoint)
-  tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_checkpoint)
-
-  verifier.verify_reauthored_model(
-      original_model=transformers_verifier.TransformersModelWrapper(
-          original_model
-      ),
-      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
-      tokenizer=verifier.TokenizerWrapper(tokenizer),
-      generate_prompts=_PROMPTS.value,
+  verify_util.verify_openelm(
+      checkpoint_dir="apple/OpenELM-3B",
       max_new_tokens=_MAX_NEW_TOKENS.value,
+      prompts=_PROMPTS.value,
   )
 
 

ai_edge_torch/generative/examples/openelm/verify_util.py

@@ -0,0 +1,76 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils for verifying the OpenELM model."""
+import logging
+import os
+import pathlib
+
+from ai_edge_torch.generative.examples.openelm import openelm
+from ai_edge_torch.generative.utilities import loader
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
+
+
+DEFAULT_PROMPTS = ["What is the meaning of life?"]
+
+
+def verify_openelm(
+    checkpoint_dir: str,
+    weight_filename: str = "model.safetensors",
+    max_new_tokens: int = 30,
+    initialize_from_local: bool = True,
+    prompts: list[str] | None = None,
+) -> bool:
+  """Verifies the reauthored OpenELM model with a custom loader."""
+  logging.info("Loading the original model from: %s", checkpoint_dir)
+  original_model = transformers.AutoModelForCausalLM.from_pretrained(
+      checkpoint_dir
+  )
+
+  logging.info("Building the reauthored model from: %s", checkpoint_dir)
+  custom_loader = (
+      None
+      if initialize_from_local
+      else loader.get_custom_loader("", "safetensors")
+  )
+
+  if initialize_from_local:
+    # Locate the cached dir.
+    cached_config_file = transformers.utils.cached_file(
+        checkpoint_dir, transformers.utils.CONFIG_NAME
+    )
+    reauthored_checkpoint = pathlib.Path(cached_config_file).parent
+  else:
+    reauthored_checkpoint = os.path.join(checkpoint_dir, weight_filename)
+
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
+  reauthored_model = openelm.build_model(
+      checkpoint_path=reauthored_checkpoint,
+      custom_loader=custom_loader,
+  )
+
+  logging.info("Loading the tokenizer from: %s", checkpoint_dir)
+  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint_dir)
+  return verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
+      tokenizer=verifier.TokenizerWrapper(tokenizer),
+      generate_prompts=DEFAULT_PROMPTS if prompts is None else prompts,
+      max_new_tokens=max_new_tokens,
+      atol=1e-04,
+  )

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -110,10 +110,7 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=16384,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -132,7 +129,6 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_use_bias=False,
-      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -93,10 +93,7 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
     The model config for the decoder of a PaliGemma 3B model.
   """
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-6,
-      zero_centered=True,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-6, zero_centered=True
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
@@ -140,7 +137,6 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       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

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -118,9 +118,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       use_bias=True,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM,
-      epsilon=1e-6,
-      enable_hlfb=True,
+      type=cfg.NormalizationType.LAYER_NORM, epsilon=1e-6
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
@@ -137,7 +135,6 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/examples/phi/phi2.py

@@ -66,9 +66,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=10240,
       use_bias=True,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.LAYER_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -85,7 +83,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       final_norm_config=norm_config,
       lm_head_use_bias=True,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/examples/phi/phi3.py

@@ -162,9 +162,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True,
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -192,7 +190,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       block_configs=block_config,
       final_norm_config=norm_config,
       lm_head_share_weight_with_embedding=False,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config

ai_edge_torch/generative/examples/phi/phi4.py

@@ -112,9 +112,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
       intermediate_size=8192,
   )
-  norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM, enable_hlfb=True
-  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,
       ff_config=ff_config,
@@ -141,7 +139,6 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       embedding_dim=3072,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
       build_rope=build_rope,
   )
   return config

ai_edge_torch/generative/examples/phi/verify.py

@@ -14,15 +14,10 @@
 # ==============================================================================
 
 """Verifies the reauthored Phi-2 model."""
-import logging
 
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.phi import phi2
-from ai_edge_torch.generative.utilities import transformers_verifier
-from ai_edge_torch.generative.utilities import verifier
-import kagglehub
-import transformers
+from ai_edge_torch.generative.examples.phi import verify_util
 
 
 _PROMPTS = flags.DEFINE_multi_string(
@@ -38,25 +33,12 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
 
 
 def main(_):
-  checkpoint = kagglehub.model_download("Microsoft/phi/transformers/2")
-  logging.info("Loading the original model from: %s", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
-
-  logging.info("Building the reauthored model from: %s", checkpoint)
-  reauthored_model = phi2.build_model(checkpoint)
-
-  logging.info("Loading the tokenizer from: %s", checkpoint)
-  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
-
-  verifier.verify_reauthored_model(
-      original_model=transformers_verifier.TransformersModelWrapper(
-          original_model
-      ),
-      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
-      tokenizer=verifier.TokenizerWrapper(tokenizer),
-      generate_prompts=_PROMPTS.value,
+  verify_util.verify_phi(
+      version="v2",
+      checkpoint_dir="microsoft/phi-2",
       max_new_tokens=_MAX_NEW_TOKENS.value,
-      atol=1e-03,
+      prompts=_PROMPTS.value,
+      atol=1e-02,
   )
 
 

ai_edge_torch/generative/examples/phi/verify_phi3.py

@@ -15,15 +15,10 @@
 
 """Verifies the reauthored Phi-3.5 model."""
 
-import logging
-import pathlib
 
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.phi import phi3
-from ai_edge_torch.generative.utilities import transformers_verifier
-from ai_edge_torch.generative.utilities import verifier
-import transformers
+from ai_edge_torch.generative.examples.phi import verify_util
 
 
 _PROMPTS = flags.DEFINE_multi_string(
@@ -39,29 +34,11 @@ _MAX_NEW_TOKENS = flags.DEFINE_integer(
 
 
 def main(_):
-  checkpoint = "microsoft/Phi-3.5-mini-instruct"
-  logging.info("Loading the original model from: %s", checkpoint)
-  original_model = transformers.AutoModelForCausalLM.from_pretrained(checkpoint)
-
-  # 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 = phi3.build_model(reauthored_checkpoint)
-
-  logging.info("Loading the tokenizer from: %s", checkpoint)
-  tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
-
-  verifier.verify_reauthored_model(
-      original_model=transformers_verifier.TransformersModelWrapper(
-          original_model
-      ),
-      reauthored_model=verifier.ReauthoredModelWrapper(reauthored_model),
-      tokenizer=verifier.TokenizerWrapper(tokenizer),
-      generate_prompts=_PROMPTS.value,
+  verify_util.verify_phi(
+      version="v3",
+      checkpoint_dir="microsoft/Phi-3.5-mini-instruct",
       max_new_tokens=_MAX_NEW_TOKENS.value,
+      prompts=_PROMPTS.value,
   )