ai-edge-torch-nightly 0.5.0.dev20250413__py3-none-any.whl → 0.5.0.dev20250415__py3-none-any.whl

ai_edge_torch/generative/utilities/converter.py

@@ -57,7 +57,7 @@ def define_conversion_flags(model_name: str):
   )
   flags.DEFINE_string(
       'output_name_prefix',
-      'qwen',
+      f'{model_name}',
       'The prefix of the output tflite model name.',
   )
   flags.DEFINE_multi_integer(

ai_edge_torch/generative/utilities/verifier.py

@@ -181,7 +181,7 @@ def verify_with_input_ids(
     original_model: ModelWrapper,
     reauthored_model: ReauthoredModelWrapper,
     input_ids: List[int],
-    kv_cache_max_len: int = 1024,
+    kv_cache_max_len: int = 128,
     rtol: float = 1e-05,
     atol: float = 1e-05,
 ):
@@ -273,6 +273,8 @@ def verify_reauthored_model(
     rtol: float = 1e-05,
     atol: float = 1e-05,
     continue_on_failure: bool = False,
+    verify_inputs: bool = True,
+    verify_prompts: bool = True,
 ) -> bool:
   """Verifies the reauthored model against the original model.
 
@@ -301,33 +303,37 @@ def verify_reauthored_model(
   """
   failure_count = 0
 
-  for input_ids in forward_input_ids:
-    logging.info("Verifying the reauthored model with input IDs: %s", input_ids)
-    try:
-      verify_with_input_ids(
-          original_model, reauthored_model, input_ids, rtol=rtol, atol=atol
+  if verify_inputs:
+    for input_ids in forward_input_ids:
+      logging.info(
+          "Verifying the reauthored model with input IDs: %s", input_ids
       )
-    except AssertionError as e:
-      logging.error("*** FAILED *** verify with input IDs: %s", input_ids)
-      failure_count += 1
-      if not continue_on_failure:
-        return False
-    else:
-      logging.info("*** PASSED *** verify with input IDs: %s", input_ids)
-
-  for prompts in generate_prompts:
-    logging.info("Verifying the reauthored model with prompts: %s", prompts)
-    try:
-      verify_model_with_prompts(
-          original_model, reauthored_model, tokenizer, prompts, max_new_tokens
-      )
-    except AssertionError as e:
-      logging.error("*** FAILED *** verify with prompts: %s", prompts)
-      failure_count += 1
-      if not continue_on_failure:
-        return False
-    else:
-      logging.info("*** PASSED *** verify with prompts: %s", prompts)
+      try:
+        verify_with_input_ids(
+            original_model, reauthored_model, input_ids, rtol=rtol, atol=atol
+        )
+      except AssertionError as e:
+        logging.error("*** FAILED *** verify with input IDs: %s", input_ids)
+        failure_count += 1
+        if not continue_on_failure:
+          return False
+      else:
+        logging.info("*** PASSED *** verify with input IDs: %s", input_ids)
+
+  if verify_prompts:
+    for prompts in generate_prompts:
+      logging.info("Verifying the reauthored model with prompts: %s", prompts)
+      try:
+        verify_model_with_prompts(
+            original_model, reauthored_model, tokenizer, prompts, max_new_tokens
+        )
+      except AssertionError as e:
+        logging.error("*** FAILED *** verify with prompts: %s", prompts)
+        failure_count += 1
+        if not continue_on_failure:
+          return False
+      else:
+        logging.info("*** PASSED *** verify with prompts: %s", prompts)
 
   if failure_count == 0:
     logging.info("*** PASSED *** verify_reauthored_model")

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250413"
+__version__ = "0.5.0.dev20250415"

{ai_edge_torch_nightly-0.5.0.dev20250413.dist-info → ai_edge_torch_nightly-0.5.0.dev20250415.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250413
+Version: 0.5.0.dev20250415
 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.dev20250413.dist-info → ai_edge_torch_nightly-0.5.0.dev20250415.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=UZYJ3ZA4AyEgqcgQ3Q3zennjEkHbpYJ1WBQ_XxZt2-o,706
+ai_edge_torch/version.py,sha256=F_xrSGW8cTj5_LRehy1awbLrOMYd-sOQPyVRXV5mNNI,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=GPDsXhfECjDzOut4vh_d9qWcyfpxobFMBTsC7MyJbM0,5557
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -186,7 +186,7 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=-v2Vj7Qdd3Gy
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 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=87Tzj-gLydx8_xnHxKlCbMmM1XHShstpKi8RH3xY7Xw,9757
+ai_edge_torch/generative/utilities/converter.py,sha256=swtz69oyMOxSaCEYST_Gzd5sjGZ1qOBAfd_0xl207Nk,9766
 ai_edge_torch/generative/utilities/export_config.py,sha256=8-795nyd3M34LkGhgW7hwHlJyTc2Oz1iipHK8yBhdFs,1633
 ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
 ai_edge_torch/generative/utilities/model_builder.py,sha256=ZYX1TxpFdj573du2QCyHJlFjx4q1m12R74fp4Gwl92A,6343
@@ -195,7 +195,7 @@ ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWt
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
 ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
 ai_edge_torch/generative/utilities/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
-ai_edge_torch/generative/utilities/verifier.py,sha256=gsBGv-WgeP73Bmao29CiTIQOC1YZ43IUJcGzytpcZyM,12095
+ai_edge_torch/generative/utilities/verifier.py,sha256=RSMQ8eda63VHM-5KmquKfogmTPyhGvGnqkoz9i4bppY,12270
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=JsVmYrM_JEuN_smMHXUsRlo3Liapp7UyktbPpPARwDk,5386
 ai_edge_torch/hlfb/mark_pattern/fx_utils.py,sha256=YCtMgu-4w2BQ5fpnlpWC6IauKPf_tVqc7Ff91OTqlSw,1796
@@ -245,8 +245,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 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.dev20250413.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250413.dist-info/METADATA,sha256=cXJe6i4sOy9shfnsOge_0KwXzOMaHPHMSfIlYukSm6Q,2051
-ai_edge_torch_nightly-0.5.0.dev20250413.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250413.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250413.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250415.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250415.dist-info/METADATA,sha256=hwDAXpg_N8uBvl134iezzps8jSfwgjs1Q3zrX0As4IM,2051
+ai_edge_torch_nightly-0.5.0.dev20250415.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250415.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250415.dist-info/RECORD,,