[LiteRT Torch] End-to-end NPU export pipeline, configuration manager, and master notebook. Upgrade calibration, loader, and quantization to use PreTrainedTokenizerFast and LitertLmBundle.

PiperOrigin-RevId: 958581673

Author: YunanAZ

litert_torch/generative/export_hf/experimental/calib/calibrate.py

@@ -14,9 +14,11 @@
 # ==============================================================================
 """Loader for models."""
 
+import tempfile
 from typing import Tuple
 from litert_torch.generative.export_hf.experimental.calib import sampling_executor as tfl_sampling_executor
 from litert_torch.generative.export_hf.experimental.calib import tokenizer as tokenizer_lib
+from litert_torch.generative.export_hf.experimental.litertlm_bundle import litertlm_bundle as litertlm_utils
 
 
 def _infer_model_specs(
@@ -42,8 +44,8 @@ def _infer_drafter_step(decode_model_path: str):
 
 
 def load_models(
-    model_path: str | Tuple[str, str],
-    embedder_model_path: str,
+    model_path: str | Tuple[str, str] | None,
+    embedder_model_path: str | None,
     spm_path: str | None,
     transformers_model_path: str | None,
     max_kv_cache_size: int | None,
@@ -58,8 +60,48 @@ def load_models(
     enable_min_max_calibration_update: bool = True,
     ema_smoothing_factor: float = 0.95,
     use_profiler_based_calibration: bool = False,
+    input_litertlm: str | None = None,
 ) -> tfl_sampling_executor.TflSamplingExecutorConfig:
   """Loads the models."""
+  unpacked = None
+  if input_litertlm or (
+      isinstance(model_path, str) and model_path.endswith('.litertlm')
+  ):
+    bundle_path = input_litertlm or model_path
+    assert isinstance(bundle_path, str)
+    unpack_dir = tempfile.mkdtemp(prefix='litertlm_unpacked_')
+    unpacked = litertlm_utils.unpack_litertlm(bundle_path, unpack_dir)
+    if (
+        not model_path
+        or model_path == (None, None)
+        or (isinstance(model_path, str) and model_path == bundle_path)
+        or (
+            isinstance(model_path, tuple)
+            and model_path[0] is None
+            and model_path[1] is None
+        )
+    ):
+      model_path = unpacked.get('tf_lite_prefill_decode')
+    if not embedder_model_path:
+      embedder_model_path = unpacked.get('tf_lite_embedder')
+    if not auxiliary_model_path:
+      auxiliary_model_path = unpacked.get('tf_lite_aux')
+    if not per_layer_embedder_model_path:
+      per_layer_embedder_model_path = unpacked.get('tf_lite_per_layer_embedder')
+    if not spm_path and not transformers_model_path:
+      spm_path = unpacked.get('SP_Tokenizer')
+      transformers_model_path = unpacked.get('transformers_model_path')
+
+  if not model_path:
+    raise ValueError(
+        'Must specify model_path or provide a valid input_litertlm bundle.'
+    )
+  if not embedder_model_path:
+    raise ValueError(
+        'Must specify embedder_model_path or include embedder in'
+        ' input_litertlm.'
+    )
+
   if isinstance(model_path, tuple):
     prefill_model_path, decode_model_path = model_path
     decode_model_path = decode_model_path or prefill_model_path
@@ -158,7 +200,11 @@ def load_models(
   )
   tokenizer_config = tokenizer_lib.TokenizerConfig(
       vocab_path=spm_path,
-      transformers_model_path=transformers_model_path
+      transformers_model_path=transformers_model_path if not unpacked else None,
+      tokenizer_json_path=unpacked.get('tokenizer_json_path')
+      if unpacked
+      else None,
+      chat_template=unpacked.get('chat_template') if unpacked else None,
   )
   return tfl_sampling_executor.TflSamplingExecutorConfig(
       prefill_model_entries=prefill_model_entries,
@@ -181,4 +227,5 @@ def load_models(
       enable_min_max_calibration_update=enable_min_max_calibration_update,
       ema_smoothing_factor=ema_smoothing_factor,
       use_profiler_based_calibration=use_profiler_based_calibration,
+      stop_tokens=unpacked.get('stop_token_ids') if unpacked else None,
   )

litert_torch/generative/export_hf/experimental/calib/loader.py

@@ -107,7 +107,25 @@ def get_example_prompt(
     tokenizer: tokenizer_lib.Tokenizer | None = None,
 ) -> str | tokenizer_lib.Request:
   """Gets the prompt from the example."""
+  if isinstance(example, str):
+    prompt = example
+    if enable_formatting:
+      prompt = PROMPT_TEMPLATE_PREFIX + prompt + PROMPT_TEMPLATE_SUFFIX
+    return prompt
+
   if isinstance(example, dict):
+    if 'text' in example and isinstance(example['text'], str):
+      prompt = example['text']
+      if enable_formatting:
+        prompt = PROMPT_TEMPLATE_PREFIX + prompt + PROMPT_TEMPLATE_SUFFIX
+      return prompt
+
+    if 'prompt' in example and isinstance(example['prompt'], str):
+      prompt = example['prompt']
+      if enable_formatting:
+        prompt = PROMPT_TEMPLATE_PREFIX + prompt + PROMPT_TEMPLATE_SUFFIX
+      return prompt
+
     if 'messages' in example:
       user_messages = [
           msg for msg in example['messages'] if msg.get('role') == 'user'
@@ -118,6 +136,7 @@ def get_example_prompt(
             tokenize=False,
             add_generation_prompt=True,
         )
+        assert isinstance(prompt, str)
         print(f'\n--- Formatted prompt using chat template:\n{prompt}\n---')
         return prompt
       else: