ai-edge-torch-nightly 0.7.0.dev20250919__py3-none-any.whl → 0.7.0.dev20250920__py3-none-any.whl

ai_edge_torch/generative/utilities/converter.py

@@ -18,18 +18,20 @@
 import enum
 import os
 import pathlib
-from typing import Optional, Union
+import tempfile
+from typing import Any, Optional, Union
 from absl import flags
 from ai_edge_torch._convert import converter as converter_utils
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.quantize import quant_recipes
-from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities import export_config as export_config_lib
+from ai_edge_torch.generative.utilities import litertlm_builder
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 
-ExportConfig = export_config.ExportConfig
+ExportConfig = export_config_lib.ExportConfig
 
 
 class ExportableModule(torch.nn.Module):
@@ -325,6 +327,7 @@ def convert_to_tflite(
       loras,
       export_config,
   )
+  return output_file
 
 
 def _export_helper(
@@ -457,3 +460,51 @@ def _export_helper(
       quant_config=quant_config,
   )
   edge_model.export(output_file)
+
+
+def convert_to_litert(
+    pytorch_model: torch.nn.Module,
+    output_path: str,
+    output_name_prefix: str,
+    prefill_seq_len: Union[int, list[int]],
+    kv_cache_max_len: int,
+    pixel_values_size: torch.Size = None,
+    pixel_seq_len: int = 0,
+    quantize: str = 'dynamic_int8',
+    config: cfg.ModelConfig = None,
+    lora_ranks: Optional[list[int]] = None,
+    export_config: ExportConfig = None,
+    output_format: str = 'tflite',
+    **kwargs,
+):
+  """Converts a nn.Module model to multi-signature tflite model and pack it."""
+  with tempfile.TemporaryDirectory() as workdir:
+    if output_format == 'litertlm':
+      tflite_model_output_path = workdir
+    else:
+      tflite_model_output_path = output_path
+    tflite_model_path = convert_to_tflite(
+        pytorch_model,
+        tflite_model_output_path,
+        output_name_prefix,
+        prefill_seq_len,
+        kv_cache_max_len,
+        pixel_values_size,
+        pixel_seq_len,
+        quantize,
+        config,
+        lora_ranks,
+        export_config,
+    )
+    if output_format == 'litertlm':
+      tokenizer_model_path = kwargs.pop('tokenizer_model_path', None)
+      hf_tokenizer_model_path = kwargs.pop('hf_tokenizer_model_path', None)
+      litertlm_builder.build_litertlm(
+          tflite_model_path=tflite_model_path,
+          workdir=workdir,
+          output_path=output_path,
+          context_length=kv_cache_max_len,
+          tokenizer_model_path=tokenizer_model_path,
+          hf_tokenizer_model_path=hf_tokenizer_model_path,
+          **kwargs,
+      )

ai_edge_torch/generative/utilities/litertlm_builder.py

@@ -0,0 +1,129 @@
+# 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.
+# ==============================================================================
+
+"""Utilities for building LiteRT-LM files."""
+
+
+import os
+import pathlib
+
+try:
+  # pylint: disable=g-import-not-at-top
+  from ai_edge_litert.internal import llm_metadata_pb2
+  from ai_edge_litert.internal import litertlm_builder
+  # pylint: enable=g-import-not-at-top
+
+  _litertlm_builder_available = True
+except ImportError:
+  llm_metadata_pb2 = None
+  litertlm_builder = None
+  _litertlm_builder_available = False
+
+
+def is_litertlm_builder_available() -> bool:
+  return _litertlm_builder_available
+
+
+def build_litertlm(
+    tflite_model_path: str,
+    workdir: str,
+    output_path: str,
+    context_length: int,
+    model_prompt_prefix: str | None,
+    model_prompt_suffix: str | None,
+    user_prompt_prefix: str | None,
+    user_prompt_suffix: str | None,
+    tokenizer_model_path: str | None,
+    hf_tokenizer_model_path: str | None,
+    start_token: str | None = None,
+    start_token_id: int | None = None,
+    stop_tokens: str | list[str] | None = None,
+    stop_token_ids: list[int] | None = None,
+    **kwargs,
+):
+  """Builds a LiteRT-LM file from a TFlite model."""
+  del kwargs
+
+  if not is_litertlm_builder_available():
+    raise ValueError('LiteRT-LM builder is not available.')
+  assert llm_metadata_pb2 is not None
+  assert litertlm_builder is not None
+
+  llm_metadata = llm_metadata_pb2.LlmMetadata()
+
+  if start_token_id is not None:
+    llm_metadata.start_token.token_ids.ids.append(start_token_id)
+  elif start_token is not None:
+    llm_metadata.start_token.token_str = start_token
+
+  stop_tokens_list = []
+  if stop_tokens is not None:
+    if isinstance(stop_tokens, str):
+      stop_tokens_list.append(stop_tokens)
+    else:
+      stop_tokens_list.extend(stop_tokens)
+  if stop_token_ids is not None:
+    stop_tokens_list.extend(stop_token_ids)
+  for stop_token in stop_tokens_list:
+    if isinstance(stop_token, str):
+      tu = llm_metadata.stop_tokens.add()
+      tu.token_str = stop_token
+    else:
+      assert isinstance(stop_token, int)
+      tu = llm_metadata.stop_tokens.add()
+      tu.token_ids.ids.append(stop_token)
+
+  if model_prompt_prefix is not None:
+    llm_metadata.prompt_templates.model.prefix = model_prompt_prefix
+  if model_prompt_suffix is not None:
+    llm_metadata.prompt_templates.model.suffix = model_prompt_suffix
+  if user_prompt_prefix is not None:
+    llm_metadata.prompt_templates.user.prefix = user_prompt_prefix
+  if user_prompt_suffix is not None:
+    llm_metadata.prompt_templates.user.suffix = user_prompt_suffix
+
+  llm_metadata.max_num_tokens = context_length
+
+  llm_metadata_path = os.path.join(workdir, 'llm_metadata.pb')
+  with open(llm_metadata_path, 'wb') as f:
+    f.write(llm_metadata.SerializeToString())
+
+  builder = litertlm_builder.LitertLmFileBuilder()
+  builder.add_system_metadata(
+      litertlm_builder.Metadata(
+          key='Authors',
+          value='',
+          dtype=litertlm_builder.DType.STRING,
+      )
+  )
+  builder.add_llm_metadata(llm_metadata_path)
+  if tokenizer_model_path:
+    builder.add_sentencepiece_tokenizer(tokenizer_model_path)
+  elif hf_tokenizer_model_path:
+    builder.add_hf_tokenizer(hf_tokenizer_model_path)
+  else:
+    raise ValueError(
+        'Either tokenizer_model_path or hf_tokenizer_model_path must be'
+        ' provided.'
+    )
+  builder.add_tflite_model(
+      tflite_model_path, litertlm_builder.TfLiteModelType.PREFILL_DECODE
+  )
+
+  file_name = pathlib.Path(tflite_model_path).stem + '.litertlm'
+  if os.path.dirname(file_name):
+    os.makedirs(os.path.dirname(file_name), exist_ok=True)
+  with open(os.path.join(output_path, file_name), 'wb') as f:
+    builder.build(f)

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.7.0.dev20250919"
+__version__ = "0.7.0.dev20250920"

{ai_edge_torch_nightly-0.7.0.dev20250919.dist-info → ai_edge_torch_nightly-0.7.0.dev20250920.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.7.0.dev20250919
+Version: 0.7.0.dev20250920
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.7.0.dev20250919.dist-info → ai_edge_torch_nightly-0.7.0.dev20250920.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=i9nybsJyTD-Fo-H_Taf-7IBN4Dvc9BuLdb-dlAA0_y0,806
+ai_edge_torch/version.py,sha256=AMujyxObvKlUOBjKtYZjM5FWraCbXiQmR9pEr_hpKI4,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -208,8 +208,9 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=NkEwrjO8vIcd
 ai_edge_torch/generative/test/test_quantize.py,sha256=kKJ01wscTC2t_Ylr7huO5gNKES01gm3dT1gx52z15PA,7356
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=TKHUuC0b6swR3dYyijLz_IDCdmCKQJUT-nvXFO5M-0U,15885
+ai_edge_torch/generative/utilities/converter.py,sha256=Q2C9DqYWqiyuu8KoL69SAxDL-RNJ7K8_Z45H2FpsAn0,17563
 ai_edge_torch/generative/utilities/export_config.py,sha256=qjkEbjcvi2AgQikZS5qfgR95Z5z9pm07KX-RN5ibfNE,2280
+ai_edge_torch/generative/utilities/litertlm_builder.py,sha256=0cNuaqhc7cQcAa4NRalUXyoPQUQC9O3-aHAJEDV1Mps,4265
 ai_edge_torch/generative/utilities/loader.py,sha256=drgKBmNibuc3PCdc0kU0pVcp2Nt1_mjLYh67RyXOn7U,15952
 ai_edge_torch/generative/utilities/model_builder.py,sha256=xBvcTxihB9TN88UtQiXA9sAITQgf-pA77R-VZlLgUeU,6950
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
@@ -269,8 +270,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.7.0.dev20250919.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20250919.dist-info/METADATA,sha256=GWs1_9J5oxQPmSFJwW3lrHLNW8qMHLX5QnjM1cITwkI,2074
-ai_edge_torch_nightly-0.7.0.dev20250919.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.7.0.dev20250919.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20250919.dist-info/RECORD,,
+ai_edge_torch_nightly-0.7.0.dev20250920.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.7.0.dev20250920.dist-info/METADATA,sha256=IReTV7SfbskOGXNOj3XypKPULQxj4aIKg6QMaxGDMaY,2074
+ai_edge_torch_nightly-0.7.0.dev20250920.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.7.0.dev20250920.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.7.0.dev20250920.dist-info/RECORD,,