ai-edge-torch-nightly 0.5.0.dev20250520__py3-none-any.whl → 0.6.0.dev20250521__py3-none-any.whl

ai_edge_torch/generative/utilities/converter.py

@@ -350,10 +350,10 @@ def _export_helper(
       )
 
   prefill_masks = None
-  if flags.FLAGS.mask_as_input:
+  if export_config.mask_as_input:
     prefill_masks = _build_mask(
-        flags.FLAGS.prefill_seq_lens,
-        flags.FLAGS.kv_cache_max_len,
+        prefill_seq_lens,
+        config.kv_cache_max_len,
         config.causal_mask_value,
     )
     if not isinstance(prefill_masks, list):
@@ -424,7 +424,7 @@ def _export_helper(
         'input_pos': decode_input_pos,
         'kv_cache': decode_kv,
     }
-    if flags.FLAGS.mask_as_input:
+    if export_config.mask_as_input:
       # Note that the decode mask is not a correct causal mask, but it is okay
       # for the conversion purpose because only the shape matters in conversion.
       # A correct causal mask of decode for a given token position of decode, it
@@ -433,7 +433,7 @@ def _export_helper(
       #  torch.triu(mask, diagonal=decode_position).unsqueeze(0).unsqueeze(0)
       #
       sample_kwargs['mask'] = _build_mask(
-          1, flags.FLAGS.kv_cache_max_len, config.causal_mask_value
+          1, config.kv_cache_max_len, config.causal_mask_value
       )
     if lora is not None:
       sample_kwargs['lora'] = lora

ai_edge_torch/generative/utilities/export_config.py

@@ -43,6 +43,9 @@ class ExportConfig:
   kvcache_cls: type = kv_utils.KVCache
   # The batch size of the decode signature.
   decode_batch_size: int = 1
+  # If true, the mask will be passed in as input. Otherwise, mask will be
+  # built by the model internally.
+  mask_as_input: bool = False
 
 
 def get_from_flags() -> ExportConfig:
@@ -51,5 +54,7 @@ def get_from_flags() -> ExportConfig:
 
   if flags.FLAGS.transpose_kv_cache:
     export_config.kvcache_layout = kv_utils.KV_LAYOUT_TRANSPOSED
+  if flags.FLAGS.mask_as_input:
+    export_config.mask_as_input = flags.FLAGS.mask_as_input
 
   return export_config

ai_edge_torch/generative/utilities/loader.py

@@ -19,8 +19,8 @@ import os
 from typing import Callable, Dict, List, Optional, Tuple
 
 from ai_edge_torch.generative.layers import model_config
-import safetensors
 from safetensors import safe_open
+from safetensors.torch import load_file
 import torch
 
 
@@ -47,7 +47,7 @@ def get_custom_loader(
 
   if checkpoint_format:
     if checkpoint_format == "safetensors":
-      return safetensors.torch.load_file
+      return load_file
     if checkpoint_format == "pt":
       return lambda path: torch.load(path, weights_only=True)
     raise ValueError(f"Unsupported checkpoint format: {checkpoint_format}")
@@ -55,7 +55,7 @@ def get_custom_loader(
   if os.path.splitext(checkpoint_path)[1] in [".bin", ".pt", ".ckpt"]:
     return lambda path: torch.load(path, weights_only=True)
   if checkpoint_path.endswith(".safetensors"):
-    return safetensors.torch.load_file
+    return load_file
   raise ValueError(f"Unsupported checkpoint format: {checkpoint_path}")
 
 

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250520"
+__version__ = "0.6.0.dev20250521"

{ai_edge_torch_nightly-0.5.0.dev20250520.dist-info → ai_edge_torch_nightly-0.6.0.dev20250521.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250520
+Version: 0.6.0.dev20250521
 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.5.0.dev20250520.dist-info → ai_edge_torch_nightly-0.6.0.dev20250521.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=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
-ai_edge_torch/version.py,sha256=4yV1q9jK9Zr0i0SQM4PpfioywnIovZWfNODUFlxFS-I,706
+ai_edge_torch/version.py,sha256=lmyCstaeVZjTAbBP4s9Z02tpX00ynyLPsymBY2tCe4A,706
 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
@@ -205,9 +205,9 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=vQWmpzMkJ2hP
 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=VRI960xo86g6lGLc_II3vDovFMa2DGIxnAZgE2GfSiM,15530
-ai_edge_torch/generative/utilities/export_config.py,sha256=5IvR3grlMd4mWO5c_Y4x9Fk1b1xa57MzlYNE8XUaN28,2049
-ai_edge_torch/generative/utilities/loader.py,sha256=ODAdOnwQXscVPiUM6ssFWqDtD-Hl-h814X1EH1c0tuw,15969
+ai_edge_torch/generative/utilities/converter.py,sha256=DuoPb8Uhbxa32uUvr6grV5lssmUJdx298QwYz8cG_1Y,15512
+ai_edge_torch/generative/utilities/export_config.py,sha256=qjkEbjcvi2AgQikZS5qfgR95Z5z9pm07KX-RN5ibfNE,2280
+ai_edge_torch/generative/utilities/loader.py,sha256=oGgEc2tHRsVqSN3mgvcngXQrpV0a7cwTpJ3LmMVnyF0,15954
 ai_edge_torch/generative/utilities/model_builder.py,sha256=tBfOcsI_NcneggHqkCSydYN3ZgmkzPc6nW0AJrA81wI,6461
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
@@ -264,8 +264,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.5.0.dev20250520.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250520.dist-info/METADATA,sha256=P4YDKSZOCPj-hx7bnU6EWLvigx-dpgHd_cIARQd4Fss,2074
-ai_edge_torch_nightly-0.5.0.dev20250520.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250520.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250520.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250521.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250521.dist-info/METADATA,sha256=_UC8q7Xe3xMUCwKKbF4CJ5hewK9PLIJ26ksKCAeWjik,2074
+ai_edge_torch_nightly-0.6.0.dev20250521.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250521.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250521.dist-info/RECORD,,