Add verify_util to run model numerical validation with a custom loader.

PiperOrigin-RevId: 759916326

Author: ai-edge-bot

ai_edge_torch/generative/examples/amd_llama_135m/verify.py

@@ -15,9 +15,15 @@
 
 """Verifies the reauthored AMD-Llama-135M model."""
 
+import logging
+import pathlib
+
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.amd_llama_135m import verify_util
+from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
 
 
 _PROMPTS = flags.DEFINE_multi_string(
@@ -33,10 +39,32 @@
 
 
 def main(_):
-  verify_util.verify_amd_llama_135m(
-      "amd/AMD-Llama-135m",
+  checkpoint = "amd/AMD-Llama-135m"
+  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 = amd_llama_135m.build_model(str(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,
       max_new_tokens=_MAX_NEW_TOKENS.value,
-      prompts=_PROMPTS.value,
+      atol=1e-04,
   )
 
 

ai_edge_torch/generative/examples/amd_llama_135m/verify_util.py

@@ -15,9 +15,15 @@
 
 """Verifies the reauthored DeepSeek R1 distilled 1.5B model."""
 
+import logging
+import pathlib
+
 from absl import app
 from absl import flags
-from ai_edge_torch.generative.examples.deepseek import verify_util
+from ai_edge_torch.generative.examples.deepseek import deepseek
+from ai_edge_torch.generative.utilities import transformers_verifier
+from ai_edge_torch.generative.utilities import verifier
+import transformers
 
 
 _PROMPTS = flags.DEFINE_multi_string(
@@ -33,10 +39,30 @@
 
 
 def main(_):
-  verify_util.verify_deepseek_r1_distill_1_5b(
-      "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
+  checkpoint = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+  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 = deepseek.build_model(str(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,
       max_new_tokens=_MAX_NEW_TOKENS.value,
-      prompts=_PROMPTS.value,
+      atol=1e-04,
   )
 
 

ai_edge_torch/generative/examples/deepseek/verify.py

@@ -17,13 +17,11 @@
 
 import logging
 import os
-from typing import Callable, Dict, List, Tuple
+from typing import List, Tuple
 
-from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.examples.gemma import gemma2
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
-from ai_edge_torch.generative.utilities import loader
 from ai_edge_torch.generative.utilities import verifier
 from gemma import config as gemma_config
 from gemma import model as gemma_model
@@ -109,7 +107,6 @@ def verify_reauthored_gemma_model(
     generate_prompts: List[str],
     forward_input_ids: List[List[int]],
     weight_filename: str = "model.ckpt",
-    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
     tokenizer_filename: str = "tokenizer.model",
     max_new_tokens: int = 20,
     mask_as_input: bool = False,
@@ -128,14 +125,7 @@ def verify_reauthored_gemma_model(
 
   logging.info("Loading the original model from: %s", checkpoint)
   original_model = gemma_model.GemmaForCausalLM(config).eval()
-  checkpoint_path = os.path.join(checkpoint, weight_filename)
-  if custom_loader is None:
-    original_model.load_weights(checkpoint_path)
-  else:
-    original_model.load_state_dict(
-        custom_loader(checkpoint_path)["model_state_dict"],
-        strict=False,
-    )
+  original_model.load_weights(os.path.join(checkpoint, weight_filename))
 
   return verifier.verify_reauthored_model(
       original_model=GemmaWrapper(original_model),
@@ -154,62 +144,27 @@ def verify_reauthored_gemma_model(
 
 
 def verify_gemma2(
-    checkpoint_dir: str,
-    weight_filename: str,
+    gemma2_model_path: str,
     prompts: List[str],
     max_new_tokens: int,
     mask_as_input: bool = False,
     kv_layout: kv_utils.KVLayout = kv_utils.KV_LAYOUT_DEFAULT,
-    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored Gemma2 model.
 
   Return True if the verification passes, False otherwise.
   """
-  checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
-  logging.info("Building the reauthored model from: %s", checkpoint_path)
-  reauthored_model = gemma2.build_2b_model(checkpoint_path, custom_loader)
+  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=checkpoint_dir,
+      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]],
-      weight_filename=weight_filename,
-      custom_loader=custom_loader,
       max_new_tokens=max_new_tokens,
       mask_as_input=mask_as_input,
       kv_layout=kv_layout,
       atol=1e-04,
   )
-
-
-def verify_gemma1_with_custom_loader(checkpoint_dir: str) -> bool:
-  """Verifies the reauthored Gemma1 model with a custom loader."""
-  weight_filename = "gemma-2b-it.ckpt"
-  checkpoint_path = os.path.join(checkpoint_dir, weight_filename)
-  custom_loader = loader.get_custom_loader(checkpoint_path)
-  reauthored_model = gemma1.build_2b_model(checkpoint_path, custom_loader)
-  return verify_reauthored_gemma_model(
-      checkpoint=checkpoint_dir,
-      variant="2b",
-      reauthored_model=reauthored_model,
-      weight_filename=weight_filename,
-      custom_loader=custom_loader,
-      generate_prompts=["What is the meaning of life?"],
-      forward_input_ids=[[1, 2, 3, 4]],
-      max_new_tokens=30,
-  )
-
-
-def verify_gemma2_with_custom_loader(checkpoint_dir: str) -> bool:
-  """Verifies the reauthored Gemma2 model with a custom loader."""
-  return verify_gemma2(
-      checkpoint_dir=checkpoint_dir,
-      weight_filename="model.ckpt",
-      prompts=["What is the meaning of life?"],
-      max_new_tokens=30,
-      mask_as_input=True,
-      custom_loader=loader.get_custom_loader("", checkpoint_format="pt"),
-  )

ai_edge_torch/generative/examples/deepseek/verify_util.py

@@ -22,7 +22,6 @@
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
-from ai_edge_torch.generative.utilities import loader
 from ai_edge_torch.generative.utilities import verifier
 from gemma import config as gemma_config
 from gemma import model as gemma_model
@@ -261,15 +260,3 @@ def verify_gemma3(
       custom_loader=custom_loader,
       atol=1e-04,
   )
-
-
-def verify_gemma3_with_custom_loader(checkpoint: str) -> bool:
-  """Verifies the reauthored Gemma3 model with a custom loader."""
-  return verify_gemma3(
-      checkpoint=checkpoint,
-      prompts=["What is the meaning of life?"],
-      max_new_tokens=30,
-      variant="1b",
-      weight_filename="model.ckpt",
-      custom_loader=loader.get_custom_loader("", checkpoint_format="pt"),
-  )