生成 AI モデルのクエリ

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

従来の ML または Python モデルのクエリ要求については、「カスタム モデルのエンドポイントを提供するクエリ」をご覧ください。

Mosaic AI Model Serving では、生成 AI のモデルにアクセスするための Foundation Models API と外部モデルがサポートされています。 Model Serving では、OpenAI と互換性のある統合された API と SDK を使って、それらのクエリが実行されます。 これにより、サポートされているすべてのクラウドとプロバイダーで運用環境の生成 AI モデルの実験とカスタマイズを行うことができます。

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

メソッド	詳細
OpenAI クライアント	OpenAI クライアントを使用して、Mosaic AI Model Serving エンドポイントによってホストされているモデルに対してクエリを実行します。 モデル サービング エンドポイント名を model 入力として指定します。 Foundation Model API または外部モデルで使用できるチャット、埋め込み、入力候補モデルでサポートされます。
提供 UI	[提供エンドポイント] ページから [エンドポイントに対するクエリを実行] を選択します。 JSON 形式のモデル入力データを入力し、[要求の送信] をクリックします。 モデルに入力例がログされている場合は、[例の表示] を使用してそれを読み込みます。
REST API	REST API を使用してモデルを呼び出してクエリを実行します。 詳細については、「POST /serving-endpoints/{name}/invocations」を参照してください。 複数のモデルを提供するエンドポイントへのスコアリング要求については、「エンドポイントの背後にある個々のモデルに対するクエリの実行」を参照してください。
MLflow デプロイ SDK	モデルに対してクエリを実行するには、MLflow Deployments SDK の predict() 関数を使用します。
Databricks Python SDK	Databricks Python SDK は、REST API 上のレイヤーです。 認証などの低レベルの詳細を処理するため、モデルとの対話が容易になります。
SQL 関数	SQL 関数の ai_query を使用して、SQL から直接モデル推論を呼び出します。 「ai_query() を使用した提供モデルに対するクエリの実行」を参照してください。

要件

• モデルの提供のエンドポイント。
• サポートされているリージョン内の Databricks ワークスペース。
  • Foundation Model API のリージョン
  • 外部モデルのリージョン
• OpenAI クライアント、REST API、または MLflow デプロイ SDK を介してスコアリング要求を送信するには、Databricks API トークンが必要です。

重要

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

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

パッケージをインストールする

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

OpenAI クライアント

OpenAI クライアントを使用するには、クラスターに databricks-sdk[openai] パッケージをインストールする必要があります。 Databricks SDK には、生成 AI モデルに対してクエリを実行するように自動的に構成された承認を使用して OpenAI クライアントを構築するためのラッパーが用意されています。 ノートブックまたはローカル ターミナルで次を実行してください。

!pip install databricks-sdk[openai]

以下は、Databricks Notebook にパッケージをインストールする場合にのみ必要です

dbutils.library.restartPython()

REST API

Serving REST API へのアクセスは、Databricks Runtime for Machine Learning で利用できます。

MLflow デプロイ SDK

!pip install mlflow

以下は、Databricks Notebook にパッケージをインストールする場合にのみ必要です

dbutils.library.restartPython()

Databricks Python SDK

Databricks SDK for Python は、Databricks Runtime 13.3 LTS 以上を使用するすべての Azure Databricks クラスターに既にインストールされています。 Databricks Runtime 12.2 LTS 以下を使用する Azure Databricks クラスターの場合、まず Databricks SDK for Python をインストールする必要があります。 Databricks SDK for Python を参照してください。

チャット入力候補モデルに対してクエリを実行する

チャット モデルのクエリの例を次に示します。 この例は、Model Serving の機能 (Foundation Model API または外部モデル) のいずれかを使って利用できるチャット モデルのクエリに適用されます。

バッチ推論の例については、「 Foundation Model API によってプロビジョニングされたスループットを使用したバッチ推論を参照してください。

OpenAI クライアント

以下は、ワークスペースで提供される Foundation Model API 従量課金トークン エンドポイント databricks-dbrx-instruct によって利用可能な DBRX Instruct モデルのチャット要求です。

OpenAI クライアントを使用するには、モデル サービング エンドポイント名を model 入力として指定します。

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
)

ワークスペースの外部で基礎モデルのクエリを実行するには、OpenAI クライアントを直接使用する必要があります。 OpenAI クライアントを Databricks に接続するために、Databricks ワークスペース インスタンスも必要です。 次の例では、Databricks API トークンがあり、コンピューティングに openai がインストールされていることを前提としています。

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

重要

次の例では、基盤モデルを提供する提供エンドポイントに対するクエリを実行するために REST API パラメーターを使用します。 これらのパラメーターはパブリック プレビューであり、定義が変更される可能性があります。 POST /serving-endpoints/{name}/invocations を参照してください。

以下は、ワークスペースで提供される Foundation Model API 従量課金トークン エンドポイント databricks-dbrx-instruct によって利用可能な DBRX Instruct モデルのチャット要求です。

curl \
-u token:$DATABRICKS_TOKEN \
-X POST \
-H "Content-Type: application/json" \
-d '{
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful assistant."
    },
    {
      "role": "user",
      "content": " What is a mixture of experts model?"
    }
  ]
}' \
https://<workspace_host>.databricks.com/serving-endpoints/databricks-dbrx-instruct/invocations \

MLflow デプロイ SDK

重要

次の例では、MLflow デプロイ SDK の predict() API を使用します。

以下は、ワークスペースで提供される Foundation Model API 従量課金トークン エンドポイント databricks-dbrx-instruct によって利用可能な DBRX Instruct モデルのチャット要求です。

import mlflow.deployments

# Only required when running this example outside of a Databricks Notebook
export DATABRICKS_HOST="https://<workspace_host>.databricks.com"
export DATABRICKS_TOKEN="dapi-your-databricks-token"

client = mlflow.deployments.get_deploy_client("databricks")

chat_response = client.predict(
    endpoint="databricks-dbrx-instruct",
    inputs={
        "messages": [
            {
              "role": "user",
              "content": "Hello!"
            },
            {
              "role": "assistant",
              "content": "Hello! How can I assist you today?"
            },
            {
              "role": "user",
              "content": "What is a mixture of experts model??"
            }
        ],
        "temperature": 0.1,
        "max_tokens": 20
    }
)

Databricks Python SDK

以下は、ワークスペースで提供される Foundation Model API 従量課金トークン エンドポイント databricks-dbrx-instruct によって利用可能な DBRX Instruct モデルのチャット要求です。

このコードは、ワークスペース内のノートブックで実行する必要があります。 「Azure Databricks ノートブックから Databricks SDK for Python を使用する」を参照してください。

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()
response = w.serving_endpoints.query(
    name="databricks-dbrx-instruct",
    messages=[
        ChatMessage(
            role=ChatMessageRole.SYSTEM, content="You are a helpful assistant."
        ),
        ChatMessage(
            role=ChatMessageRole.USER, content="What is a mixture of experts model?"
        ),
    ],
    max_tokens=128,
)
print(f"RESPONSE:\n{response.choices[0].message.content}")

LangChain

LangChain を使用して基盤モデル エンドポイントにクエリを実行するには、ChatDatabricks ChatModel クラスを使用し、endpoint を指定します。

次の例では、LangChain の ChatDatabricks ChatModel クラスを使用して、Foundation Model API のトークン単位の支払いエンドポイントである databricks-dbrx-instruct にクエリを実行します。

%pip install langchain-databricks

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_databricks import ChatDatabricks

messages = [
    SystemMessage(content="You're a helpful assistant"),
    HumanMessage(content="What is a mixture of experts model?"),
]

llm = ChatDatabricks(endpoint_name="databricks-dbrx-instruct")
llm.invoke(messages)

SQL

重要

次の例では、組み込みの SQL 関数である ai_query を使用します。 この関数はパブリック プレビューであり、定義が変更される可能性があります。 「ai_query() を使用した提供モデルに対するクエリの実行」を参照してください。

以下は、ワークスペースで提供される Foundation Model API 従量課金トークン エンドポイント databricks-meta-llama-3-1-70b-instruct によって利用可能な meta-llama-3-1-70b-instruct のチャット要求です。

Note

ai_query() 関数は、DBRX または DBRX Instruct モデルを提供するクエリ エンドポイントをサポートしていません。

SELECT ai_query(
    "databricks-meta-llama-3-1-70b-instruct",
    "Can you explain AI in ten words?"
  )

たとえば、次に示すのは REST API 使用時のチャット モデルで想定されている要求形式です。 外部モデルに対しては、特定のプロバイダーとエンドポイント構成で有効な追加のパラメータを含めることができます。 「その他のクエリ パラメーター」を参照してください。

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

REST API を使用して行う要求に対して想定されている応答形式は、次のとおりです。

{
  "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
}

埋め込みモデルに対してクエリを実行する

Foundation Model API によって利用可能となる gte-large-en モデルに対する埋め込み要求を次に示します。 この例は、Model Serving の機能 (Foundation Model API または外部モデル) のいずれかを使って利用できる埋め込みモデルのクエリに適用されます。

OpenAI クライアント

OpenAI クライアントを使用するには、モデル サービング エンドポイント名を model 入力として指定します。

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"
)

ワークスペースの外部で基礎モデルのクエリを実行するには、次に示すように、OpenAI クライアントを直接使用する必要があります。 次の例では、Databricks API トークンと openai がコンピューティングにインストールされていることを前提としています。 OpenAI クライアントを Databricks に接続するために、Databricks ワークスペース インスタンスも必要です。

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"
)

REST API

重要

次の例では、基盤モデルまたは外部モデルを提供する提供エンドポイントのクエリを実行するために REST API パラメーターを使っています。 これらのパラメーターはパブリック プレビューであり、定義が変更される可能性があります。 POST /serving-endpoints/{name}/invocations を参照してください。

curl \
-u token:$DATABRICKS_TOKEN \
-X POST \
-H "Content-Type: application/json" \
-d  '{ "input": "Embed this sentence!"}' \
https://<workspace_host>.databricks.com/serving-endpoints/databricks-gte-large-en/invocations

MLflow デプロイ SDK

重要

次の例では、MLflow デプロイ SDK の predict() API を使用します。

import mlflow.deployments

export DATABRICKS_HOST="https://<workspace_host>.databricks.com"
export DATABRICKS_TOKEN="dapi-your-databricks-token"

client = mlflow.deployments.get_deploy_client("databricks")

embeddings_response = client.predict(
    endpoint="databricks-gte-large-en",
    inputs={
        "input": "Here is some text to embed"
    }
)

Databricks Python SDK

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()
response = w.serving_endpoints.query(
    name="databricks-gte-large-en",
    input="Embed this sentence!"
)
print(response.data[0].embedding)

LangChain

LangChain で Databricks Foundation Model API モデルを埋め込みモデルとして使用するには、DatabricksEmbeddings クラスをインポートし、次のように endpoint パラメーターを指定します。

%pip install langchain-databricks

from langchain_databricks import DatabricksEmbeddings

embeddings = DatabricksEmbeddings(endpoint="databricks-gte-large-en")
embeddings.embed_query("Can you explain AI in ten words?")

SQL

重要

次の例では、組み込みの SQL 関数である ai_query を使用します。 この関数はパブリック プレビューであり、定義が変更される可能性があります。 「ai_query() を使用した提供モデルに対するクエリの実行」を参照してください。

SELECT ai_query(
    "databricks-gte-large-en",
    "Can you explain AI in ten words?"
  )

次に示すのは埋め込みモデルで想定されている要求形式です。 外部モデルに対しては、特定のプロバイダーとエンドポイント構成で有効な追加のパラメータを含めることができます。 「その他のクエリ パラメーター」を参照してください。

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

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

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

埋め込み処理が正規化されているかどうかを確認する

モデルによって生成された埋め込み値が正規化されているかどうかを確認するには、次のコマンドを使用します。

  import numpy as np

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

テキスト入力候補モデルに対してクエリを実行する

OpenAI クライアント

重要

OpenAI クライアントを使用した Foundation Model API のトークンごとの支払いを使用して使用できるテキスト補完モデルのクエリはサポートされていません。 このセクションで示すように、OpenAI クライアントを使用した外部モデルのクエリのみがサポートされています。

OpenAI クライアントを使用するには、モデル サービング エンドポイント名を model 入力として指定します。 次の例では、OpenAI クライアントを使って、Anthropic によってホストされている claude-2 入力候補モデルにクエリを実行します。 OpenAI クライアントを使うには、model フィールドに、クエリ対象のモデルをホストするモデル提供エンドポイントの名前を設定します。

この例では、Anthropic モデル プロバイダーから外部モデルにアクセスするように構成された、以前に作成したエンドポイント anthropic-completions-endpoint を使用します。 外部モデル エンドポイントを作成する方法を参照してください。

クエリを実行できるその他のモデルとそのプロバイダーについては、「サポートされているモデル」を参照してください。

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)

REST API

外部モデルを使用して使用可能になった入力候補モデルのクエリを実行するための入力候補要求を次に示します。

重要

次の例では、REST API パラメーターを使用して、外部モデルにサービスを提供するエンドポイントにクエリを実行します。 これらのパラメーターはパブリック プレビューであり、定義が変更される可能性があります。 POST /serving-endpoints/{name}/invocations を参照してください。

curl \
-u token:$DATABRICKS_TOKEN \
-X POST \
-H "Content-Type: application/json" \
-d '{"prompt": "What is a quoll?", "max_tokens": 64}' \
https://<workspace_host>.databricks.com/serving-endpoints/<completions-model-endpoint>/invocations

MLflow デプロイ SDK

外部モデルを使用して使用可能になった入力候補モデルのクエリを実行するための入力候補要求を次に示します。

重要

次の例では、MLflow デプロイ SDK の predict() API を使用します。

import os
import mlflow.deployments

# Only required when running this example outside of a Databricks Notebook

os.environ['DATABRICKS_HOST'] = "https://<workspace_host>.databricks.com"
os.environ['DATABRICKS_TOKEN'] = "dapi-your-databricks-token"

client = mlflow.deployments.get_deploy_client("databricks")

completions_response = client.predict(
    endpoint="<completions-model-endpoint>",
    inputs={
        "prompt": "What is the capital of France?",
        "temperature": 0.1,
        "max_tokens": 10,
        "n": 2
    }
)

# Print the response
print(completions_response)

Databricks Python SDK

T:外部モデルを使用して使用可能になった入力候補モデルのクエリを実行するための入力候補要求を次に示します。

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()
response = w.serving_endpoints.query(
    name="<completions-model-endpoint>",
    prompt="Write 3 reasons why you should train an AI model on domain specific data sets."
)
print(response.choices[0].text)

SQL

重要

次の例では、組み込みの SQL 関数である ai_query を使用します。 この関数はパブリック プレビューであり、定義が変更される可能性があります。 「ai_query() を使用した提供モデルに対するクエリの実行」を参照してください。

SELECT ai_query(
    "<completions-model-endpoint>",
    "Can you explain AI in ten words?"
  )

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

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

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

{
  "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 プレイグラウンドを使ってサポートされているモデル LLM とチャットする

AI プレイグラウンドを使って、サポートされている大規模言語モデルを操作できます。 AI プレイグラウンドは、Azure Databricks ワークスペースから LLM をテスト、ダイアログを表示、比較できるチャットのような環境です。

その他のリソース

• モデルの監視とデバッグのための推論テーブル
• Foundation Model API を使用したバッチ推論のプロビジョニング済みスループット
• Databricks Foundation Model API
• Mosaic AI Model Serving の外部モデル
• チュートリアル: 外部モデル エンドポイントを作成して OpenAI モデルにクエリを実行する
• トークン単位の支払いでサポートされているモデル
• Foundation Model REST API リファレンス