基盤モデルに対してクエリを実行する

[アーティクル]
05/21/2024

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

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

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

チャット入力候補モデルに対してクエリを実行する
埋め込みモデルに対してクエリを実行する
テキスト入力候補モデルに対してクエリを実行する

Databricks Model Serving には、基盤モデルにスコアリング要求を送信するための以下のオプションが用意されています。

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

要件

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

重要

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

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

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

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

OpenAI クライアント

OpenAI クライアントを使用するには、クラスターに openai パッケージをインストールする必要があります。ノートブックまたはローカルターミナルで次を実行してください。

!pip install 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 GenAI SDK

 !pip install databricks-genai

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

 dbutils.library.restartPython()

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

チャットモデルのクエリの例を次に示します。

バッチ推論の例については、Foundation Model API を使用したバッチ推論に関する記事を参照してください。

OpenAI クライアント

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

OpenAI クライアントを使用するには、モデル提供エンドポイント名を model 入力として指定します。次の例では、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.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 を参照してください。

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 を使用します。


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 GenAI SDK

from databricks_genai_inference import ChatCompletion

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

response = ChatCompletion.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=128)
print(f"response.message:{response.message}")

LangChain

LangChain を使用して基盤モデルエンドポイントにクエリを実行するには、次のいずれかを行うことができます。

Databricks LLM クラスをインポートし、endpoint_name と transform_input_fn を指定してください。
ChatDatabricks ChatModel クラスをインポートし、endpoint を指定してください。

次の例では、LangChain の Databricks LLM クラスを使用して、Foundation Model API のトークン単位の支払いエンドポイントである databricks-dbrx-instruct にクエリを実行します。 Foundation Model API は要求ディクショナリの messages を想定していますが、LangChain Databricks LLM は既定で要求ディクショナリの prompt を提供します。 transform_input 関数を使用して、要求ディクショナリを想定される形式に準備してください。

from langchain.llms import Databricks
from langchain_core.messages import HumanMessage, SystemMessage

def transform_input(**request):
  request["messages"] = [
    {
      "role": "user",
      "content": request["prompt"]
    }
  ]
  del request["prompt"]
  return request

llm = Databricks(endpoint_name="databricks-dbrx-instruct", transform_input_fn=transform_input)
llm("What is a mixture of experts model?")

次の例では、ChatDatabricks ChatModel クラスを使用し、endpoint を指定します。

from langchain.chat_models import ChatDatabricks
from langchain_core.messages import HumanMessage, SystemMessage

messages = [
    SystemMessage(content="You're a helpful assistant"),
    HumanMessage(content="What is a mixture of experts model?"),
]
chat_model = ChatDatabricks(endpoint="databricks-dbrx-instruct", max_tokens=500)
chat_model.invoke(messages)

SQL

重要

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

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

Note

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

SELECT ai_query(
    "databricks-llama-2-70b-chat",
    "Can you explain AI in ten words?"
  )

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

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

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

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

チャットセッション

Databricks GenAI SDK には、マルチラウンドチャット会話を管理するための ChatSession クラスが用意されています。次の機能が提供されます。

関数	Return	説明
`reply (string)`		新しいユーザーメッセージを受け取ります
`last`	string	アシスタントからの最後のメッセージ
`history`	dict のリスト	チャット履歴のメッセージ (ロールを含む)。
`count`	int	これまでに実行されたチャットラウンドの数。

ChatSession を初期化するには、ChatCompletion と同じ引数のセットを使用します。これらの引数はチャットセッション全体で使用されます。


from databricks_genai_inference import ChatSession

chat = ChatSession(model="llama-2-70b-chat", system_message="You are a helpful assistant.", max_tokens=128)
chat.reply("Knock, knock!")
chat.last # return "Hello! Who's there?"
chat.reply("Guess who!")
chat.last # return "Okay, I'll play along! Is it a person, a place, or a thing?"

chat.history
# return: [
#     {'role': 'system', 'content': 'You are a helpful assistant.'},
#     {'role': 'user', 'content': 'Knock, knock.'},
#     {'role': 'assistant', 'content': "Hello! Who's there?"},
#     {'role': 'user', 'content': 'Guess who!'},
#     {'role': 'assistant', 'content': "Okay, I'll play along! Is it a person, a place, or a thing?"}
# ]

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

Foundation Model API によって利用可能となる bge-large-en モデルに対する埋め込み要求を次に示します。

OpenAI クライアント

OpenAI クライアントを使用するには、モデル提供エンドポイント名を model 入力として指定します。次の例では、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.embeddings.create(
  model="databricks-bge-large-en",
  input="what is databricks"
)

REST API

重要


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-bge-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-bge-large-en",
    inputs={
        "input": "Here is some text to embed"
    }
)

Databricks GenAI SDK


from databricks_genai_inference import Embedding

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

response = Embedding.create(
    model="bge-large-en",
    input="3D ActionSLAM: wearable person tracking in multi-floor environments")
print(f'embeddings: {response.embeddings}')

LangChain

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

from langchain.embeddings import DatabricksEmbeddings

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

SQL

重要


SELECT ai_query(
    "databricks-bge-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
  }
}

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

次に、Foundation Model API によって使用可能になった databricks-mpt-30b-instruct モデルの補完要求を示します。パラメーターと構文については、「完了タスク」を参照してください。

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

completion = client.completions.create(
  model="databricks-mpt-30b-instruct",
  prompt="what is databricks",
  temperature=1.0
)

REST API

重要


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/databricks-mpt-30b-instruct/invocations

MLflow デプロイ SDK

重要

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


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

completions_response = client.predict(
    endpoint="databricks-mpt-30b-instruct",
    inputs={
        "prompt": "What is the capital of France?",
        "temperature": 0.1,
        "max_tokens": 10,
        "n": 2
    }
)

Databricks GenAI SDK


from databricks_genai_inference import Completion

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

response = Completion.create(
    model="databricks-mpt-30b-instruct",
    prompt="Write 3 reasons why you should train an AI model on domain specific data sets.",
    max_tokens=128)
print(f"response.text:{response.text:}")

SQL

重要

SELECT ai_query(
    "databricks-mpt-30b-instruct",
    "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 をテスト、ダイアログを表示、比較できるチャットのような環境です。

AI プレイグラウンド

次の方法で共有

基盤モデルに対してクエリを実行する

要件

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

LangChain

SQL

チャットセッション

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

LangChain

SQL

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

SQL

AI プレイグラウンドを使ってサポートされているモデル LLM とチャットする

その他のリソース

フィードバック

フィードバック

その他のリソース

次の方法で共有

基盤モデルに対してクエリを実行する

要件

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

LangChain

SQL

チャット セッション

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

LangChain

SQL

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

OpenAI クライアント

REST API

MLflow デプロイ SDK

Databricks GenAI SDK

SQL

AI プレイグラウンドを使ってサポートされているモデル LLM とチャットする

その他のリソース

フィードバック

フィードバック

その他のリソース

チャットセッション