Created
December 8, 2025 14:35
-
-
Save monday8am/eba470685303095bb1e71bd8497ce668 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # STEP 6 — Build .litertlm using high-level API | |
| %cd /content/ai-edge-torch/ | |
| import os | |
| import json | |
| import tempfile | |
| from ai_edge_torch.generative.utilities import litertlm_builder | |
| print(f"🏗️ Building .litertlm bundle with {MODEL_NAME} metadata...") | |
| # Load Jinja chat template from checkpoint | |
| jinja_template = None | |
| tokenizer_config_path = os.path.join(checkpoint_path, "tokenizer_config.json") | |
| if os.path.exists(tokenizer_config_path): | |
| with open(tokenizer_config_path, 'r') as f: | |
| config = json.load(f) | |
| jinja_template = config.get('chat_template') | |
| if jinja_template: | |
| print(f"✅ Loaded Jinja template from checkpoint ({len(jinja_template)} chars)") | |
| else: | |
| print("⚠️ No chat_template found in tokenizer_config.json") | |
| else: | |
| print("⚠️ tokenizer_config.json not found") | |
| # Map the string MODEL_NAME to the correct LlmModelType | |
| MODEL_TO_LLM_TYPE = { | |
| "hammer2.1": litertlm_builder.litertlm_builder.LlmModelType.GENERIC, | |
| "qwen2.5": litertlm_builder.litertlm_builder.LlmModelType.QWEN2P5, | |
| "qwen3": litertlm_builder.litertlm_builder.LlmModelType.QWEN3, | |
| } | |
| # Fallback: if unknown model name, default to generic one | |
| ll_model_type = MODEL_TO_LLM_TYPE.get( | |
| MODEL_NAME.lower(), | |
| litertlm_builder.litertlm_builder.LlmModelType.GENERIC | |
| ) | |
| print(f"{ll_model_type} model type selected!") | |
| # Build bundle using high-level API (PR #858) | |
| with tempfile.TemporaryDirectory() as workdir: | |
| litertlm_builder.build_litertlm( | |
| tflite_model_path=OUTPUT_TFLITE, | |
| workdir=workdir, | |
| output_path="/content/", | |
| context_length=CONTEXT_LENGTH, | |
| model_prompt_prefix=None, | |
| model_prompt_suffix=None, | |
| user_prompt_prefix=None, | |
| user_prompt_suffix=None, | |
| tokenizer_model_path=OUTPUT_SPM, | |
| hf_tokenizer_model_path=None, | |
| start_token="<|im_start|>", | |
| stop_tokens=["<|im_end|>", "<|endoftext|>"], | |
| llm_model_type=ll_model_type, | |
| jinja_prompt_template=jinja_template, | |
| ) | |
| print(f"✅ Bundle built with {ll_model_type} metadata and Jinja template") | |
| if os.path.exists(OUTPUT_LITERT): | |
| size_mb = os.path.getsize(OUTPUT_LITERT) / (1024 ** 2) | |
| print(f"🎉 Success! Bundle created ({size_mb:.2f} MB)") | |
| else: | |
| print("❌ Conversion failed — check logs above") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment