埋め込みモデルのクエリを実行する

この記事では、埋め込みタスク用に最適化され、 Unity AI Gateway によって提供される基盤モデルのクエリ要求を記述する方法について説明します。

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

• Foundation Models API、これはDatabricks でホストされる基盤モデルと呼ばれる。
• Databricks の外部でホストされる基盤モデルと呼ばれる外部モデル。

要求事項

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

クエリのサンプル

Show query examples that send embedding requests to the databricks-gte-large-en model using the DatabricksOpenAI helper, the plain OpenAI client, and a curl REST API call.

以下は、さまざまなgte-large-enを使用して、Foundation Model API によってトークンごとの支払いを行う モデルに対する埋め込み要求です。

OpenAI クライアント

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

from databricks_openai import DatabricksOpenAI

client = DatabricksOpenAI()

response = client.embeddings.create(
  model="system.ai.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/ai-gateway/mlflow/v1"
)

response = client.embeddings.create(
  model="system.ai.gte-large-en",
  input="what is databricks"
)

SQL

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

REST API

curl \
-u token:$DATABRICKS_TOKEN \
-X POST \
-H "Content-Type: application/json" \
-d  '{
  "model": "system.ai.gte-large-en",
  "input": "Embed this sentence!"
}' \
https://<workspace_host>.databricks.com/ai-gateway/mlflow/v1/embeddings

MLflow デプロイ SDK

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="system.ai.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="system.ai.gte-large-en",
    input="Embed this sentence!"
)
print(response.data[0].embedding)

LangChain

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

%pip install databricks-langchain

from databricks_langchain import DatabricksEmbeddings

embeddings = DatabricksEmbeddings(endpoint="system.ai.gte-large-en")
embeddings.embed_query("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 モデルの種類 」を参照してください。

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

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

  import numpy as np

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

その他のリソース

• チャット モデルのクエリを実行します。
• クエリ推論モデル。
• 視覚モデルのクエリを実行します。