* Move logging funcions to seperate file.

* Add some missing BUILD rules.

PiperOrigin-RevId: 678962505

Author: haozha111

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -16,7 +16,6 @@
 """Example of building a Gemma2 model."""
 
 import os
-import pathlib
 from typing import Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
@@ -25,7 +24,6 @@
 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
 

ai_edge_torch/generative/examples/gemma/verify_gemma1.py

@@ -15,13 +15,14 @@
 
 """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
-from ai_edge_torch.generative.utilities import verifier
 import kagglehub
 
+
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
     "What is the meaning of life?",
@@ -37,7 +38,7 @@
 def main(_):
   checkpoint = kagglehub.model_download("google/gemma/pyTorch/2b-it")
 
-  verifier.log_msg("Building the reauthored model from", checkpoint)
+  logging.info("Building the reauthored model from: %s", checkpoint)
   reauthored_model = gemma1.build_2b_model(checkpoint)
 
   verify_util.verify_reauthored_gemma_model(

ai_edge_torch/generative/examples/gemma/verify_gemma2.py

@@ -15,13 +15,15 @@
 
 """Verifies the reauthored Gemma2 model."""
 
+import logging
 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?",
@@ -37,7 +39,7 @@
 def main(_):
   checkpoint = kagglehub.model_download("google/gemma-2/pyTorch/gemma-2-2b-it")
 
-  verifier.log_msg("Building the reauthored model from", checkpoint)
+  logging.info("Building the reauthored model from: %s", checkpoint)
   reauthored_model = gemma2.build_2b_model(checkpoint)
 
   verify_util.verify_reauthored_gemma_model(

ai_edge_torch/generative/examples/gemma/verify_util.py

@@ -16,6 +16,7 @@
 """Utility functions to verify the reauthored Gemma model."""
 
 import dataclasses
+import logging
 import os
 from typing import List, Tuple
 
@@ -127,7 +128,7 @@ def verify_reauthored_gemma_model(
   # Use float32 to be compatible with the reauthored model.
   config.dtype = torch.float32
 
-  verifier.log_msg("Loading the original model from", checkpoint)
+  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))
 

ai_edge_torch/generative/examples/openelm/verify.py

@@ -15,14 +15,15 @@
 
 """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 verifier
 import transformers
 
+
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
     "What is the meaning of life?",
@@ -32,7 +33,7 @@
 
 def main(_):
   checkpoint = "apple/OpenELM-3B"
-  verifier.log_msg("Loading the original model from", checkpoint)
+  logging.info("Loading the original model from: %s", checkpoint)
   wrapper_model = verifier.ModelWrapper(
       model=transformers.AutoModelForCausalLM.from_pretrained(
           checkpoint, trust_remote_code=True
@@ -44,11 +45,11 @@ def main(_):
       checkpoint, transformers.utils.CONFIG_NAME
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
-  verifier.log_msg("Building the reauthored model from", reauthored_checkpoint)
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = openelm.build_model(reauthored_checkpoint)
 
   tokenizer_checkpoint = "meta-llama/Llama-2-7b-hf"
-  verifier.log_msg("Loading the tokenizer from", tokenizer_checkpoint)
+  logging.info("Loading the tokenizer from: %s", tokenizer_checkpoint)
   tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_checkpoint)
 
   verifier.verify_reauthored_model(

ai_edge_torch/generative/examples/phi/verify.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Verifies the reauthored Phi-2 model."""
+import logging
 
 from absl import app
 from absl import flags
@@ -22,6 +23,7 @@
 import kagglehub
 import transformers
 
+
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
     "Instruct: Write an email about the weather Output:",
@@ -36,18 +38,18 @@
 
 def main(_):
   checkpoint = kagglehub.model_download("Microsoft/phi/transformers/2")
-  verifier.log_msg("Loading the original model from", checkpoint)
+  logging.info("Loading the original model from: %s", 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)
+  logging.info("Building the reauthored model from: %s", checkpoint)
   reauthored_model = phi2.build_model(checkpoint)
 
-  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  logging.info("Loading the tokenizer from: %s", checkpoint)
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(

ai_edge_torch/generative/examples/phi/verify_phi3.py

@@ -15,6 +15,7 @@
 
 """Verifies the reauthored Phi-3.5 model."""
 
+import logging
 import pathlib
 
 from absl import app
@@ -23,6 +24,7 @@
 from ai_edge_torch.generative.utilities import verifier
 import transformers
 
+
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
     "Instruct: Write an email about the weather Output:",
@@ -37,7 +39,7 @@
 
 def main(_):
   checkpoint = "microsoft/Phi-3.5-mini-instruct"
-  verifier.log_msg("Loading the original model from", checkpoint)
+  logging.info("Loading the original model from: %s", checkpoint)
   generation_config = transformers.GenerationConfig.from_pretrained(checkpoint)
   generation_config.max_new_tokens = _MAX_NEW_TOKENS.value
   wrapper_model = verifier.ModelWrapper(
@@ -50,10 +52,10 @@ def main(_):
       checkpoint, transformers.utils.CONFIG_NAME
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
-  verifier.log_msg("Building the reauthored model from", reauthored_checkpoint)
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = phi3.build_model(reauthored_checkpoint)
 
-  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  logging.info("Loading the tokenizer from: %s", checkpoint)
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(

ai_edge_torch/generative/examples/smollm/verify.py

@@ -15,6 +15,7 @@
 
 """Verifies the reauthored SmolLM-135M model."""
 
+import logging
 import pathlib
 
 from absl import app
@@ -23,6 +24,7 @@
 from ai_edge_torch.generative.utilities import verifier
 import transformers
 
+
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
     "What is the meaning of life?",
@@ -32,7 +34,7 @@
 
 def main(_):
   checkpoint = "HuggingFaceTB/SmolLM-135M"
-  verifier.log_msg("Loading the original model from", checkpoint)
+  logging.info("Loading the original model from: %s", checkpoint)
   wrapper_model = verifier.ModelWrapper(
       model=transformers.AutoModelForCausalLM.from_pretrained(checkpoint),
   )
@@ -41,10 +43,10 @@ def main(_):
       checkpoint, transformers.utils.CONFIG_NAME
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
-  verifier.log_msg("Building the reauthored model from", reauthored_checkpoint)
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = smollm.build_model(reauthored_checkpoint)
 
-  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  logging.info("Loading the tokenizer from: %s", checkpoint)
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(

ai_edge_torch/generative/examples/tiny_llama/verify.py

@@ -15,6 +15,7 @@
 
 """Verifies the reauthored TinyLlama-1.1B model."""
 
+import logging
 import pathlib
 
 from absl import app
@@ -23,6 +24,7 @@
 from ai_edge_torch.generative.utilities import verifier
 import transformers
 
+
 _PROMPTS = flags.DEFINE_multi_string(
     "prompts",
     "Show me the program to add 2 and 3.",
@@ -32,7 +34,7 @@
 
 def main(_):
   checkpoint = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-  verifier.log_msg("Loading the original model from", checkpoint)
+  logging.info("Loading the original model from: %s", checkpoint)
   wrapper_model = verifier.ModelWrapper(
       model=transformers.AutoModelForCausalLM.from_pretrained(
           checkpoint, trust_remote_code=True
@@ -43,10 +45,10 @@ def main(_):
       checkpoint, transformers.utils.CONFIG_NAME
   )
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
-  verifier.log_msg("Building the reauthored model from", reauthored_checkpoint)
+  logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = tiny_llama.build_model(reauthored_checkpoint)
 
-  verifier.log_msg("Loading the tokenizer from", checkpoint)
+  logging.info("Loading the tokenizer from: %s", checkpoint)
   tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
 
   verifier.verify_reauthored_model(

ai_edge_torch/generative/tools/batch_convert.py

@@ -17,6 +17,7 @@
 
 import dataclasses
 import enum
+import logging
 import os
 import pathlib
 from typing import Callable, Sequence
@@ -29,7 +30,6 @@
 from ai_edge_torch.generative.examples.phi import phi3
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import verifier
 import torch
 
 _CHECKPOINT_ROOT_PATH = flags.DEFINE_string(
@@ -67,12 +67,12 @@ class ConversionConfig:
 
   def print_config(self) -> None:
     """Prints the conversion config."""
-    verifier.log_msg("Model name:", self.model_name)
-    verifier.log_msg("Input checkpoint:", self.input_checkpoint)
-    verifier.log_msg("TF Lite output path:", self.tflite_output_path)
-    verifier.log_msg("Prefill seq len:", self.prefill_seq_len)
-    verifier.log_msg("KV cache max len:", self.kv_cache_max_len)
-    verifier.log_msg("Export precision:", self.export_precision)
+    logging.info("Model name: %s", self.model_name)
+    logging.info("Input checkpoint: %s", self.input_checkpoint)
+    logging.info("TF Lite output path: %s", self.tflite_output_path)
+    logging.info("Prefill seq len: %s", self.prefill_seq_len)
+    logging.info("KV cache max len: %s", self.kv_cache_max_len)
+    logging.info("Export precision: %s", self.export_precision)
 
 
 def prepare_conversion_configs() -> Sequence[ConversionConfig]:
@@ -143,7 +143,7 @@ def get_output_filename(
   if precision == ExportPrecision.INT8:
     precision_str = "q8"
   elif precision == ExportPrecision.FP32:
-    precision_str = "fp32"
+    precision_str = "f32"
   else:
     raise ValueError(f"Unsupported precision: {precision}")
   return f"{model_name}_{precision_str}_seq{prefill_seq_len}_ekv{kv_cache_max_len}.tflite"
@@ -152,8 +152,8 @@ def get_output_filename(
 def convert_models(conversion_configs: Sequence[ConversionConfig]) -> None:
   """Executes the conversion for a batch of models specified by the `conversion_configs`."""
   for config in conversion_configs:
-    verifier.log_msg(
-        "Converting model:", config.model_name, " with the following config:"
+    logging.info(
+        "Converting model: %s with the following config:", config.model_name
     )
     config.print_config()
     pytorch_model = config.model_builder(
@@ -172,7 +172,7 @@ def convert_models(conversion_configs: Sequence[ConversionConfig]) -> None:
           prefill_seq_len=config.prefill_seq_len,
           quantize=True if precision == ExportPrecision.INT8 else False,
       )
-      verifier.log_msg("Successfully converted model:", output_filename)
+      logging.info("Successfully converted model: %s", output_filename)
 
 
 def main(_):