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


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

https://qiita.com/taka_yayoi/items/431fa69430c5c6a5e741#%E9%96%A2%E6%95%B0%E3%83%99%E3%83%BC%E3%82%B9%E3%81%AE%E3%83%95%E3%83%AC%E3%83%BC%E3%83%90%E3%83%BC%E3%81%AB%E3%82%88%E3%82%8Bopenai-api%E3%81%AE%E3%82%B5%E3%83%9D%E3%83%BC%E3%83%88

機能の説明はこちらに。Credential management for OpenAI on Databricksです。これは、Databricksの[モデルサービングエンドポイント](https://qiita.com/taka_yayoi/items/f6c8c11333761f96164f)経由でOpenAI APIを呼び出す際に、[Databricksシークレット](https://qiita.com/taka_yayoi/items/338ef0c5394fe4eb87c0)に保存しているAPIキーを読み取るというものです。

https://mlflow.org/docs/latest/python_api/openai/index.html

# OpenAI APIキーの設定

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

https://platform.openai.com/account/api-keys

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

```sh:Bash
databricks secrets put --scope demo-token-takaaki.yayoi --key openai_api_key
```

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

```py: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`に、上で作成したシークレットのスコープ名を設定します。

```py:Python
import os
import mlflow
import openai

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

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

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

```py: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",
    )
```

モデルがロギングされます。
![Screenshot 2023-05-05 at 10.40.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/2dce2255-4a14-68c3-aedc-c33edea8ec93.png)

なお、上で環境変数`MLFLOW_OPENAI_SECRET_SCOPE`が設定されていると`openai.yaml`が一緒にロギングされます。このファイルにどのシークレットを参照するのかが記録されます。
![Screenshot 2023-05-05 at 10.40.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/68fb0763-7d84-5f29-7318-e85ca531f318.png)

# モデルサービング

あとはこれをモデルレジストリに登録してモデルサービングされるようにします。
![Screenshot 2023-05-05 at 10.43.00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/46d25532-d95c-3a24-4c6a-946bf62ae031.png)
![Screenshot 2023-05-05 at 10.43.14.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/c10f4200-9cf5-eacd-3e52-b891984cea5d.png)

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

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

回答が返ってきました！
![Screenshot 2023-05-05 at 10.45.23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0e0bfba7-9042-8052-fcda-458c804a1ee1.png)

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

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

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

```py: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",
    )
```

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

日本語で応対してくれるようになりました。
![Screenshot 2023-05-05 at 10.57.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0af48ff8-b3dc-fcdc-b029-022fed1c66f1.png)

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

```py: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:Bash
streamlit run chatbot.py
```

画面経由でOpenAI APIを活用できるようになりました！
![Screenshot 2023-05-05 at 11.35.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3e433bc0-9c34-e9bd-a529-58e0031a6f23.png)

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

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

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)