开始在 Databricks 上查询 LLM

本文介绍如何开始使用基础模型 API 在 Databricks 上处理和查询 LLM。

开始在 Databricks 上处理并查询 LLM 模型的最简单方法是以按令牌付费方式使用基础模型 API。 API 从 Databricks 工作区的“处理 UI”中自动提供的按令牌付费终结点提供对常用基础模型的访问。 请参阅按令牌付费支持的模型。

还可以借助 AI 操场使用按令牌付费模型进行测试和聊天。 请参阅与 LLM 聊天并使用 AI 操场制作 GenAI 应用原型。

对于生产工作负载，尤其是那些具有微调模型或需要性能保证的工作负载，Databricks 建议在预配的吞吐量终结点上使用基础模型 API。

要求

• 支持基础模型 API 按令牌付费的区域中的 Databricks 工作区。
• Databricks 个人访问令牌，用于使用 OpenAI 客户端查询和访问 Mosaic AI 模型服务终结点。

重要

作为适用于生产场景的安全最佳做法，Databricks 建议在生产期间使用计算机到计算机 OAuth 令牌来进行身份验证。

对于测试和开发，Databricks 建议使用属于服务主体（而不是工作区用户）的个人访问令牌。 若要为服务主体创建令牌，请参阅管理服务主体的令牌。

开始使用基础模型 API

以下示例应在 Databricks 笔记本中运行。 该代码示例查询了 Meta Llama 3.1 405B Instruct 模型，该模型在按标记付费终结点 databricks-meta-llama-3-1-405b-instruct 上提供。

在此示例中，通过使用托管要查询的模型的模型处理终结点的名称填充 model 字段，使用 OpenAI 客户端查询模型。 使用个人访问令牌填充 DATABRICKS_TOKEN 和 Databricks 工作区实例，从而将 OpenAI 客户端连接到 Databricks。

from openai import OpenAI
import os

DATABRICKS_TOKEN = os.environ.get("DATABRICKS_TOKEN")

client = OpenAI(
  api_key=DATABRICKS_TOKEN, # your personal access token
  base_url='https://<workspace_id>.databricks.com/serving-endpoints', # your Databricks workspace instance
)

chat_completion = client.chat.completions.create(
  messages=[
    {
      "role": "system",
      "content": "You are an AI assistant",
    },
    {
      "role": "user",
      "content": "What is a mixture of experts model?",
    }
  ],
  model="databricks-meta-llama-3-1-405b-instruct",
  max_tokens=256
)

print(chat_completion.choices[0].message.content)

注意

如果收到 ImportError: cannot import name 'OpenAI' from 'openai' 消息，请使用 !pip install -U openai 升级 openai 版本。 安装包后，运行 dbutils.library.restartPython()。

预期输出：

{
  "id": "xxxxxxxxxxxxx",
  "object": "chat.completion",
  "created": "xxxxxxxxx",
  "model": "databricks-meta-llama-3-1-405b-instruct",
  "choices": [
    {
      "index": 0,
      "message":
        {
          "role": "assistant",
          "content": "A Mixture of Experts (MoE) model is a machine learning technique that combines the predictions of multiple expert models to improve overall performance. Each expert model specializes in a specific subset of the data, and the MoE model uses a gating network to determine which expert to use for a given input."
        },
      "finish_reason": "stop"
    }
  ],
  "usage":
    {
      "prompt_tokens": 123,
      "completion_tokens": 23,
      "total_tokens": 146
    }
}

后续步骤

• 使用 AI 操场在熟悉的聊天界面中试用不同的模型。
• 查询生成 AI 模型。
• 使用外部模型访问 Databricks 外部托管的模型。
• 了解如何使用预配吞吐量终结点部署微调的模型。
• 探索监视模型质量和终结点运行状况的方法。