Allow all verify_util helpers to take a custom loader object.

PiperOrigin-RevId: 765248503

Author: haozha111

ai_edge_torch/generative/examples/amd_llama_135m/verify_util.py

@@ -16,11 +16,13 @@
 import logging
 import os
 import pathlib
+from typing import Callable, Dict
 
 from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
 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 torch
 import transformers
 
 
@@ -31,8 +33,9 @@ def verify_amd_llama_135m(
     checkpoint_dir: str,
     weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
-    initialize_from_local: bool = True,
     prompts: list[str] | None = None,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored AMD-Llama-135M model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -41,11 +44,8 @@ def verify_amd_llama_135m(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/examples/deepseek/verify_util.py

@@ -16,11 +16,13 @@
 import logging
 import os
 import pathlib
+from typing import Callable, Dict
 
 from ai_edge_torch.generative.examples.deepseek import deepseek
 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 torch
 import transformers
 
 
@@ -31,8 +33,9 @@ def verify_deepseek_r1_distill_1_5b(
     checkpoint_dir: str,
     weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
-    initialize_from_local: bool = True,
     prompts: list[str] | None = None,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored DeepSeek R1 distilled 1.5B model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -41,11 +44,8 @@ def verify_deepseek_r1_distill_1_5b(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/examples/llama/verify_util.py

@@ -16,11 +16,13 @@
 import logging
 import os
 import pathlib
+from typing import Callable, Dict
 
 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 torch
 import transformers
 
 _BUILDER = {
@@ -36,8 +38,9 @@ def verify_llama_3_2(
     checkpoint_dir: str,
     weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
-    initialize_from_local: bool = True,
     prompts: list[str] | None = None,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored Llama 3.2 model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -46,11 +49,8 @@ def verify_llama_3_2(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/examples/openelm/verify_util.py

@@ -16,11 +16,13 @@
 import logging
 import os
 import pathlib
+from typing import Callable, Dict
 
 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 torch
 import transformers
 
 
@@ -31,8 +33,9 @@ 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,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored OpenELM model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -41,11 +44,8 @@ def verify_openelm(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/examples/phi/verify_util.py

@@ -13,14 +13,17 @@
 # limitations under the License.
 # ==============================================================================
 """Utils for verifying the Phi model."""
+
 import logging
 import os
 import pathlib
+from typing import Callable, Dict
 
 from ai_edge_torch.generative.examples.phi import phi2, phi3, phi4
 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 torch
 import transformers
 
 
@@ -38,9 +41,10 @@ def verify_phi(
     checkpoint_dir: str,
     weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
-    initialize_from_local: bool = True,
     prompts: list[str] | None = None,
     atol: float = 1e-04,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored Phi model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -49,11 +53,8 @@ def verify_phi(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/examples/qwen/verify_util.py

@@ -13,14 +13,17 @@
 # limitations under the License.
 # ==============================================================================
 """Utils for verifying the Qwen model."""
+
 import logging
 import os
 import pathlib
+from typing import Callable, Dict
 
 from ai_edge_torch.generative.examples.qwen import qwen, qwen3
 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 torch
 import transformers
 
 
@@ -50,8 +53,9 @@ def verify_qwen(
     checkpoint_dir: str,
     weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
-    initialize_from_local: bool = True,
     prompts: list[str] | None = None,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored Llama 3.2 model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -60,11 +64,8 @@ def verify_qwen(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/examples/smollm/verify_util.py

@@ -16,11 +16,12 @@
 import logging
 import os
 import pathlib
-
+from typing import Callable, Dict
 from ai_edge_torch.generative.examples.smollm import smollm
 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 torch
 import transformers
 
 _BUILDER = {
@@ -36,8 +37,9 @@ def verify_smollm_135m(
     checkpoint_dir: str,
     weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
-    initialize_from_local: bool = True,
     prompts: list[str] | None = None,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored SmoLLM model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -46,11 +48,9 @@ def verify_smollm_135m(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/examples/tiny_llama/verify_util.py

@@ -13,14 +13,17 @@
 # limitations under the License.
 # ==============================================================================
 """Utils for verifying the TinyLlama model."""
+
 import logging
 import os
 import pathlib
+from typing import Callable, Dict
 
 from ai_edge_torch.generative.examples.tiny_llama import tiny_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 torch
 import transformers
 
 
@@ -31,8 +34,9 @@ def verify_tiny_llama(
     checkpoint_dir: str,
     weight_filename: str = "model.safetensors",
     max_new_tokens: int = 30,
-    initialize_from_local: bool = True,
     prompts: list[str] | None = None,
+    initialize_from_local: bool = True,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] | None = None,
 ) -> bool:
   """Verifies the reauthored TinyLlama model with a custom loader."""
   logging.info("Loading the original model from: %s", checkpoint_dir)
@@ -41,11 +45,8 @@ def verify_tiny_llama(
   )
 
   logging.info("Building the reauthored model from: %s", checkpoint_dir)
-  custom_loader = (
-      None
-      if initialize_from_local
-      else loader.get_custom_loader("", "safetensors")
-  )
+  if custom_loader is None and not initialize_from_local:
+    custom_loader = loader.get_custom_loader("", "safetensors")
 
   if initialize_from_local:
     # Locate the cached dir.

ai_edge_torch/generative/utilities/loader.py

@@ -23,14 +23,13 @@
 from safetensors.torch import load_file
 import torch
 
-
 def get_custom_loader(
     checkpoint_path: str,
     checkpoint_format: Optional[str] = None,
 ) -> Callable[[str], Dict[str, torch.Tensor]]:
   """Returns a custom loader for the given checkpoint path.
 
-  Those customer loaders can either support state dictionary or safetensors, and
+  Those custome loaders can either support state dictionary or safetensors, and
   the actual data might be fetched from a remote source.
 
   Args: