通过

使用 ai_query 批量执行 LLM 推理

  • 项目

重要

此功能目前以公共预览版提供。

本文介绍如何通过结合使用内置 Databricks SQL 函数 ai_query 与使用基础模型 API 预配吞吐量的终结点来执行批处理推理。 本文中的示例和指南适用于使用大型语言模型（LLM）处理多个输入的批处理推理工作负荷。

可以将 ai_query 与 SQL 或 PySpark 配合使用来运行批处理推理工作负荷。 若要对数据运行批处理推理，请在 ai_query中指定以下内容：

  • Unity Catalog 的输入表和输出表
  • 预配的吞吐量终结点名称
  • 模型提示和任何模型参数

有关此 AI 函数的更多详细信息，请参阅ai_query函数。

要求

  • 支持基础模型 API 的区域中的工作区。
  • 以下项之一：
    • 通用计算，计算大小为 i3.2xlarge 或更大，运行 Databricks Runtime 15.4 ML LTS 或更高版本，且至少有两个工作节点。
    • 中型及更大型的 SQL 仓库。
  • 现有的模型服务终结点。 请参阅 预配的吞吐量基础模型 API 来创建预配的吞吐量终结点。
  • Unity 目录中 Delta 表的查询权限，其中包含要使用的数据。
  • 将表属性中的 pipelines.channel 设置为“预览”以使用 ai_query()。 有关示例查询，请参阅示例。

使用 ai_query 和 SQL

下面是使用 ai_query 和 SQL 的批处理推理示例。 此示例包括 modelParameters 和 max_tokens、temperature，并展示如何使用 concat()将模型提示与输入列连接起来。 有多种方法可以执行串联，例如使用 ||、 concat()或 format_string()。


CREATE OR REPLACE TABLE ${output_table_name} AS (
  SELECT
      ${input_column_name},
      AI_QUERY(
        "${endpoint}",
        CONCAT("${prompt}", ${input_column_name}),
        modelParameters => named_struct('max_tokens', ${num_output_tokens},'temperature', ${temperature})
      ) as response
    FROM ${input_table_name}
    LIMIT ${input_num_rows}
)

使用 ai_query 和 PySpark

如果喜欢使用 Python，还可以使用 ai_query 和 PySpark 运行批处理推理，如下所示：

df_out = df.selectExpr("ai_query('{endpoint_name}', CONCAT('{prompt}', {input_column_name}), modelParameters => named_struct('max_tokens', ${num_output_tokens},'temperature', ${temperature})) as {output_column_name}")

df_out.write.mode("overwrite").saveAsTable(output_table_name)

使用 Python 的批量推断示例笔记本

以下示例笔记本创建预配的吞吐量终结点，并使用 Python 和 Meta Llama 3.1 70B 模型运行批量 LLM 推理。 它还提供有关对批处理推理工作负荷进行基准测试以及创建提供终结点的预配吞吐量模型的指导。

使用预配吞吐量终结点笔记本的 LLM 批量推理

获取笔记本

以下示例笔记本创建了一个预置吞吐量端点，并使用 Python 运行批量 LLM 推理。您可以从 GTE 大型（英语）或 BGE 大型（英语）嵌入模型中进行选择。

使用预配吞吐量终结点笔记本的 LLM 批量推理嵌入

获取笔记本