DatabricksでLLMのクエリを使い始める

この記事では、Foundation Model APIsを使用してDatabricksでLLMを提供およびクエリする方法について説明します。

Databricks上でLLMモデルの提供やクエリを開始する最も簡単な方法は、基盤モデルAPIトークン単位の従量課金で使用することです。APIは、DatabricksワークスペースのサービングUIで自動的に利用可能なトークン単位の従量課金のエンドポイントから、一般的な基盤モデルへのアクセスを提供します。「トークン単位の従量課金でサポートされるモデル」を参照してください。

AI Playgroundを使用して、トークン単位の従量課金モデルをテストしたりチャットしたりすることもできます。AI Playgroundを使用してサポートされているLLMとチャットするを参照してください。

本番運用のワークロード、特に微調整されたモデルやパフォーマンス保証を必要とするワークロードの場合、Databricksはプロビジョニングされたスループットのエンドポイントで基盤モデルAPIを使用するようにアップグレードすることを推奨します。

要件

重要

運用シナリオにおけるセキュリティのベストプラクティスとして、Databricksでは本番環境での認証にマシン間OAuthトークンを使用することを推奨しています。

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

FoundationモデルAPIの使用を開始する

次の例では、トークン単位の従量課金のエンドポイントdatabricks-dbrx-instructで提供されるdatabricks-dbrx-instructモデルをクエリします。 DBRX Instructモデルの詳細については、こちらを参照してください。

この例では、OpenAIクライアントを使用して、クエリー対象のモデルをホストするモデルサービング エンドポイントの名前をmodelフィールドに入力することでモデルをクエリーします。個人用アクセストークンを使用してDATABRICKS_TOKENDatabricksワークスペースのインスタンスを入力し、OpenAI クライアントをDatabricksに接続します。

注:

以下のメッセージImportError: cannot import name 'OpenAI' from 'openai'が表示されたら、openaiのバージョンを!pip install -U openaiを使ってアップグレードしてください。

from openai import OpenAI
import os

DATABRICKS_TOKEN = os.environ.get("DATABRICKS_TOKEN")

client = OpenAI(
  api_key=DATABRICKS_TOKEN, # your personal access token
  base_url='https://<workspace_id>.databricks.com/serving-endpoints', # your Databricks workspace instance
)

chat_completion = client.chat.completions.create(
  messages=[
    {
      "role": "system",
      "content": "You are an AI assistant",
    },
    {
      "role": "user",
      "content": "What is a mixture of experts model?",
    }
  ],
  model="databricks-llama-3-instruct",
  max_tokens=256
)

print(chat_completion.choices[0].message.content)

期待される出力:

{
  "id": "xxxxxxxxxxxxx",
  "object": "chat.completion",
  "created": "xxxxxxxxx",
  "model": "databricks-dbrx-instruct",
  "choices": [
    {
      "index": 0,
      "message":
        {
          "role": "assistant",
          "content": "A Mixture of Experts (MoE) model is a machine learning technique that combines the predictions of multiple expert models to improve overall performance. Each expert model specializes in a specific subset of the data, and the MoE model uses a gating network to determine which expert to use for a given input."
        },
      "finish_reason": "stop"
    }
  ],
  "usage":
    {
      "prompt_tokens": 123,
      "completion_tokens": 23,
      "total_tokens": 146
    }
}

次のステップ