埋め込みモデルのクエリ

この記事では、エンべディング タスク用に最適化された基盤モデルのクエリ要求を記述し、それらをモデルサービング エンドポイントに送信する方法について説明します。

この記事の例は、次のいずれかを使用して使用可能になる基盤モデルのクエリに適用されます。

• Databricksがホストする 基盤モデルAPI と呼ばれる 基盤モデル 。
• Databricks の外部でホストされている基盤モデル と呼ばれる外部モデル 。

必要条件

• 要件を参照してください。
• 選択したクエリ クライアント オプションに基づいて、適切なパッケージをクラスターにインストールします。

クエリの例

以下は、さまざまな クライアントオプションを使用して、基盤モデルAPIトークン単位の従量課金によって利用可能になったgte-large-enモデルのエンべディング要求です。

OpenAI client

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

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

ワークスペースの外部で基盤モデルをクエリするには、以下に示すように、OpenAI クライアントを直接使用する必要があります。次の例では、 Databricks API トークンと openai がコンピュートにインストールされていることを前提としています。 また、OpenAI クライアントを Databricks に接続するには、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"
)

SQL

important

次の例では、組み込み SQL 関数 ai_query を使用しています。この機能は パブリック プレビュー 段階であり、定義は変更される可能性があります。

SQL

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

REST API

important

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

Bash

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

important

次の例では、MLflow Deployments SDK の predict() API を使用しています。

Python

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

Python

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で 基盤モデルAPIモデルを使用するには、DatabricksEmbeddings クラスをインポートし、次のようにendpoint パラメーターを指定します。

Bash

%pip install databricks-langchain

Python

from databricks_langchain import DatabricksEmbeddings

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

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

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

追加のリソース

• チャット モデルに対してクエリを実行します。
• 推論モデルのクエリ。
• ビジョン モデルのクエリ。