LLM 推断指南

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

试试吧！

此任务支持 Gemma 的以下变体：Gemma-2 2B、Gemma 2B 和 Gemma 7B。Gemma 是一系列先进的轻量级开放式模型，其开发采用了与 Gemini 模型相同的研究成果和技术。它还支持以下外部模型：Phi-2、Falcon-RW-1B 和 StableLM-3B。

除了受支持的模型之外，您还可以使用 Google 的 AI Edge Torch 将 PyTorch 模型导出为多签名 LiteRT (tflite) 模型，这些模型会与分词器参数捆绑在一起，以创建与 LLM 推理 API 兼容的任务软件包。使用 AI Edge Torch 转换的模型只能在 CPU 后端上运行，因此仅限于 Android 和 iOS。

开始使用

如需开始使用此任务，请按照适用于目标平台的以下任一实现指南操作。以下平台专用指南将引导您完成此任务的基本实现，其中包含使用可用模型和建议的配置选项的代码示例：

任务详情

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

功能

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

  1. 文本到文本生成 - 根据输入的文本提示生成文本。
  2. LLM 选择 - 应用多个模型，根据您的特定用例量身定制应用。您还可以对模型进行重新训练，并应用自定义权重。
  3. LoRA 支持 - 使用 LoRA 模型扩展和自定义 LLM 功能，方法是基于您的所有数据集进行训练，或从开源社区获取准备好的预构建 LoRA 模型（与使用 AI Edge Torch Generative API 转换的模型不兼容）。
任务输入 任务输出
LLM Inference API 接受以下输入：
  • 文本提示（例如问题、电子邮件主题、要总结的文档）
 LLM 推理 API 会输出以下结果：
  • 根据输入提示生成的文本（例如，问题的答案、电子邮件草稿、文档摘要）

配置选项

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

选项名称 说明 值范围 默认值
modelPath 模型在项目目录中的存储路径。 路径 不适用
maxTokens 模型处理的词元（输入词元 + 输出词元）数量上限。 整数 512
topK 模型在生成的每个步骤中考虑的令牌数。 将预测限制为前 k 个概率最高的 token。 整数 40
temperature 生成期间引入的随机性程度。温度越高，生成的文本越富有创造力，温度越低，生成的文本越具可预测性。 浮点数 0.8
randomSeed 文本生成期间使用的随机种子。 整数 0
loraPath 设备本地 LoRA 模型的绝对路径。注意：此功能仅适用于 GPU 型号。 路径 不适用
resultListener 设置结果监听器以异步接收结果。 仅在使用异步生成方法时适用。 不适用 不适用
errorListener 设置可选的错误监听器。 不适用 不适用

模型

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

在初始化 LLM Inference API 之前，请下载模型并将该文件存储在项目目录中。您可以使用预转换的模型，也可以将模型转换为与 MediaPipe 兼容的格式。

LLM Inference API 与两类模型兼容，其中一些模型需要转换。使用下表确定您的模型所需的步骤方法。

模型 转化方法 兼容的平台 文件类型
支持的模型 Gemma 2B、Gemma 7B、Gemma-2 2B、Phi-2、StableLM、Falcon MediaPipe Android、iOS、网站 .bin
其他 PyTorch 模型 所有 PyTorch LLM 模型 AI Edge Torch Generative 库 Android、iOS .task

我们在 Kaggle 上托管了 Gemma 2B、Gemma 7B 和 Gemma-2 2B 的转换后的 .bin 文件。这些模型可以使用我们的 LLM 推理 API 直接部署。如需了解如何转换其他模型，请参阅模型转换部分。

Gemma-2 2B

Gemma-2 2B 是 Gemma 系列先进的轻量级开放模型中的最新模型，采用了与 Gemini 模型相同的研究成果和技术。该模型包含 20 亿个参数和开放权重。Gemma-2 2B 以同类模型中最先进的推理能力而闻名。

下载 Gemma-2 2B

Gemma-2 2B 型号有以下变体：

您还可以在将模型添加到应用之前对其进行调优并添加新的权重。如需详细了解如何调优和自定义 Gemma，请参阅调优 Gemma。从 Kaggle 模型下载 Gemma-2 2B 后，该模型已采用适合与 Mediapipe 任务搭配使用的格式。

Gemma 2B

Gemma 2B 是一系列先进的轻量级开放模型之一，其开发采用了与 Gemini 模型相同的研究成果和技术。该模型包含 20 亿个参数和开放权重。此模型非常适合处理各种文本生成任务，包括问答、摘要和推理。

下载 Gemma 2B

Gemma 2B 模型有以下变体：

您还可以在将模型添加到应用之前对其进行调优并添加新的权重。如需详细了解如何调优和自定义 Gemma，请参阅调优 Gemma。从 Kaggle 模型下载 Gemma 2B 后，该模型已采用适合与 Mediapipe 任务搭配使用的格式。

Gemma 7B

Gemma 7B 是一个较大的 Gemma 模型，具有 70 亿个参数和开放权重。该模型在各种文本生成任务（包括问答、摘要和推理）方面更强大。Gemma 7B 仅在 Web 上受支持。

下载 Gemma 7B

Gemma 7B 模型有一种变体：

从 Kaggle 模型下载 Gemma 7B 后，该模型已采用适合与 MediaPipe 搭配使用的格式。

Falcon 1B

Falcon-1B 是一个拥有 10 亿参数的因果解码器专用模型，基于 RefinedWeb 的 3,500 亿个词元进行了训练。

下载 Falcon 1B

LLM Inference API 要求下载并在本地存储以下文件：

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

下载 Falcon 模型文件后，您可以使用转换脚本将模型转换为 MediaPipe 格式。请按照适用于受支持型号的转化脚本部分中的步骤操作。

StableLM 3B

StableLM-3B 是一个拥有 30 亿参数的 decoder-only 语言模型，在 4 个 epoch 内基于多样化的英语和代码数据集内的 1 万亿个词元进行了预训练。

下载 StableLM 3B

LLM Inference API 要求下载并在本地存储以下文件：

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

下载 StableLM 模型文件后，就可以使用转换脚本将模型转换为 MediaPipe 格式了。请按照适用于受支持型号的转化脚本部分中的步骤操作。

Phi-2

Phi-2 是一个拥有 27 亿参数的 Transformer 模型。该模型是使用各种 NLP 合成文本和过滤后的网站训练而成的。该模型最适合用于采用问答、聊天和代码格式的提示。

下载 Phi-2

LLM Inference API 要求下载并在本地存储以下文件：

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

下载 Phi-2 模型文件后，您可以使用转换脚本将模型转换为 MediaPipe 格式。请按照适用于受支持型号的转化脚本部分中的步骤操作。

生成式 PyTorch 模型

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

如果您打算使用 AI Edge Torch Generative API 转换 PyTorch 模型，请按照适用于 PyTorch 模型的 Torch Generative 转换器部分中的步骤操作。

模型转换

借助 MediaPipe LLM Inference API，您可以在设备端运行各种各样的大型语言模型。这包括已预先转换为与 MediaPipe 兼容的格式的模型，以及可使用转换脚本或 AI Edge Torch 库转换的其他模型。

LLM Inference API 接受 .bin 和 .task 文件格式的模型。预转换的模型和使用转换脚本转换的模型将是 .bin 文件，而使用 AI Edge Torch 库转换的模型将是 .task 文件。请勿手动更改转换后的模型的文件格式。

LLM Inference API 包含三个模型转换路径：

预转换的模型

Gemma-2 2B、Gemma 2B 和 Gemma 7B 模型以 MediaPipe 格式提供为预转换的模型。这些模型无需用户执行任何额外的转换步骤，并且可以使用 LLM 推理 API 按原样运行。

您可以从 Kaggle 模型下载 Gemma-2 2B：

您可以从 Kaggle 模型下载 Gemma 2B 的变体：

您可以从 Kaggle 模型下载 Gemma 7B：

如需详细了解 Gemma 模型，请参阅 Gemma-2 2B、Gemma 2B 和 Gemma 7B 的文档。

适用于受支持模型的转化脚本

MediaPipe 软件包提供了一个转换脚本，用于将以下外部模型转换为与 MediaPipe 兼容的格式：

如需详细了解受支持的外部模型，请参阅 Falcon 1B、StableLM 3B 和 Phi-2 的文档。

模型转换过程需要 MediaPipe PyPI 软件包。0.10.11 之后，所有 MediaPipe 软件包中都提供了转换脚本。

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

$ python3 -m pip install mediapipe

使用 genai.converter 库转换模型：

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)

如需转换 LoRA 模型，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 模型。

参数 说明 可接受的值
input_ckpt model.safetensors 或 pytorch.bin 文件的路径。请注意，有时模型 safetensors 格式会分片到多个文件，例如 model-00001-of-00003.safetensors、model-00001-of-00003.safetensors。您可以指定文件模式，例如 model*.safetensors。 路径
ckpt_format 模型文件格式。 {"safetensors", "pytorch"}
model_type 要转换的 LLM。 {"PHI_2", "FALCON_RW_1B", "STABLELM_4E1T_3B", "GEMMA_2B"}
backend 用于运行模型的处理器（代理）。 {"cpu", "gpu"}
output_dir 用于存放各层权重文件的输出目录的路径。 路径
output_tflite_file 输出文件的路径。例如“model_cpu.bin”或“model_gpu.bin”。此文件仅与 LLM Inference API 兼容，不能用作通用的 `tflite` 文件。 路径
vocab_model_file 存储 tokenizer.json 和 tokenizer_config.json 文件的目录的路径。对于 Gemma，请指向单个 tokenizer.model 文件。 路径
lora_ckpt 存储 LoRA 适配器权重的 safetensors 文件的 LoRA ckpt 的路径。 路径
lora_rank 表示 LoRA ckpt 排名的整数。必须提供此值才能转换 Lora 权重。如果未提供，转换器会假定没有 LoRA 权重。注意：只有 GPU 后端支持 LoRA。 整数
lora_output_tflite_file LoRA 权重的输出 tflite 文件名。 路径

适用于 PyTorch 模型的 Torch 生成式转换器

您可以使用 AI Edge Torch Generative API 将 PyTorch 生成式模型转换为与 MediaPipe 兼容的格式。您可以使用该 API 编写、转换和量化 PyTorch LLM，以便与 LLM Inference API 搭配使用。Torch 生成式转换器仅适用于 CPU，并且需要至少 64 GB RAM 的 Linux 机器。

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

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

如需创建任务软件包，请使用捆绑脚本创建任务软件包。捆绑过程会将映射的模型与其他元数据（例如分词器参数）来运行端到端推理。

模型捆绑流程需要 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)
参数 说明 可接受的值
tflite_model AI Edge 导出的 TFLite 模型的路径。 路径
tokenizer_model SentencePiece 词解析器模型的路径。 路径
start_token 模型专用开始令牌。提供的词解析器模型中必须包含起始令牌。 STRING
stop_tokens 模型专用停止令牌。提供的词解析器模型中必须包含停止标记。 LIST[STRING]
output_filename 输出任务软件包文件的名称。 路径

LoRA 自定义

Mediapipe LLM Inference API 可配置为支持大语言模型的低秩自适应 (LoRA)。利用经过微调的 LoRA 模型，开发者可以通过经济高效的训练流程自定义 LLM 的行为。

LLM Inference 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 模型。例如，对于 Gemma-2B，使用 monsterapi/gemma-2b-lora-maths-orca-200k；对于 Phi-2，使用 lole25/phi-2-sft-ultrachat-lora。

使用准备好的训练数据集进行训练并保存模型后，您会获得一个包含经过微调的 LoRA 模型权重的 adapter_model.safetensors 文件。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 模型推理。

Android 在初始化期间支持静态 LoRA。如需加载 LoRA 模型，用户需要指定 LoRA 模型路径以及基础 LLM。

// 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() 方法。