Upload scripts/bundle_litertlm.py with huggingface_hub

scripts/bundle_litertlm.py

@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+"""
+Bundle a Strategy-1 KV-cache TFLite + SentencePiece tokenizer into a
+.litertlm file compatible with Google AI Edge / LiteRT-LM runtime.
+
+Embeds:
+  - LlmMetadata proto: Gemma3 model type, 2K max tokens, TranslateGemma
+    Jinja chat template, BOS/EOS/end_of_turn stop tokens
+  - TFLite model  (model_type=prefill_decode)
+  - SentencePiece tokenizer
+
+Usage:
+  python bundle_litertlm.py \
+    --tflite /path/to/model.tflite \
+    --tokenizer /path/to/tokenizer.model \
+    --output /path/to/output.litertlm \
+    [--max-tokens 2048]
+"""
+
+import argparse
+import sys
+import tempfile
+from pathlib import Path
+
+# Make litert_lm package importable from /tmp/litert-lm-pkg
+sys.path.insert(0, "/tmp/litert-lm-pkg")
+
+from litert_lm_builder import litertlm_builder
+from litert_lm.runtime.proto import (
+    llm_metadata_pb2,
+    llm_model_type_pb2,
+    token_pb2,
+)
+
+
+# TranslateGemma 4B IT Jinja chat template (from tokenizer_config.json)
+# Requires structured content: {type, source_lang_code, target_lang_code, text}
+TRANSLATE_GEMMA_JINJA_TEMPLATE = """\
+{%- set languages = {
+    "aa": "Afar", "ab": "Abkhazian", "af": "Afrikaans", "ak": "Akan",
+    "am": "Amharic", "an": "Aragonese", "ar": "Arabic", "as": "Assamese",
+    "az": "Azerbaijani", "ba": "Bashkir", "be": "Belarusian", "bg": "Bulgarian",
+    "bm": "Bambara", "bn": "Bengali", "bo": "Tibetan", "br": "Breton",
+    "bs": "Bosnian", "ca": "Catalan", "ce": "Chechen", "co": "Corsican",
+    "cs": "Czech", "cv": "Chuvash", "cy": "Welsh", "da": "Danish",
+    "de": "German", "dv": "Divehi", "dz": "Dzongkha", "ee": "Ewe",
+    "el": "Greek", "en": "English", "eo": "Esperanto", "es": "Spanish",
+    "et": "Estonian", "eu": "Basque", "fa": "Persian", "ff": "Fulah",
+    "fi": "Finnish", "fo": "Faroese", "fr": "French", "fy": "Western Frisian",
+    "ga": "Irish", "gd": "Scottish Gaelic", "gl": "Galician", "gn": "Guarani",
+    "gu": "Gujarati", "gv": "Manx", "ha": "Hausa", "he": "Hebrew",
+    "hi": "Hindi", "hr": "Croatian", "ht": "Haitian", "hu": "Hungarian",
+    "hy": "Armenian", "ia": "Interlingua", "id": "Indonesian", "ig": "Igbo",
+    "ii": "Sichuan Yi", "ik": "Inupiaq", "io": "Ido", "is": "Icelandic",
+    "it": "Italian", "iu": "Inuktitut", "ja": "Japanese", "jv": "Javanese",
+    "ka": "Georgian", "ki": "Kikuyu", "kk": "Kazakh", "kl": "Kalaallisut",
+    "km": "Central Khmer", "kn": "Kannada", "ko": "Korean", "ks": "Kashmiri",
+    "ku": "Kurdish", "kw": "Cornish", "ky": "Kyrgyz", "la": "Latin",
+    "lb": "Luxembourgish", "lg": "Ganda", "ln": "Lingala", "lo": "Lao",
+    "lt": "Lithuanian", "lu": "Luba-Katanga", "lv": "Latvian", "mg": "Malagasy",
+    "mi": "Maori", "mk": "Macedonian", "ml": "Malayalam", "mn": "Mongolian",
+    "mr": "Marathi", "ms": "Malay", "mt": "Maltese", "my": "Burmese",
+    "nb": "Norwegian Bokmål", "nd": "North Ndebele", "ne": "Nepali",
+    "nl": "Dutch", "nn": "Norwegian Nynorsk", "no": "Norwegian",
+    "nr": "South Ndebele", "nv": "Navajo", "ny": "Chichewa", "oc": "Occitan",
+    "om": "Oromo", "or": "Oriya", "os": "Ossetian", "pa": "Punjabi",
+    "pl": "Polish", "ps": "Pashto", "pt": "Portuguese", "qu": "Quechua",
+    "rm": "Romansh", "rn": "Rundi", "ro": "Romanian", "ru": "Russian",
+    "rw": "Kinyarwanda", "sa": "Sanskrit", "sc": "Sardinian", "sd": "Sindhi",
+    "se": "Northern Sami", "sg": "Sango", "si": "Sinhala", "sk": "Slovak",
+    "sl": "Slovenian", "sn": "Shona", "so": "Somali", "sq": "Albanian",
+    "sr": "Serbian", "ss": "Swati", "st": "Southern Sotho", "su": "Sundanese",
+    "sv": "Swedish", "sw": "Swahili", "ta": "Tamil", "te": "Telugu",
+    "tg": "Tajik", "th": "Thai", "ti": "Tigrinya", "tk": "Turkmen",
+    "tl": "Tagalog", "tn": "Tswana", "to": "Tonga", "tr": "Turkish",
+    "ts": "Tsonga", "tt": "Tatar", "ug": "Uyghur", "uk": "Ukrainian",
+    "ur": "Urdu", "uz": "Uzbek", "ve": "Venda", "vi": "Vietnamese",
+    "vo": "Volapük", "wa": "Walloon", "wo": "Wolof", "xh": "Xhosa",
+    "yi": "Yiddish", "yo": "Yoruba", "za": "Zhuang", "zh": "Chinese",
+    "zu": "Zulu"
+} -%}
+{{ bos_token }}
+{%- if (messages[0]['role'] != 'user') -%}
+    {{ raise_exception("Conversations must start with a user prompt.") }}
+{%- endif -%}
+{%- for message in messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- if message['content'] is none or message['content'] is not string -%}
+            {{ raise_exception("Assistant role must provide content as a string") }}
+        {%- endif -%}
+        {{ '<start_of_turn>model\\n'}}
+        {{ message["content"] | trim }}
+    {%- elif (message['role'] == 'user') -%}
+        {%- set content = message["content"] if message["content"] is string else message["content"][0] -%}
+        {%- if content is string -%}
+            {%- set source_lang = "English" -%}
+            {%- set source_lang_code = "en" -%}
+            {%- set target_lang = "Spanish" -%}
+            {%- set target_lang_code = "es" -%}
+            {{ '<start_of_turn>user\\nYou are a professional ' + source_lang + ' (' + source_lang_code + ') to ' +
+               target_lang + ' (' + target_lang_code + ') translator. Your goal is to accurately convey the meaning ' +
+               'and nuances of the original ' + source_lang + ' text while adhering to ' + target_lang + ' grammar, ' +
+               'vocabulary, and cultural sensitivities.\\n' +
+               'Produce only the ' + target_lang + ' translation, without any additional explanations or ' +
+               'commentary. Please translate the following ' + source_lang + ' text into ' + target_lang + ':\\n\\n\\n' +
+               content | trim
+            }}
+        {%- else -%}
+            {%- set source_lang_code = content["source_lang_code"] | replace("_", "-") -%}
+            {%- set source_lang = languages.get(source_lang_code, source_lang_code) -%}
+            {%- set target_lang_code = content["target_lang_code"] | replace("_", "-") -%}
+            {%- set target_lang = languages.get(target_lang_code, target_lang_code) -%}
+            {{ '<start_of_turn>user\\nYou are a professional ' + source_lang + ' (' + source_lang_code + ') to ' +
+               target_lang + ' (' + target_lang_code + ') translator. Your goal is to accurately convey the meaning ' +
+               'and nuances of the original ' + source_lang + ' text while adhering to ' + target_lang + ' grammar, ' +
+               'vocabulary, and cultural sensitivities.\\n'
+            }}
+            {%- if content["type"] == 'text' -%}
+                {{ 'Produce only the ' + target_lang + ' translation, without any additional explanations or ' +
+                   'commentary. Please translate the following ' + source_lang + ' text into ' + target_lang + ':\\n\\n\\n' +
+                   content["text"] | trim
+                }}
+            {%- endif -%}
+        {%- endif -%}
+    {%- else -%}
+        {{ raise_exception("Conversations must only contain user or assistant roles.") }}
+    {%- endif -%}
+    {{ '<end_of_turn>\\n' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'<start_of_turn>model\\n'}}
+{%- endif -%}"""
+
+
+def build_llm_metadata_proto(max_tokens: int) -> bytes:
+    meta = llm_metadata_pb2.LlmMetadata()
+    meta.max_num_tokens = max_tokens
+
+    # Model type: Gemma3 (text-only variant — no vision config needed for TranslateGemma text mode)
+    meta.llm_model_type.gemma3.CopyFrom(llm_model_type_pb2.Gemma3())
+
+    # Start token: BOS = token id 2
+    meta.start_token.token_ids.ids.append(2)
+
+    # Stop tokens: EOS (id=1) and end_of_turn (id=106)
+    eos = meta.stop_tokens.add()
+    eos.token_ids.ids.append(1)
+
+    eot = meta.stop_tokens.add()
+    eot.token_ids.ids.append(106)
+
+    # Embed the Jinja template
+    meta.jinja_prompt_template = TRANSLATE_GEMMA_JINJA_TEMPLATE
+
+    return meta.SerializeToString()
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Bundle TFLite + tokenizer into .litertlm")
+    ap.add_argument("--tflite", required=True)
+    ap.add_argument("--tokenizer", required=True, help="SentencePiece .model file")
+    ap.add_argument("--output", required=True)
+    ap.add_argument("--max-tokens", type=int, default=2048)
+    ap.add_argument("--quant", default="int8", help="Quantization label for metadata")
+    args = ap.parse_args()
+
+    tflite_path = Path(args.tflite)
+    tokenizer_path = Path(args.tokenizer)
+    output_path = Path(args.output)
+
+    if not tflite_path.exists():
+        print(f"[x] TFLite not found: {tflite_path}", file=sys.stderr)
+        sys.exit(1)
+    if not tokenizer_path.exists():
+        print(f"[x] Tokenizer not found: {tokenizer_path}", file=sys.stderr)
+        sys.exit(1)
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Write LlmMetadata to temp file
+    meta_bytes = build_llm_metadata_proto(args.max_tokens)
+    with tempfile.NamedTemporaryFile(suffix=".pb", delete=False) as f:
+        meta_file = Path(f.name)
+        f.write(meta_bytes)
+
+    print(f"[+] Building .litertlm: {output_path.name}")
+    print(f"    TFLite: {tflite_path} ({tflite_path.stat().st_size / 1e9:.2f} GB)")
+    print(f"    Tokenizer: {tokenizer_path}")
+    print(f"    Max tokens: {args.max_tokens}")
+
+    Metadata = litertlm_builder.Metadata
+    DType = litertlm_builder.DType
+
+    builder = litertlm_builder.LitertLmFileBuilder()
+    builder.add_system_metadata(Metadata(key="model_name", value=f"TranslateGemma-4B-IT-{args.quant}", dtype=DType.STRING))
+    builder.add_system_metadata(Metadata(key="authors", value="google", dtype=DType.STRING))
+    builder.add_system_metadata(Metadata(key="quantization", value=args.quant, dtype=DType.STRING))
+
+    builder.add_tflite_model(
+        str(tflite_path),
+        model_type=litertlm_builder.TfLiteModelType.PREFILL_DECODE,
+    )
+    builder.add_sentencepiece_tokenizer(str(tokenizer_path))
+    builder.add_llm_metadata(str(meta_file))
+
+    with open(output_path, "wb") as f:
+        builder.build(f)
+    meta_file.unlink(missing_ok=True)
+
+    size = output_path.stat().st_size
+    print(f"[+] Written: {output_path} ({size / 1e9:.2f} GB)")
+
+
+if __name__ == "__main__":
+    main()