メインコンテンツまでスキップ

基盤モデルのクエリー

この記事では、基盤モデルのクエリ要求を書式設定し、モデルサービングエンドポイントに送信する方法について説明します。 Databricks によってホストされている基盤モデルと、Databricks の外部でホストされている基盤モデルに対してクエリを実行できます。

従来の ML モデルまたは Python モデルのクエリ要求については、 カスタム モデルのクエリ サービング エンドポイントを参照してください。

Mosaic AI Model Serving は、基盤モデルにアクセスするための基盤モデル API と外部モデルをサポートしています。モデルサービングは、統一されたOpenAI互換の API とクエリ SDK を使用します。 これにより、エクスペリメントを行い、サポートされているクラウドやプロバイダー間での本番運用のための基盤モデルをカスタマイズすることが可能になります。

Mosaic AI Model Serving は、基盤モデルまたは外部モデルを提供するエンドポイントにスコアリング要求を送信するための次のオプションを提供します。

メソッド

詳細

オープンAIクライアント

OpenAI クライアントを使用して、 Mosaic AI Model Servingエンドポイントによってホストされているモデルをクエリします。 model入力としてモデルサーバーのエンドポイント名を指定します。 基盤モデルAPIまたは外部モデルによって提供されるチャット、埋め込み、補完モデルがサポートされます。

SQL 関数

ai_query SQL 関数を使用して、SQL から直接モデル推論を呼び出します。例: 基盤モデルのクエリを参照してください。

サービングUI

サービングエンドポイント ページから エンドポイントのクエリー を選択します。JSON形式のモデル入力データを挿入し、 リクエストを送信 をクリックします。 モデルに入力例がログに記録されている場合は、 例を表示 を使用して読み込みます。

REST API

REST API を使用してモデルを呼び出し、クエリを実行します。 詳細は POST /serving-endpoints/{name}/invocations を参照してください。 複数のモデルを提供するエンドポイントへの要求のスコアリングについては、エンドポイントの背後にある個々のモデルのクエリを参照してください。

MLflow Deployments SDK

MLflow Deployments SDK の predict() 関数を使用して、モデルに対してクエリを実行します。

Databricks Python SDK

Databricks Python SDK は、REST API の上位にあるレイヤーです。 認証などの低レベルの詳細を処理するため、モデルとの対話が容易になります。

必要条件​

important

本番運用シナリオのセキュリティのベスト プラクティスとして、 Databricks では、本番運用中の認証に マシン間 OAuth トークン を使用することをお勧めします。

テストと開発のために、 Databricks ワークスペース ユーザーではなく 、サービスプリンシパル に属する個人用アクセス トークンを使用することをお勧めします。 サービスプリンシパルのトークンを作成するには、「 サービスプリンシパルのトークンの管理」を参照してください。

パッケージのインストール​

クエリ方法を選択したら、まず適切なパッケージをクラスターにインストールする必要があります。

To use the OpenAI client, the databricks-sdk[openai] package needs to be installed on your cluster. Databricks SDK provides a wrapper for constructing the OpenAI client with authorization automatically configured to query generative AI models. Run the following in your notebook or your local terminal:

!pip install databricks-sdk[openai]>=0.35.0

The following is only required when installing the package on a Databricks Notebook

Python
dbutils.library.restartPython()

チャットコンプリーションモデルのクエリ​

チャット モデルに対してクエリを実行する例を次に示します。 この例は、モデル サービス機能 (基盤モデルAPIまたは外部モデル) のいずれかを使用して使用可能になったチャット モデルのクエリに適用されます。

バッチ推論の例については、「AI関数を使用したバッチLLM 推論の実行 」を参照してください。

The following is a chat request for the DBRX Instruct model made available by the Foundation Model APIs pay-per-token endpoint, databricks-dbrx-instruct in your workspace.

To use the OpenAI client, specify the model serving endpoint name as the model input.

Python

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
openai_client = w.serving_endpoints.get_open_ai_client()

response = openai_client.chat.completions.create(
    model="databricks-dbrx-instruct",
    messages=[
      {
        "role": "system",
        "content": "You are a helpful assistant."
      },
      {
        "role": "user",
        "content": "What is a mixture of experts model?",
      }
    ],
    max_tokens=256
)

To query foundation models outside of your workspace, you must use the OpenAI client directly. You also need your Databricks workspace instance to connect the OpenAI client to Databricks. The following example assumes you have a Databricks API token and openai installed on your compute.

Python

import os
import openai
from openai import OpenAI

client = OpenAI(
    api_key="dapi-your-databricks-token",
    base_url="https://example.staging.cloud.databricks.com/serving-endpoints"
)

response = client.chat.completions.create(
    model="databricks-dbrx-instruct",
    messages=[
      {
        "role": "system",
        "content": "You are a helpful assistant."
      },
      {
        "role": "user",
        "content": "What is a mixture of experts model?",
      }
    ],
    max_tokens=256
)

例として、REST API を使用する場合のチャットモデルの想定されるリクエスト形式を次に示します。 外部モデルの場合、特定のプロバイダーとエンドポイント構成に有効な追加のパラメーターを含めることができます。 追加のクエリ・パラメーターを参照してください。

Bash
{
  "messages": [
    {
      "role": "user",
      "content": "What is a mixture of experts model?"
    }
  ],
  "max_tokens": 100,
  "temperature": 0.1
}

以下は、REST API を使用して行われたリクエストに対して想定されるレスポンス形式です。

JSON
{
  "model": "databricks-dbrx-instruct",
  "choices": [
    {
      "message": {},
      "index": 0,
      "finish_reason": null
    }
  ],
  "usage": {
    "prompt_tokens": 7,
    "completion_tokens": 74,
    "total_tokens": 81
  },
  "object": "chat.completion",
  "id": null,
  "created": 1698824353
}

エンベディングモデルのクエリ​

以下は、基盤モデルAPIによって利用可能になった gte-large-en モデルのエンベディングのリクエストです。 この例は、モデルサービング機能 (基盤モデルAPIまたは外部モデル) のいずれかを使用して使用可能になったエンベディングモデルのクエリに適用されます。

To use the OpenAI client, specify the model serving endpoint name as the model input.

Python

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
openai_client = w.serving_endpoints.get_open_ai_client()

response = openai_client.embeddings.create(
  model="databricks-gte-large-en",
  input="what is databricks"
)

To query foundation models outside your workspace, you must use the OpenAI client directly, as demonstrated below. The following example assumes you have a Databricks API token and openai installed on your compute. You also need your Databricks workspace instance to connect the OpenAI client to Databricks.

Python

import os
import openai
from openai import OpenAI

client = OpenAI(
    api_key="dapi-your-databricks-token",
    base_url="https://example.staging.cloud.databricks.com/serving-endpoints"
)

response = client.embeddings.create(
  model="databricks-gte-large-en",
  input="what is databricks"
)

以下は、エンベディングモデルで想定される要求形式です。 外部モデルの場合、特定のプロバイダーとエンドポイント構成に有効な追加のパラメーターを含めることができます。 追加のクエリ・パラメーターを参照してください。

Bash

{
  "input": [
    "embedding text"
  ]
}

想定される応答形式は次のとおりです。

JSON
{
  "object": "list",
  "data": [
    {
      "object": "embedding",
      "index": 0,
      "embedding": []
    }
  ],
  "model": "text-embedding-ada-002-v2",
  "usage": {
    "prompt_tokens": 2,
    "total_tokens": 2
  }
}

エンベディングが正規化されているかどうかを確認する​

モデルによって生成されたエンべディングが正規化されているかどうかを確認するには、次を使用します。

Python

  import numpy as np

  def is_normalized(vector: list[float], tol=1e-3) -> bool:
      magnitude = np.linalg.norm(vector)
      return abs(magnitude - 1) < tol

テキストコンプリーションモデルのクエリ​

important

Querying text completion models made available using Foundation Model APIs pay-per-token using the OpenAI client is not supported. Only querying external models using the OpenAI client is supported as demonstrated in this section.

To use the OpenAI client, specify the model serving endpoint name as the model input. The following example queries the claude-2 completions model hosted by Anthropic using the OpenAI client. To use the OpenAI client, populate the model field with the name of the model serving endpoint that hosts the model you want to query.

This example uses a previously created endpoint, anthropic-completions-endpoint, configured for accessing external models from the Anthropic model provider. See how to create external model endpoints.

See Supported models for additional models you can query and their providers.

Python

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
openai_client = w.serving_endpoints.get_open_ai_client()

completion = openai_client.completions.create(
model="anthropic-completions-endpoint",
prompt="what is databricks",
temperature=1.0
)
print(completion)

以下は、完了モデルで想定される要求形式です。 外部モデルの場合、特定のプロバイダーとエンドポイント構成に有効な追加のパラメーターを含めることができます。 追加のクエリ・パラメーターを参照してください。

Bash
{
  "prompt": "What is mlflow?",
  "max_tokens": 100,
  "temperature": 0.1,
  "stop": [
    "Human:"
  ],
  "n": 1,
  "stream": false,
  "extra_params":
  {
    "top_p": 0.9
  }
}

想定される応答形式は次のとおりです。

JSON
{
  "id": "cmpl-8FwDGc22M13XMnRuessZ15dG622BH",
  "object": "text_completion",
  "created": 1698809382,
  "model": "gpt-3.5-turbo-instruct",
  "choices": [
    {
      "text": "MLflow is an open-source platform for managing the end-to-end machine learning lifecycle. It provides tools for tracking experiments, managing and deploying models, and collaborating on projects. MLflow also supports various machine learning frameworks and languages, making it easier to work with different tools and environments. It is designed to help data scientists and machine learning engineers streamline their workflows and improve the reproducibility and scalability of their models.",
      "index": 0,
      "logprobs": null,
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 5,
    "completion_tokens": 83,
    "total_tokens": 88
  }
}

AI Playground を使用してサポートされている LLM とチャットする​

サポートされている大規模言語モデルは、 AI Playgroundを使用して操作できます。 AI Playground は、Databricks ワークスペースから LLM をテスト、プロンプト、比較できるチャットのような環境です。

AI playground

追加のリソース​

に最終更新