LLM Inference guide

The LLM Inference API lets you run large language models (LLMs) completely on-device, which you can use to perform a wide range of tasks, such as generating text, retrieving information in natural language form, and summarizing documents. The task provides built-in support for multiple text-to-text large language models, so you can apply the latest on-device generative AI models to your apps and products.

Try it!arrow_forward

The task supports the following variants of Gemma: Gemma-2 2B, Gemma 2B, and Gemma 7B. Gemma is a family of lightweight, state-of-the-art open models built from the same research and technology used to create the Gemini models. It also supports the following external models: Phi-2, Falcon-RW-1B and StableLM-3B.

In addition to the supported models, you can use Google's AI Edge Torch to export PyTorch models into multi-signature LiteRT (tflite) models, which are bundled with tokenizer parameters to create Task Bundles that are compatible with the LLM Inference API. Models converted with AI Edge Torch can only run on the CPU backend and are therefore limited to Android and iOS.

Get Started

Start using this task by following one of these implementation guides for your target platform. These platform-specific guides walk you through a basic implementation of this task, with code examples that use an available model and the recommended configuration options:

• Web:

  • Guide
  • Code example

• Android:

  • Guide
  • Code example

• iOS

  • Guide
  • Code example

Task details

This section describes the capabilities, inputs, outputs, and configuration options of this task.

Features

The LLM Inference API contains the following key features:

1. Text-to-text generation - Generate text based on an input text prompt.
2. LLM selection - Apply multiple models to tailor the app for your specific use cases. You can also retrain and apply customized weights to the model.
3. LoRA support - Extend and customize the LLM capability with LoRA model either by training on your all dataset, or taking prepared prebuilt LoRA models from the open-source community (not compatible with models converted with the AI Edge Torch Generative API).

Configurations options

This task has the following configuration options:

Models

The LLM Inference API supports many text-to-text large language models, including built-in support for several models that are optimized to run on browsers and mobile devices. These lightweight models can be used to run inferences completely on-device.

Before initializing the LLM Inference API, download a model and store the file within your project directory. You can use a pre-converted model or convert a model to a MediaPipe-compatible format.

The LLM Inference API is compatible with two categories types of models, some of which require model conversion. Use the table to identify the required steps method for your model.

We are hosting the converted .bin files for Gemma 2B, Gemma 7B, and Gemma-2 2B on Kaggle. These models can be directly deployed using our LLM Inference API. To learn how you can convert other models, see the Model Conversion section.

Gemma-2 2B

Gemma-2 2B is the latest model in the Gemma family of lightweight, state-of-the-art open models built from the same research and technology used to create the Gemini models. The model contains 2B parameters and open weights. Gemma-2 2B is known for state-of-the art reasoning skills for models in its class.

Download Gemma-2 2B

The Gemma-2 2B models are available in the following variants:

• gemma2-2b-it-cpu-int8: Gemma-2 2B 8-bit model with CPU compatibility.
• gemma2-2b-it-gpu-int8: Gemma-2 2B 8-bit model with GPU compatibility.

You can also tune the model and add new weights before adding it to the app. For more information on tuning and customizing Gemma, see Tuning Gemma. After downloading Gemma-2 2B from Kaggle Models, the model is already in the appropriate format to use with MediaPipe Tasks.

Gemma 2B

Gemma 2B is a part of a family of lightweight, state-of-the-art open models built from the same research and technology used to create the Gemini models. The model contains 2B parameters and open weights. This model is well-suited for a variety of text generation tasks, including question answering, summarization, and reasoning.

Download Gemma 2B

The Gemma 2B models are available in the following variants:

• gemma-2b-it-cpu-int4: Gemma 2B 4-bit model with CPU compatibility.
• gemma-2b-it-cpu-int8: Gemma 2B 8-bit model with CPU compatibility.
• gemma-2b-it-gpu-int4: Gemma 2B 4-bit model with GPU compatibility.
• gemma-2b-it-gpu-int8: Gemma 2B 8-bit model with GPU compatibility.

You can also tune the model and add new weights before adding it to the app. For more information on tuning and customizing Gemma, see Tuning Gemma. After downloading Gemma 2B from Kaggle Models, the model is already in the appropriate format to use with MediaPipe Tasks.

Gemma 7B

Gemma 7B is a larger Gemma model with 7B parameters and open weights. The model is more powerful for a variety of text generation tasks, including question answering, summarization, and reasoning. Gemma 7B is only supported on Web.

Download Gemma 7B

The Gemma 7B model comes in one variant:

• gemma-1.1-7b-it-gpu-int8: Gemma 7B 8-bit model with GPU compatibility.

After downloading Gemma 7B from Kaggle Models, the model is already in the appropriate format to use with MediaPipe.

Falcon 1B

Falcon-1B is a 1 billion parameter causal decoder-only model trained on 350B tokens of RefinedWeb.

Download Falcon 1B

The LLM Inference API requires the following files to be downloaded and stored locally:

• tokenizer.json
• tokenizer_config.json
• pytorch_model.bin

After downloading the Falcon model files, the model is ready to be converted to the MediaPipe format with a conversion script. Follow the steps in the Conversion script for supported models section.

StableLM 3B

StableLM-3B is a 3 billion parameter decoder-only language model pre-trained on 1 trillion tokens of diverse English and code datasets for 4 epochs.

Download StableLM 3B

The LLM Inference API requires the following files to be downloaded and stored locally:

• tokenizer.json
• tokenizer_config.json
• model.safetensors

After downloading the StableLM model files, the model is ready to be converted to the MediaPipe format with a conversion script. Follow the steps in the Conversion script for supported models section.

Phi-2

Phi-2 is a 2.7 billion parameter Transformer model. It was trained using various NLP synthetic texts and filtered websites. The model is best suited for prompts using the Question-Answer, chat, and code format.

Download Phi-2

The LLM Inference API requires the following files to be downloaded and stored locally:

• tokenizer.json
• tokenizer_config.json
• model-00001-of-00002.safetensors
• model-00002-of-00002.safetensors

After downloading the Phi-2 model files, the model is ready to be converted to the MediaPipe format with a conversion script. Follow the steps in the Conversion script for supported models section.

Generative PyTorch Models

PyTorch generative models can be converted to a MediaPipe-compatible format with the AI Edge Torch Generative API. You can use the API to convert PyTorch models into multi-signature LiteRT (TensorFlow Lite) models. For more details on mapping and exporting models, visit the AI Edge Torch GitHub page.

If you intend to use the AI Edge Torch Generative API to convert a PyTorch model, follow the steps in the Torch Generative converter for PyTorch models section.

Model conversion

The MediaPipe LLM Inference API lets you run a wide variety of large language models on-device. This includes models that have been pre-converted to a MediaPipe-compatible format, as well as other models that can be converted with a conversion script or the AI Edge Torch library.

The LLM Inference API accepts models in the .bin and .task file formats. Pre-converted models and models converted with the conversion script will be .bin files, while models converted with the AI Edge Torch library will be .task files. Do not manually alter the file formats of your converted models.

The LLM Inference API contains three model conversion paths:

• Pre-converted models (Gemma 2B, Gemma 7B, Gemma-2 2B): No conversion required.
• Supported models (Phi-2, StableLM, Falcon): MediaPipe conversion script.
• Other PyTorch models (All PyTorch LLM models): AI Edge Torch Generative API.

Pre-converted models

The Gemma-2 2B, Gemma 2B and Gemma 7B models are available as pre-converted models in the MediaPipe format. These models do not require any additional conversion steps from the user and can be run as-is with the LLM Inference API.

You can download the Gemma-2 2B from Kaggle Models:

• gemma2-2b-it-cpu-int8: Gemma-2 2B 8-bit model with CPU compatibility.
• gemma2-2b-it-gpu-int8: Gemma-2 2B 8-bit model with GPU compatibility.

You can download variants of Gemma 2B from Kaggle Models:

• gemma-2b-it-cpu-int4: Gemma 2B 4-bit model with CPU compatibility.
• gemma-2b-it-cpu-int8: Gemma 2B 8-bit model with CPU compatibility.
• gemma-2b-it-gpu-int4: Gemma 2B 4-bit model with GPU compatibility.
• gemma-2b-it-gpu-int8: Gemma 2B 8-bit model with GPU compatibility.

You can download the Gemma 7B from Kaggle Models:

• gemma-1.1-7b-it-gpu-int8: Gemma 7B 8-bit model with GPU compatibility.

For more information on the Gemma models, see the documentation on Gemma-2 2B, Gemma 2B and Gemma 7B.

Conversion script for supported models

The MediaPipe package offers a conversion script to convert the following external models into a MediaPipe-compatible format:

• Falcon 1B
• StableLM 3B
• Phi-2

For more information on the supported external models, see the documentation on Falcon 1B, StableLM 3B, and Phi-2.

The model conversion process requires the MediaPipe PyPI package. The conversion script is available in all MediaPipe packages after 0.10.11.

Install and import the dependencies with the following:

$ python3 -m pip install mediapipe

Use the genai.converter library to convert the model:

import mediapipe as mp
from mediapipe.tasks.python.genai import converter

config = converter.ConversionConfig(
  input_ckpt=INPUT_CKPT,
  ckpt_format=CKPT_FORMAT,
  model_type=MODEL_TYPE,
  backend=BACKEND,
  output_dir=OUTPUT_DIR,
  combine_file_only=False,
  vocab_model_file=VOCAB_MODEL_FILE,
  output_tflite_file=OUTPUT_TFLITE_FILE,
)

converter.convert_checkpoint(config)

To convert the LoRA model, the ConversionConfig should specify the base model options as well as additional LoRA options. Notice that since the API only supports LoRA inference with GPU, the backend must be set to 'gpu'.

import mediapipe as mp
from mediapipe.tasks.python.genai import converter

config = converter.ConversionConfig(
  # Other params related to base model
  ...
  # Must use gpu backend for LoRA conversion
  backend='gpu',
  # LoRA related params
  lora_ckpt=LORA_CKPT,
  lora_rank=LORA_RANK,
  lora_output_tflite_file=LORA_OUTPUT_TFLITE_FILE,
)

converter.convert_checkpoint(config)

The converter will output two TFLite flatbuffer files, one for the base model and the other for the LoRA model.

Torch Generative converter for PyTorch models

PyTorch generative models can be converted to a MediaPipe-compatible format with the AI Edge Torch Generative API. You can use the API to author, convert, and quantize PyTorch LLMs to use with the LLM Inference API. The Torch Generative converter only converts for CPU and requires a Linux machine with at least 64 GBs of RAM.

Converting a PyTorch model with the AI Edge Torch Generative API involves the following:

1. Download the PyTorch model checkpoints
2. Use the AI Edge Torch Generative API to author, convert, and quantize the model to a MediaPipe-compatible file format (.tflite).
3. Create a Task Bundle (.task) from the tflite file and the model tokenizer.

To create a Task Bundle, use the bundling script to create a Task Bundle. The bundling process packs the mapped model with additional metadata (e.g., Tokenizer Parameters) needed to run end-to-end inference.

The model bundling process requires the MediaPipe PyPI package. The conversion script is available in all MediaPipe packages after 0.10.14.

Install and import the dependencies with the following:

$ python3 -m pip install mediapipe

Use the genai.bundler library to bundle the model:

import mediapipe as mp
from mediapipe.tasks.python.genai import bundler

config = bundler.BundleConfig(
    tflite_model=TFLITE_MODEL,
    tokenizer_model=TOKENIZER_MODEL,
    start_token=START_TOKEN,
    stop_tokens=STOP_TOKENS,
    output_filename=OUTPUT_FILENAME,
    enable_bytes_to_unicode_mapping=ENABLE_BYTES_TO_UNICODE_MAPPING,
)
bundler.create_bundle(config)

LoRA customization

LoRA support of the LLM Inference API works for all Gemma variants and Phi-2 models for the GPU backend, with LoRA weights applicable to attention layers only. This initial implementation serves as an experimental API for future developments with plans to support more models and various types of layers in the coming updates.

Prepare LoRA models

Follow the instructions on HuggingFace to train a fine tuned LoRA model on your own dataset with supported model types, Gemma or Phi-2. Gemma-2 2B, Gemma 2B and Phi-2 models are both available on HuggingFace in the safetensors format. Since LLM Inference API only supports LoRA on attention layers, only specify attention layers while creating the LoraConfig as following:

# For Gemma
from peft import LoraConfig
config = LoraConfig(
    r=LORA_RANK,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
)

# For Phi-2
config = LoraConfig(
    r=LORA_RANK,
    target_modules=["q_proj", "v_proj", "k_proj", "dense"],
)

For testing, there are publicly accessible fine-tuned LoRA models which fit LLM Inference API available on HuggingFace. For example, monsterapi/gemma-2b-lora-maths-orca-200k for Gemma-2B and lole25/phi-2-sft-ultrachat-lora for Phi-2.

After training on the prepared dataset and saving the model, you obtain an adapter_model.safetensors file containing the fine-tuned LoRA model weights. The safetensors file is the LoRA checkpoint used in the model conversion.

As the next step, you need convert the model weights into a TensorFlow Lite Flatbuffer using the MediaPipe Python Package. The ConversionConfig should specify the base model options as well as additional LoRA options. Notice that since the API only supports LoRA inference with GPU, the backend must be set to 'gpu'.

import mediapipe as mp
from mediapipe.tasks.python.genai import converter

config = converter.ConversionConfig(
  # Other params related to base model
  ...
  # Must use gpu backend for LoRA conversion
  backend='gpu',
  # LoRA related params
  lora_ckpt=LORA_CKPT,
  lora_rank=LORA_RANK,
  lora_output_tflite_file=LORA_OUTPUT_TFLITE_FILE,
)

converter.convert_checkpoint(config)

The converter will output two TFLite flatbuffer files, one for the base model and the other for the LoRA model.

LoRA model inference

The Web, Android and iOS LLM Inference API are updated to support LoRA model inference.

// Set the configuration options for the LLM Inference task
val options = LlmInferenceOptions.builder()
        .setModelPath('<path to base model>')
        .setMaxTokens(1000)
        .setTopK(40)
        .setTemperature(0.8)
        .setRandomSeed(101)
        .setLoraPath('<path to LoRA model>')
        .build()

// Create an instance of the LLM Inference task
llmInference = LlmInference.createFromOptions(context, options)

To run LLM inference with LoRA, use the same generateResponse() or generateResponseAsync() methods as the base model.