python/colabs/Getting Started with LiteRT-LM Python API.ipynb

{"cells":[{"cell_type":"code","source":["# Copyright 2026 The AI Edge Quantizer Authors.\n","#\n","# Licensed under the Apache License, Version 2.0 (the \"License\");\n","# you may not use this file except in compliance with the License.\n","# You may obtain a copy of the License at\n","#\n","#     http://www.apache.org/licenses/LICENSE-2.0\n","#\n","# Unless required by applicable law or agreed to in writing, software\n","# distributed under the License is distributed on an \"AS IS\" BASIS,\n","# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n","# See the License for the specific language governing permissions and\n","# limitations under the License.\n","# =============================================================================="],"metadata":{"id":"9GAYqvNyuEVC"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["<table class=\"tfo-notebook-buttons\" align=\"left\">\n","  <td>\n","    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-ai-edge/LiteRT-LM/blob/main/python/colabs/Getting Started with LiteRT-LM Python API.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n","  </td>\n","  <td>\n","    <a target=\"_blank\" href=\"https://github.com/google-ai-edge/ai-edge-quantizer/blob/main/python/colabs/Getting Started with LiteRT-LM Python API.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n","  </td>\n","</table>"],"metadata":{"id":"G9NEOdJOuHDE"}},{"cell_type":"markdown","metadata":{"id":"title"},"source":["# Getting Started with LiteRT-LM Python API\n","\n","This notebook walks through the [LiteRT-LM Python API](https://ai.google.dev/edge/litert-lm/python) with [Gemma 4 E2B](https://huggingface.co/litert-community/gemma-4-E2B-it-litert-lm).\n","\n","## 1. Installation\n","\n","First, let's install the LiteRT-LM API from PyPI. We will need the Hugging Face package to download the model."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"install"},"outputs":[],"source":["!pip install litert-lm huggingface_hub"]},{"cell_type":"markdown","metadata":{"id":"download_intro"},"source":["## 2. Download the [Gemma 4 E2B model](https://huggingface.co/litert-community/gemma-4-E2B-it-litert-lm)"]},{"cell_type":"code","metadata":{"id":"70db5841"},"source":["from huggingface_hub import hf_hub_download\n","\n","# Download the model file using your authenticated session\n","model_path = hf_hub_download(\n","    repo_id=\"litert-community/gemma-4-E2B-it-litert-lm\",\n","    filename=\"gemma-4-E2B-it.litertlm\",\n","    local_dir=\".\",\n","    local_dir_use_symlinks=False\n",")\n","\n","# Rename to the expected filename if necessary\n","import os\n","if os.path.exists(\"gemma-4-E2B-it.litertlm\"):\n","    os.rename(\"gemma-4-E2B-it.litertlm\", \"model.litertlm\")\n","print(\"Model downloaded successfully as model.litertlm\")"],"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"sync_intro"},"source":["## 3. Sending Messages (Synchronous)\n","\n","You can send a simple string or a full message structure synchronously.\n","\n","Note that the first initialization of `model.litertlm` will take longer. Future initialization will be faster with the cache."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"send_sync","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1775154422354,"user_tz":420,"elapsed":23630,"user":{"displayName":"","userId":""}},"outputId":"b1f6b5f1-5bec-4669-ece4-5524c6c69838"},"outputs":[{"output_type":"stream","name":"stdout","text":["Response: The capital of France is **Paris**.\n"]}],"source":["import litert_lm\n","\n","with litert_lm.Engine(\"model.litertlm\") as engine:\n","  with engine.create_conversation() as conversation:\n","    response = conversation.send_message(\"What is the capital of France?\")\n","    print(f\"Response: {response['content'][0]['text']}\")"]},{"cell_type":"markdown","metadata":{"id":"async_intro"},"source":["## 4. Sending Messages (Asynchronous/Streaming)\n","\n","Streaming the response as it's generated."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"send_async","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1775154447870,"user_tz":420,"elapsed":11203,"user":{"displayName":"","userId":""}},"outputId":"6add942d-75f1-46ba-e1b4-ae1521f7d709"},"outputs":[{"output_type":"stream","name":"stdout","text":["Here's one for you:\n","\n","Why don't scientists trust atoms?\n","\n","**Because they make up everything!**"]}],"source":["import litert_lm\n","\n","with litert_lm.Engine(\"model.litertlm\") as engine:\n","  with engine.create_conversation() as conversation:\n","    for chunk in conversation.send_message_async(\"Tell me a joke.\"):\n","      print(chunk[\"content\"][0][\"text\"], end=\"\", flush=True)"]},{"cell_type":"markdown","metadata":{"id":"multimodality_intro"},"source":["## 5. Multi-Modality\n","\n","Gemma 3n supports multi-modality. Here's how to interact with it using audio or images (if supported)."]},{"cell_type":"code","source":["# Download the audio file for this example\n","!wget https://github.com/google-ai-edge/LiteRT-LM/raw/refs/heads/main/runtime/testdata/have_a_wonderful_day.wav -O have_a_wonderful_day.wav"],"metadata":{"id":"URKsr2FYYD33"},"execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"multimodality_code","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1775154498847,"user_tz":420,"elapsed":10081,"user":{"displayName":"","userId":""}},"outputId":"0b97d63c-2805-4580-ee10-2b564e120211"},"outputs":[{"output_type":"stream","name":"stdout","text":["The audio says, \"Have a wonderful day.\"\n"]}],"source":["import litert_lm\n","\n","user_message = {\n","    \"role\": \"user\",\n","    \"content\": [\n","        {\"type\": \"audio\", \"path\": \"/content/have_a_wonderful_day.wav\"},\n","        {\"type\": \"text\", \"text\": \"Describe this audio. What does it say?\"},\n","    ],\n","}\n","\n","# Initialize the engine with Audio backend.\n","with litert_lm.Engine(\"model.litertlm\", audio_backend=litert_lm.Backend.CPU) as engine:\n","  with engine.create_conversation() as conversation:\n","    for chunk in conversation.send_message_async(user_message):\n","      print(chunk[\"content\"][0][\"text\"], end=\"\", flush=True)"]},{"cell_type":"markdown","metadata":{"id":"tools_intro"},"source":["## 6. Defining and Using Tools\n","\n","You can define Python functions as tools that the model can call automatically. Note that this requires models with tool support."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"tools_code","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1775154486043,"user_tz":420,"elapsed":22061,"user":{"displayName":"","userId":""}},"outputId":"35548f46-a730-4fd8-aab7-b78b10031ece"},"outputs":[{"output_type":"stream","name":"stdout","text":["[tool call] 'product' with numbers [12.34, 98.76]\n","The product of 12.34 and 98.76 is 1218.6984."]}],"source":["import litert_lm\n","\n","def product(numbers: float) -> float:\n","    \"\"\"Compute the product of the given numbers.\n","\n","    Args:\n","      numbers: The list of numbers.\n","    \"\"\"\n","    print(\"[tool call] 'product' with numbers\", numbers)\n","\n","    result = 1\n","    for x in numbers:\n","        result *= x\n","    return result\n","\n","with litert_lm.Engine(\"model.litertlm\") as engine:\n","  with engine.create_conversation(tools=[product]) as conversation:\n","    # The model will call add_numbers automatically if it needs to sum values\n","    for chunk in conversation.send_message_async(\"What is the product of 12.34 and 98.76?\"):\n","      print(chunk[\"content\"][0][\"text\"], end=\"\", flush=True)"]}],"metadata":{"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.0"},"colab":{"provenance":[{"file_id":"https://github.com/google-ai-edge/LiteRT-LM/blob/main/python/colabs/Getting%20Started%20with%20LiteRT-LM%20Python%20API.ipynb","timestamp":1775154556374}],"last_runtime":{"build_target":"","kind":"local"}}},"nbformat":4,"nbformat_minor":0}