ai-edge-torch-nightly 0.3.0.dev20241002__py3-none-any.whl → 0.3.0.dev20241003__py3-none-any.whl

ai_edge_torch/generative/examples/llama/convert_3b_to_tflite.py

@@ -1,68 +0,0 @@
-# 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 converting Llama 3.2 3B model to multi-signature tflite model."""
-
-import os
-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 converter
-
-_CHECKPOINT_PATH = flags.DEFINE_string(
-    'checkpoint_path',
-    os.path.join(pathlib.Path.home(), 'Downloads/llm_data/llama'),
-    'The path to the model checkpoint, or directory holding the checkpoint.',
-)
-_TFLITE_PATH = flags.DEFINE_string(
-    'tflite_path',
-    '/tmp/',
-    'The tflite file path to export.',
-)
-_PREFILL_SEQ_LEN = flags.DEFINE_integer(
-    'prefill_seq_len',
-    1024,
-    'The maximum size of prefill input tensor.',
-)
-_KV_CACHE_MAX_LEN = flags.DEFINE_integer(
-    'kv_cache_max_len',
-    1280,
-    'The maximum size of KV cache buffer, including both prefill and decode.',
-)
-_QUANTIZE = flags.DEFINE_bool(
-    'quantize',
-    True,
-    'Whether the model should be quantized.',
-)
-
-
-def main(_):
-  pytorch_model = llama.build_3b_model(
-      _CHECKPOINT_PATH.value, kv_cache_max_len=_KV_CACHE_MAX_LEN.value
-  )
-  quant_suffix = 'q8' if _QUANTIZE.value else 'f32'
-  output_filename = f'llama_3b_{quant_suffix}_seq{_PREFILL_SEQ_LEN.value}_ekv{_KV_CACHE_MAX_LEN.value}.tflite'
-  converter.convert_to_tflite(
-      pytorch_model,
-      tflite_path=os.path.join(_TFLITE_PATH.value, output_filename),
-      prefill_seq_len=_PREFILL_SEQ_LEN.value,
-      quantize=_QUANTIZE.value,
-  )
-
-
-if __name__ == '__main__':
-  app.run(main)

ai_edge_torch/generative/examples/llama/verify_3b.py

@@ -1,73 +0,0 @@
-# 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 Llama 3.2-3B 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
-
-
-_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 = "meta-llama/Llama-3.2-3B-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 = llama.build_3b_model(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,
-      max_new_tokens=_MAX_NEW_TOKENS.value,
-      atol=1e-04,
-  )
-
-
-if __name__ == "__main__":
-  app.run(main)