DatabricksでMLflow 2.3のOpenAI APIのサポートを試す

こちらのOpenAIのサポートを試します。一番後回しになったのは、Databricksでのモデルサービングの機能デプロイを待っていたためでした。今日から利用できます。

機能の説明はこちらに。Credential management for OpenAI on Databricksです。これは、Databricksのモデルサービングエンドポイント経由でOpenAI APIを呼び出す際に、Databricksシークレットに保存しているAPIキーを読み取るというものです。

OpenAI APIキーの設定

当たり前ではありますがAPIキーが必要です。以下のサイトでAPIキーを発行します。

これを平文で使用するのはリスクが高いので、シークレットに格納します。この作業はDatabricks CLIが設定されているローカルマシンで実行します。以下の例ではシークレットのスコープはdemo-token-takaaki.yayoi、キーはopenai_api_keyです。スコープ名は任意ですがキー名は固定です。

Bash

databricks secrets put --scope demo-token-takaaki.yayoi --key openai_api_key

エディタが開きますので、上で取得したAPIキーを入力して保存します。

Python

import os
import mlflow
import openai

os.environ["MLFLOW_OPENAI_SECRET_SCOPE"] = "demo-token-takaaki.yayoi"

モデルのロギング

Databricksノートブックを作成して、MLflowやOpenAIなどをインストールします。

%pip install tiktoken
%pip install openai
%pip install mlflow==2.3

次のステップが重要です。環境変数MLFLOW_OPENAI_SECRET_SCOPEに、上で作成したシークレットのスコープ名を設定します。

Python

import os
import mlflow
import openai

os.environ["MLFLOW_OPENAI_SECRET_SCOPE"] = "demo-token-takaaki.yayoi"

この環境変数が設定されている状態でOpenAIのモデルをロギングすると、モデルサービングエンドポイントでOpenAI APIを呼び出す際にシークレットに格納されているAPIキーを読み込むようになります。

ここではMLflowのエキスパートになってくださいとお願いしています。

Python

with mlflow.start_run():
    model_info = mlflow.openai.log_model(
        model="gpt-3.5-turbo",
        task=openai.ChatCompletion,
        messages=[{"role": "system", "content": "You are an MLflow expert"}],
        artifact_path="model",
    )

モデルがロギングされます。

なお、上で環境変数MLFLOW_OPENAI_SECRET_SCOPEが設定されているとopenai.yamlが一緒にロギングされます。このファイルにどのシークレットを参照するのかが記録されます。

モデルサービング

あとはこれをモデルレジストリに登録してモデルサービングされるようにします。

デプロイが完了したら右上のクエリーエンドポイント(エンドポイントに対するクエリーなので後で文言修正しておきます)をクリックします。ここでクイックにモデルに対する問い合わせを行うことができます。

JSON

{"inputs":
{"role": ["user"], "content": ["What is MLflow?"]}
}

回答が返ってきました！

カスタムチャットボットの作成

OpenAI APIを呼び出すだけですのでモデルのトレーニングを行なっている訳ではないのですが、MLflowと組み合わせることでOpenAI APIを活用したいろいろなカスタムアプリケーションを構築する際に捗ります。

日本語で対応するようにしてみます。

Python

with mlflow.start_run():
    model_info = mlflow.openai.log_model(
        model="gpt-3.5-turbo",
        task=openai.ChatCompletion,
        messages=[{"role": "system", "content": "あなたはDatabricksの専門家です"}],
        artifact_path="model",
    )

上と同じようにこのモデルをロギングし、モデルレジストリへの登録、モデルサービングを行います。

日本語で応対してくれるようになりました。

さらにGUIを作成します。いつものようにstreamlit使います。

chatbot.py

import streamlit as st 
import numpy as np 
from PIL import Image
import base64
import io

import os
import requests
import numpy as np
import pandas as pd

import json

st.header('OpenAI Chatbot on Databrikcs')

def score_model(prompt):
  # 1. パーソナルアクセストークンを設定してください
  # 今回はデモのため平文で記載していますが、実際に使用する際には環境変数経由で取得する様にしてください。
  token = "<Databricksパーソナルアクセストークン>"
  #token = os.environ.get("DATABRICKS_TOKEN")

  # 2. モデルエンドポイントのURLを設定してください
  url = '<モデルエンドポイントURL>'
  headers = {'Authorization': f'Bearer {token}'}
  
  data_json_str = f"""
  {{"inputs":
{{"role": ["user"], "content": ["{prompt}"]}}
}}
"""

  data_json = json.loads(data_json_str)
  response = requests.request(method='POST', headers=headers, url=url, json=data_json)

  if response.status_code != 200:
    raise Exception(f'Request failed with status {response.status_code}, {response.text}')
  return response.json()

prompt = st.text_input("プロンプト")

if prompt != "":
  response = score_model(prompt)
  st.write(response['predictions'][0])

Bash

streamlit run chatbot.py

画面経由でOpenAI APIを活用できるようになりました！

足回りが整うと色々試行錯誤しやすくなるので、こういう機能は大好きです。

Databricksクイックスタートガイド

Databricksクイックスタートガイド

Databricks無料トライアル

Databricks無料トライアル