LLM 推断指南

借助 LLM Inference API，您可以在设备上完全运行大语言模型 (LLM)，从而执行各种任务，例如生成文本、以自然语言形式检索信息以及总结文档。该任务内置了对多种文本到文本大语言模型的支持，因此您可以将最新的设备端生成式 AI 模型应用于您的应用和产品。

试试看！arrow_forward

该任务内置了对各种 LLM 的支持。LiteRT 社区页面上托管的模型采用 MediaPipe 友好格式，无需任何额外的转换或编译步骤。

您可以使用 AI Edge Torch 将 PyTorch 模型导出为多签名 LiteRT (tflite) 模型，这些模型与分词器参数捆绑在一起，以创建任务包。使用 AI Edge Torch 转换的模型与 LLM 推理 API 兼容，并且可以在 CPU 后端上运行，因此适合 Android 和 iOS 应用。

开始使用

请按照以下适用于目标平台的实现指南之一，开始使用此任务。以下特定于平台的指南将引导您完成此任务的基本实现，并提供使用可用模型和建议的配置选项的代码示例：

• Web：

  • 指南
  • 代码示例

• Android：

  • 指南
  • 代码示例

• iOS

  • 指南
  • 代码示例

任务详情

本部分介绍此任务的功能、输入、输出和配置选项。

功能

LLM Inference API 包含以下主要功能：

1. 文本到文本生成 - 根据输入的文本提示生成文本。
2. LLM 选择 - 应用多个模型，根据您的特定使用场景量身定制应用。您还可以重新训练模型并应用自定义权重。
3. LoRA 支持 - 通过 LoRA 模型扩展和自定义 LLM 功能，方法是训练所有数据集，或从开源社区获取预建的 LoRA 模型（与使用 AI Edge Torch Generative API 转换的模型不兼容）。

配置选项

此任务具有以下配置选项：

模型

LLM Inference API 支持许多文本到文本的大语言模型，包括对多种模型的内置支持，这些模型经过优化，可在浏览器和移动设备上运行。这些轻量级模型可用于完全在设备上运行推断。

在初始化 LLM 推理 API 之前，请下载模型并将文件存储在项目目录中。您可以使用 LiteRT 社区 HuggingFace 代码库中预先转换的模型，也可以使用 AI Edge Torch Generative Converter 将模型转换为 MediaPipe 兼容格式。

如果您还没有可与 LLM 推理 API 搭配使用的 LLM，请先使用以下模型之一。

Gemma-3n

Gemma-3n E2B 和 E4B 是 Gemma 系列中的最新模型，该系列是一系列先进的轻量级开放模型，其开发采用了与 Gemini 模型相同的研究成果和技术。Gemma 3n 模型旨在在低资源设备上高效执行。它们能够处理文本、图片和音频输入，并生成文本输出。

Gemma 3n 模型使用选择性参数激活技术来降低资源要求。这项技术可让模型以 20 亿和 40 亿个参数的有效规模运行，低于它们包含的参数总数

下载 Gemma-3n E2B

下载 Gemma-3n E4B

HuggingFace 的 Gemma-3n E2B 和 E4B 模型采用 .litertlm 格式，可随时与适用于 Android 和 Web 的 LLM Inference API 搭配使用。

Gemma-3 1B

Gemma-3 1B 是 Gemma 系列轻量级、先进的开放模型中最轻量级的模型，采用与 Gemini 模型相同的研究成果和技术构建而成。该模型包含 10 亿个参数，并且权重开放。

下载 Gemma-3 1B

HuggingFace 的 Gemma-3 1B 模型以 .task/.litertlm 格式提供，可随时用于 Android 和 Web 应用的 LLM Inference API。

使用 LLM 推理 API 运行 Gemma-3 1B 时，请相应地配置以下选项：

• preferredBackend：使用此选项可在 CPU 或 GPU 后端之间进行选择。此选项仅适用于 Android 设备。
• supportedLoraRanks：LLM 推理 API 无法配置为支持使用 Gemma-3 1B 模型进行低秩自适应 (LoRA)。请勿使用 supportedLoraRanks 或 loraRanks 选项。
• maxTokens：maxTokens 的值必须与模型中内置的上下文大小相匹配。这也可以称为键值 (KV) 缓存或上下文长度。
• numResponses：必须始终为 1。此选项仅适用于网页版。

在 Web 应用上运行 Gemma-3 1B 时，初始化可能会导致当前线程长时间阻塞。如果可能，请始终从工作线程运行模型。

Gemma-2 2B

Gemma-2 2B 是 Gemma-2 的 2B 变体，可在所有平台上运行。

下载 Gemma-2 2B

该模型包含 20 亿个参数，并且权重开放。Gemma-2 2B 以其同类模型中出色的推理能力而闻名。

PyTorch 模型转换

借助 AI Edge Torch Generative API，可以将 PyTorch 生成式模型转换为 MediaPipe 兼容格式。您可以使用该 API 将 PyTorch 模型转换为多签名 LiteRT (TensorFlow Lite) 模型。如需详细了解如何映射和导出模型，请访问 AI Edge Torch GitHub 页面。

使用 AI Edge Torch Generative API 转换 PyTorch 模型涉及以下步骤：

1. 下载 PyTorch 模型检查点。
2. 使用 AI Edge Torch Generative API 将模型创作、转换和量化为 MediaPipe 兼容的文件格式 (.tflite)。
3. 根据 tflite 文件和模型分词器创建任务软件包 (.task/.litertlm)。

Torch Generative 转换器仅适用于 CPU，并且需要至少具有 64 GB RAM 的 Linux 机器。

如需创建任务包，请使用打包脚本创建任务包。打包过程会将映射的模型与额外的元数据（例如，运行端到端推理所需的令牌化器参数。

模型捆绑流程需要 MediaPipe PyPI 软件包。转换脚本在 0.10.14 之后的所有 MediaPipe 软件包中均可用。

使用以下命令安装并导入依赖项：

$ python3 -m pip install mediapipe

使用 genai.bundler 库来捆绑模型：

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 自定义

LLM 推理 API 的 LoRA 支持适用于 GPU 后端的所有 Gemma 变体和 Phi-2 模型，LoRA 权重仅适用于注意力层。此初始实现用作实验性 API，以便在未来的开发中支持更多模型和各种类型的层。

准备 LoRA 模型

按照 HuggingFace 上的说明，使用支持的模型类型（Gemma 或 Phi-2）在您自己的数据集上训练微调的 LoRA 模型。Gemma-2 2B、Gemma 2B 和 Phi-2 模型均以 safetensors 格式在 HuggingFace 上提供。由于 LLM Inference API 仅支持注意力层上的 LoRA，因此在创建 LoraConfig 时，请仅指定注意力层，如下所示：

# 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"],
)

对于测试，HuggingFace 上提供了适合 LLM 推理 API 的可公开访问的微调 LoRA 模型。例如，monsterapi/gemma-2b-lora-maths-orca-200k 适用于 Gemma-2B，lole25/phi-2-sft-ultrachat-lora 适用于 Phi-2。

在准备好的数据集上训练并保存模型后，您会获得一个 adapter_model.safetensors 文件，其中包含经过微调的 LoRA 模型权重。safetensors 文件是模型转换中使用的 LoRA 检查点。

下一步，您需要使用 MediaPipe Python 软件包将模型权重转换为 TensorFlow Lite Flatbuffer。ConversionConfig 应指定基础模型选项以及其他 LoRA 选项。请注意，由于该 API 仅支持使用 GPU 进行 LoRA 推理，因此后端必须设置为 '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)

转换器将输出两个 TFLite flatbuffer 文件，一个用于基础模型，另一个用于 LoRA 模型。

LoRA 模型推理

Web、Android 和 iOS LLM 推理 API 已更新，可支持 LoRA 模型推理。

// 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)

如需使用 LoRA 运行 LLM 推理，请使用与基础模型相同的 generateResponse() 或 generateResponseAsync() 方法。