通过

开始在 Databricks 上查询 LLM

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

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

还可以借助 AI 操场使用按令牌付费模型进行测试和聊天。 请参阅 使用 AI Playground 与 LLMs 聊天并原型设计生成式 AI 应用。

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

要求

重要

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

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

开始使用基础模型 API

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

在此示例中，通过使用 OpenAI 客户端查询模型，并在 model 字段中填入托管要查询模型的服务终结点的名称。 使用个人访问令牌填充 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' 消息，请使用 openai 升级 !pip install -U 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
    }
}