Upload scripts/bundle_litertlm.py with huggingface_hub

scripts/bundle_litertlm.py

@@ -33,106 +33,29 @@ from litert_lm.runtime.proto import (
 )
 
 
-# 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 -%}"""
+# Simple Jinja template compatible with LiteRT-LM runtime (no .get(), no complex tests).
+# Handles plain text input from Google AI Edge Gallery.
+# Uses the exact prompt format TranslateGemma was trained with (en→es default).
+# Users who need other language pairs should prefix their message with the pair,
+# e.g. "Translate English to French:\n\nHello"
+TRANSLATE_GEMMA_JINJA_TEMPLATE = \
+"{{ bos_token }}" \
+"{% for message in messages %}" \
+"{% if message['role'] == 'user' %}" \
+"<start_of_turn>user\n" \
+"You are a professional translator. " \
+"Produce only the translation of the following text, without any additional explanations or commentary:\n\n\n" \
+"{{ message['content'] | trim }}" \
+"<end_of_turn>\n" \
+"{% elif message['role'] == 'assistant' %}" \
+"<start_of_turn>model\n" \
+"{{ message['content'] | trim }}" \
+"<end_of_turn>\n" \
+"{% endif %}" \
+"{% endfor %}" \
+"{% if add_generation_prompt %}" \
+"<start_of_turn>model\n" \
+"{% endif %}"
 
 
 def build_llm_metadata_proto(max_tokens: int) -> bytes: