| import os |
| import logging |
| from litert_torch import converter |
| from mediapipe.tasks.python.genai import bundler |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") |
| log = logging.getLogger(__name__) |
|
|
| MODEL_PATH = "artifacts/deploy/ankahi-gemma4-e4b-int8" |
| TFLITE_PATH = "artifacts/litertlm/ankahi_q8.tflite" |
| TASK_PATH = "artifacts/litertlm/ankahi_q8.task" |
|
|
| def main(): |
| os.makedirs("artifacts/litertlm", exist_ok=True) |
| |
| log.info(f"Converting {MODEL_PATH} to LiteRT (.tflite)...") |
| try: |
| converter.convert( |
| model_path=MODEL_PATH, |
| output_path=TFLITE_PATH, |
| quantize="INT8", |
| prefill_seq_len=512, |
| kv_cache_max_len=1024, |
| backend="cpu" |
| ) |
| log.info(f"Successfully created TFLite model at {TFLITE_PATH}") |
| except Exception as e: |
| log.error(f"LiteRT conversion failed: {e}") |
| return |
|
|
| log.info(f"Bundling into MediaPipe Task (.task)...") |
| try: |
| |
| tokenizer_model = os.path.join(MODEL_PATH, "tokenizer.model") |
| if not os.path.exists(tokenizer_model): |
| |
| log.warning("tokenizer.model not found, searching for alternatives...") |
| tokenizer_model = os.path.join(MODEL_PATH, "tokenizer.json") |
| |
| config = bundler.BundleConfig( |
| tflite_model=TFLITE_PATH, |
| tokenizer_model=tokenizer_model, |
| start_token="<bos>", |
| stop_tokens=["<eos>", "<end_of_turn>"], |
| output_filename=TASK_PATH, |
| enable_streaming=True |
| ) |
| bundler.create_bundle(config) |
| log.info(f"Successfully created MediaPipe task at {TASK_PATH}") |
| except Exception as e: |
| log.error(f"Task bundling failed: {e}") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
