litert-torch-nightly 0.8.0.dev20260126__py3-none-any.whl → 0.9.0.dev20260127__py3-none-any.whl

litert_torch/cli.py

@@ -16,8 +16,8 @@
 
 # This is experimental and subject to change.
 
-from litert_torch.generative.export_hf import export as hf_export_lib
 import fire
+from litert_torch.generative.export_hf import export as hf_export_lib
 
 
 class CLI:
@@ -26,7 +26,7 @@ class CLI:
     self.hf_export = hf_export_lib.export
 
 
-def main(_):
+def main():
   fire.Fire(CLI())
 
 

litert_torch/generative/export_hf/core/export_lib.py

@@ -16,6 +16,8 @@
 
 import os
 import time
+
+import huggingface_hub
 from litert_torch import fx_infra
 from litert_torch._convert import converter as converter_utils
 from litert_torch.generative.export_hf.core import attention as _
@@ -86,7 +88,7 @@ def load_model(
 
   config = transformers.AutoConfig.from_pretrained(
       model_path,
-      torch_dtype=torch.float32,
+      dtype=torch.float32,
       trust_remote_code=trust_remote_code,
   )
   config._attn_implementation = 'lrt_transposed_attention'  # pylint: disable=protected-access
@@ -111,7 +113,20 @@ def load_model(
 
   verify_model_compatibility(model, config, text_model_config)
 
+  # TODO(weiyiw): Refactor into a separate function.
   tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)
+  if not hasattr(tokenizer, 'chat_template') or not tokenizer.chat_template:
+    try:
+      if utils.get_model_path_type(model_path) == 'repo_id':
+        template_file = huggingface_hub.hf_hub_download(
+            model_path, filename='chat_template.json'
+        )
+      else:
+        template_file = os.path.join(model_path, 'chat_template.json')
+      with open(template_file, 'rt') as f:
+        tokenizer.chat_template = f.read()
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      print(f'Failed to load chat template: {e}')
 
   return model, config, text_model_config, tokenizer
 

litert_torch/generative/export_hf/core/litert_lm_builder.py

@@ -27,7 +27,7 @@ _PH = 'KIMAIRA'
 def parse_chat_template(tokenizer):
   """Parses chat template."""
   if tokenizer.chat_template is None:
-    return (None, None), (None, None), (None, None)
+    return None
   try:
     messages = [
         {'role': 'system', 'content': _PH},
@@ -39,6 +39,10 @@ def parse_chat_template(tokenizer):
         add_generation_prompt=False,
     )
     sys_prompt_parts = sys_prompt.split(_PH)
+    no_sys_prompt = False
+    if len(sys_prompt_parts) == 1:
+      sys_prompt_parts = [sys_prompt_parts[0], '']
+      no_sys_prompt = True
     if len(sys_prompt_parts) != 2:
       raise ValueError(
           f'System prompt {_PH} not found in chat template: {sys_prompt}'
@@ -46,7 +50,10 @@ def parse_chat_template(tokenizer):
     if sys_prompt_parts[0].startswith(str(tokenizer.bos_token)):
       sys_prompt_parts[0] = sys_prompt_parts[0][len(tokenizer.bos_token) :]
 
-    messages.append({'role': 'user', 'content': _PH})
+    if no_sys_prompt:
+      messages = [{'role': 'user', 'content': _PH}]
+    else:
+      messages.append({'role': 'user', 'content': _PH})
     user_prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -133,20 +140,21 @@ def build_llm_metadata(
       if gen_config.temperature:
         sampler_params.temperature = gen_config.temperature
 
-  if isinstance(chat_templates, str):
-    llm_metadata.jinja_prompt_template = chat_templates
-  else:
-    sys_prompt_parts, user_prompt_parts, model_prompt_parts = chat_templates
-    pairs = []
-    if sys_prompt_parts[0] is not None:
-      pairs.append((sys_prompt_parts, llm_metadata.prompt_templates.system))
-    if user_prompt_parts[0] is not None:
-      pairs.append((user_prompt_parts, llm_metadata.prompt_templates.user))
-    if model_prompt_parts[0] is not None:
-      pairs.append((model_prompt_parts, llm_metadata.prompt_templates.model))
-    for pts, fld in pairs:
-      fld.prefix = pts[0]
-      fld.suffix = pts[1]
+  if chat_templates is not None:
+    if isinstance(chat_templates, str):
+      llm_metadata.jinja_prompt_template = chat_templates
+    else:
+      sys_prompt_parts, user_prompt_parts, model_prompt_parts = chat_templates
+      pairs = []
+      if sys_prompt_parts[0] is not None:
+        pairs.append((sys_prompt_parts, llm_metadata.prompt_templates.system))
+      if user_prompt_parts[0] is not None:
+        pairs.append((user_prompt_parts, llm_metadata.prompt_templates.user))
+      if model_prompt_parts[0] is not None:
+        pairs.append((model_prompt_parts, llm_metadata.prompt_templates.model))
+      for pts, fld in pairs:
+        fld.prefix = pts[0]
+        fld.suffix = pts[1]
 
   llm_metadata.max_num_tokens = context_length
 

litert_torch/generative/export_hf/core/utils.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Utility functions."""
 
+import os
+import re
 import torch
 
 
@@ -89,3 +91,35 @@ def has_sliding_attention(model):
     return False
   layer_types = getattr(model.config, 'layer_types', None)
   return layer_types is not None and 'sliding_attention' in layer_types
+
+
+def get_model_path_type(path_str: str) -> str:
+  """Determines if a string is a local path or a Hugging Face Repo ID.
+
+  Args:
+      path_str: The string to check.
+
+  Returns:
+      "local": If the path exists on disk.
+      "repo_id": If it looks like a Hub ID (e.g., 'meta-llama/Llama-2-7b').
+      "local_not_found": If it looks like a file path but doesn't exist.
+      "unknown": If it matches neither pattern clearly.
+  """
+  # 1. Absolute Truth: Does it exist on the disk?
+  if os.path.exists(path_str):
+    return 'local'
+
+  # 2. Heuristic: Does it have explicit path markers?
+  # Starts with "./", "/", "~", or contains Windows backslashes
+  if path_str.startswith(('.', '/', '~')) or '\\' in path_str:
+    return 'local_not_found'
+
+  # 3. Heuristic: Does it look like a Repo ID?
+  # Pattern: username/repo_name (e.g. "mistralai/Mistral-7B")
+  # or just repo_name for official models (e.g. "gpt2", "bert-base-uncased")
+  # Allowed chars: Alphanumeric, underscores, hyphens, periods.
+  repo_id_pattern = r'^(?:[\w\-\.]+\/)?[\w\-\.]+$'
+  if re.match(repo_id_pattern, path_str):
+    return 'repo_id'
+
+  return 'unknown'

litert_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of litert-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.8.0.dev20260126"
+__version__ = "0.9.0.dev20260127"

{litert_torch_nightly-0.8.0.dev20260126.dist-info → litert_torch_nightly-0.9.0.dev20260127.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litert-torch-nightly
-Version: 0.8.0.dev20260126
+Version: 0.9.0.dev20260127
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/litert-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI
@@ -41,6 +41,7 @@ Requires-Dist: jax
 Requires-Dist: torch-xla2[odml]>=0.0.1.dev20241201
 Requires-Dist: jaxtyping
 Requires-Dist: fire
+Requires-Dist: sentencepiece
 Provides-Extra: torch-xla
 Requires-Dist: torch_xla>=2.4.0; extra == "torch-xla"
 Dynamic: classifier

{litert_torch_nightly-0.8.0.dev20260126.dist-info → litert_torch_nightly-0.9.0.dev20260127.dist-info}/RECORD

@@ -1,9 +1,9 @@
 litert_torch/__init__.py,sha256=jgEpTtwnhxMhPGPTRvHJR7pzx6WU_mLbA-G1LjO7fnE,1279
 litert_torch/_config.py,sha256=zDnki83sBsQzDAea6bvzwccylWHnPUzbEyGGRh6B14w,2526
-litert_torch/cli.py,sha256=Svcs5U_HJIZYQUz5sf2Uu5__JTgxnU4-f4SBsZefgv0,985
+litert_torch/cli.py,sha256=TiguLo2O3_wY8cCnKnbXtUvGH4lzyjeSgsOnHsKd9Gg,984
 litert_torch/conftest.py,sha256=gYmFrsR4c_fjIidbyrDnek26yS0crDP6-UoyMvy-WFg,757
 litert_torch/model.py,sha256=KXFTyyfPM6AnP0JoSwsTqQR3lUQbMkTGSr3dUsfQ5Jk,5635
-litert_torch/version.py,sha256=B2Yv1xWPIduypuwtL5Z-8PlcGXPeaOafIegxJq8Dagw,804
+litert_torch/version.py,sha256=vbucmZdeLtxng93Sar6Ki8BLkqBQuisk59IB4CS7klU,804
 litert_torch/_convert/__init__.py,sha256=qdLdbj5NjhNG-QgY5O_8TzOr2XaDoWvmdY9JNPStQmw,670
 litert_torch/_convert/conversion.py,sha256=NuQEphyYp3W19IKvyTWo9pe7zt1-XmWM4zU9PDkUm54,6108
 litert_torch/_convert/conversion_utils.py,sha256=MWpB-3eN-rvQzTtXsPL30cDIK431SQuwvw3ia2K2ONM,2158
@@ -187,11 +187,11 @@ litert_torch/generative/export_hf/core/attention_test.py,sha256=RevOczfPncmbIBth
 litert_torch/generative/export_hf/core/cache.py,sha256=pfWh2SACdhNY2of2Z8KJC0wrSQ2jrkXgPHWe7PSEiuU,10263
 litert_torch/generative/export_hf/core/cache_base.py,sha256=FXMm9B8nDwC8uTyLmuBnYKLTnNtoeGN8gUnWwDCcH08,1714
 litert_torch/generative/export_hf/core/cache_test.py,sha256=y-v-oOGtRNPGWRfIfW3FcpDxvJbzrBU6Pb2o66FkUzU,6203
-litert_torch/generative/export_hf/core/export_lib.py,sha256=Gr3MqU1gBs4aSwpt3ag7dqwfCdlUXTXRmmeo3StF9mo,11622
+litert_torch/generative/export_hf/core/export_lib.py,sha256=Nvg3QiYRZcMiQd7du7w5vohazjLlZJW6YFk_WAbIpAs,12249
 litert_torch/generative/export_hf/core/exportable_module.py,sha256=XEqsV9M34OP-_vsxH7bnmxSCD6erAPl0a9I9JQM7v6k,8305
-litert_torch/generative/export_hf/core/litert_lm_builder.py,sha256=pnm5GbHmwr5KLTfONUrsMdvhwRFC2q5OkRPiRsWW_Ls,7715
+litert_torch/generative/export_hf/core/litert_lm_builder.py,sha256=f8Q2ifVyt65V-kRL0X9FRpQNKIer0R_Yx2lECZTMGPU,7965
 litert_torch/generative/export_hf/core/patches.py,sha256=i1fzs0anIFbBH-Q_PwCtp9VKXy64olJKwnGpnJUjkEo,1815
-litert_torch/generative/export_hf/core/utils.py,sha256=NxLMo4vgqG-8Hhr4ZsqDALtVV3n8rYbI1jiRaQfn-ho,2870
+litert_torch/generative/export_hf/core/utils.py,sha256=5Wgs9aAOKd2i8wmQF_IierLUuFG23v1T6zZPr-azQ7A,4018
 litert_torch/generative/export_hf/core/external_emb/__init__.py,sha256=5xWIp2ziIwapcZcjSKfeaFgBnIooa8ckhTQ7mazZC3c,670
 litert_torch/generative/export_hf/core/external_emb/exportable_module.py,sha256=mWn75lLms3BAeCTEvbkGZ2n4fxtwsqGA8PP4S8-JBdY,3058
 litert_torch/generative/export_hf/core/external_rope/__init__.py,sha256=5xWIp2ziIwapcZcjSKfeaFgBnIooa8ckhTQ7mazZC3c,670
@@ -318,9 +318,9 @@ litert_torch/testing/__init__.py,sha256=AfYP1HwTYSQmupveonEHCDV5dEyshzUgbwUrCUhb
 litert_torch/testing/export.py,sha256=3dR6oxnrdtX0MfqAfMv233cf3sHA4e0F2TBQotoo8xc,3292
 litert_torch/testing/model_coverage/__init__.py,sha256=uPXeAhWiD1O0aMDLCX7FTOSNQiea8yOtoIYPCuHEAG4,763
 litert_torch/testing/model_coverage/model_coverage.py,sha256=EPCI7PbNPb7GV28lo3qQvFdzJwJ_ZDrbCGdpeiBZhVo,4715
-litert_torch_nightly-0.8.0.dev20260126.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-litert_torch_nightly-0.8.0.dev20260126.dist-info/METADATA,sha256=IC63S5u9vXex_QMJhOklzWXMcPE-nr4JtC9S1V57NBU,2470
-litert_torch_nightly-0.8.0.dev20260126.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-litert_torch_nightly-0.8.0.dev20260126.dist-info/entry_points.txt,sha256=roYAi9hp0uYrMudMR59hGNF2pz0TSAtqNl4vQLJzxnE,55
-litert_torch_nightly-0.8.0.dev20260126.dist-info/top_level.txt,sha256=mGrsl2SYcjQSLBJX4ZXrHnFqHZe6QLRR7uk0tLfzwfM,13
-litert_torch_nightly-0.8.0.dev20260126.dist-info/RECORD,,
+litert_torch_nightly-0.9.0.dev20260127.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+litert_torch_nightly-0.9.0.dev20260127.dist-info/METADATA,sha256=bXq8cxXF9vyEVA9Kr6Q8ZKlz5x4v2L-aa1CfUBdsuAY,2499
+litert_torch_nightly-0.9.0.dev20260127.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+litert_torch_nightly-0.9.0.dev20260127.dist-info/entry_points.txt,sha256=roYAi9hp0uYrMudMR59hGNF2pz0TSAtqNl4vQLJzxnE,55
+litert_torch_nightly-0.9.0.dev20260127.dist-info/top_level.txt,sha256=mGrsl2SYcjQSLBJX4ZXrHnFqHZe6QLRR7uk0tLfzwfM,13
+litert_torch_nightly-0.9.0.dev20260127.dist-info/RECORD,,