クエリ基盤モデル

この記事では、基盤モデルのクエリ リクエストをフォーマットし、モデルサービング エンドポイントに送信する方法を学びます。

従来の機械学習または Python モデルのクエリ リクエストについては、 「カスタムモデルのサービングエンドポイントをクエリする」を参照してください。

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

Databricks モデルサービングは、基盤モデルにスコアリング リクエストを送信するための次のオプションを提供します。

手法

詳細

OpenAI クライアント

OpenAI クライアントを使用して、Databricks モデルビング エンドポイントによってホストされているモデルをクエリします。 モデルサービング エンドポイント名をmodel入力として指定します。 基盤APIsまたは外部モデルによって利用できるチャット、埋め込み、および補完モデルがサポートされています。

UI の提供

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

REST API

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

MLflow デプロイメント SDK

MLflow Deployments SDK のdetect()関数を使用してモデルをクエリします。

Databricks GenAI SDK

Databricks GenAI SDK は、REST API 上のレイヤーです。 認証やモデル ID のエンドポイント URL へのマッピングなど、低レベルの詳細を処理するため、モデルとの対話が容易になります。 SDK は、Databricks ノートブック内から使用できるように設計されています。

SQL関数

ai_query SQL 関数を使用して、SQL からモデル推論を直接呼び出します。 「 ai_query() を使用して提供されたモデルをクエリする」を参照してください。

要件

重要

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

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

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

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

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

!pip install openai

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

dbutils.library.restartPython()

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

!pip install mlflow

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

dbutils.library.restartPython()
 !pip install databricks-genai

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

 dbutils.library.restartPython()

チャット完了モデルのクエリ

以下は、チャットモデルをクエリする例です。

バッチ推論の例については、 「基盤モデルAPIsを使用したバッチ推論」を参照してください。

以下は、ワークスペースの基盤APIsペイパー ブレイク エンドポイント databricks-dbrx-instruct によって利用可能になる DBRX Instruct モデルに対するチャット リクエストです。

OpenAI クライアントを使用するには、モデルサービングのエンドポイント名をmodel入力として指定します。 次の例では、コンピュート上にDatabricks APIopenai がインストールされていることを前提としています。 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 パラメーターを使用して、基盤モデルにサービスを提供するサービス エンドポイントをクエリします。 これらのパラメーターはパブリック プレビューであり、定義は変更される可能性があります。 POST /serving-endpoints/{name}/invocations を参照してください。

以下は、ワークスペースの基盤APIsペイパー ブレイク エンドポイント 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 Deployments SDKpredict() API を使用します。

以下は、ワークスペースの基盤APIsペイパー ブレイク エンドポイント 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
    }
)

以下は、ワークスペースの基盤APIsペイパー ブレイク エンドポイント databricks-dbrx-instruct によって利用可能になる DBRX Instruct モデルに対するチャット リクエストです。

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 を使用して基盤モデル エンドポイントにクエリを実行するには、次のいずれかを実行できます。

  • Databricks LLM クラスをインポートし、 endpoint_nametransform_input_fnを指定します。

  • ChatDatabricks ChatModel クラスをインポートし、endpointを指定します。

次の例では、LangChain の Databricks LLM クラスを使用して、基盤モデルAPIsペイパートークン エンドポイント databricks-dbrx-instruct をクエリします。 基盤モデルAPIsリクエスト ディクショナリに 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 関数AIを使用します。 この関数は パブリック プレビューであり 、定義は変更される可能性があります。 「 ai_query() を使用して提供されたモデルをクエリする」を参照してください。

以下は、ワークスペースの基盤モデルAPIsペイパーベンチャー エンドポイント databricks-llama-2-70b-chat によって利用可能になる、llama-2-70b-chat へのチャット リクエストです。

注:

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クラスを提供します。 以下の機能を提供します。

関数

帰る

説明

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?"}
# ]

埋め込みモデルのクエリ

以下は、基盤モデルAPIsによって利用可能になった bge-large-en モデルの埋め込みリクエストです。

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 パラメーターを使用して、基盤モデルにサービスを提供するサービス エンドポイントをクエリします。 これらのパラメーターはパブリック プレビューであり、定義は変更される可能性があります。 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-bge-large-en/invocations

重要

次の例では、 MLflow Deployments SDKpredict() 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"
    }
)

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 で Databricks 基盤モデルAPIsモデルを埋め込みモデルとして使用するには、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 関数AIを使用します。 この関数は パブリック プレビューであり 、定義は変更される可能性があります。 「 ai_query() を使用して提供されたモデルをクエリする」を参照してください。

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

テキスト補完モデルのクエリ

以下は、基盤モデルAPIsによって利用可能になった databricks-mpt-30b-instruct モデルの完了リクエストです。 問題と構文については、 「完了タスク」を参照してください。

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

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

重要

次の例では、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/databricks-mpt-30b-instruct/invocations

重要

次の例では、 MLflow Deployments SDKpredict() 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
    }
)

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 関数AIを使用します。 この関数は パブリック プレビューであり 、定義は変更される可能性があります。 「 ai_query() を使用して提供されたモデルをクエリする」を参照してください。

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 Playground を使用してサポートされている LLM とチャットする

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

AIの遊び場