生成AIモデルへのクエリー

プレビュー

Mosaic AI Model Serving はパブリックプレビュー段階にあり、 us-east1 と us-central1でサポートされています。

この記事では、外部モデルのクエリリクエストをフォーマットし、モデルサービングエンドポイントに送信する方法について説明します。

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

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

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

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

要件

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

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

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

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

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

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

dbutils.library.restartPython()

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

!pip install mlflow

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

dbutils.library.restartPython()

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

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

チャットモデルに対してクエリを実行する例を次に示します。この例は、外部モデルを使用して利用可能になったチャットモデルのクエリに適用されます。

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="bedrock-chat-completions-endpoint",
    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を参照してください。

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/<your-external-model-endpoint>/invocations \

重要

次の例では、 MLflow Deployments 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="bedrock--chat-completions-endpoint",
    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 ノートブックから 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="bedrock-chat-completions-endpoint",
    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}")

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

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

REST API を使用して行われたリクエストに対する予想される応答形式は次のとおりです。

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

エンべディングモデルのクエリー

次の例は、外部モデルによって利用可能になった gte-large-en モデルの埋め込み要求です。

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

from databricks.sdk import WorkspaceClient

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

response = openai_client.embeddings.create(
  model="cohere-embeddings-endpoint",
  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="cohere-embeddings-endpoint",
  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/<your-embedding-model-endpoint>/invocations

重要

次の例では、 MLflow Deployments 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="cohere-embeddings-endpoint",
    inputs={
        "input": "Here is some text to embed"
    }
)

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

w = WorkspaceClient()
response = w.serving_endpoints.query(
    name="cohere-embeddings-endpoint",
    input="Embed this sentence!"
)
print(response.data[0].embedding)

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

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

重要

次の例では、 MLflow Deployments 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="openai-completions-endpoint",
    inputs={
        "prompt": "What is the capital of France?",
        "temperature": 0.1,
        "max_tokens": 10,
        "n": 2
    }
)

# Print the response
print(completions_response)

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

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

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

{
  "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モデルへのクエリー

要件

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

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

エンべディングモデルのクエリー

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

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

関連リソース