CCCMKホールディングス TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

Databricksの"Foundation Model API"と"AI Functions"が面白かったのでまとめてみます!

こんにちは、CCCMKホールディングスTECH LABの三浦です。

10月1日です。この日はちょうどこの会社に入社した日なので、毎年「これで〇年働いたのかー」と考えるのですが、今年はちょうど10年なので、例年よりも少し感慨深い気持ちになりました。

はじめに

さて、機械学習やデータエンジニアリングでよく利用するDatabricksですが、Databricks社が毎年開催しているDATA+AI Summitというイベントがあり、私も昨年現地参加をしました。

techblog.cccmkhd.co.jp

今年は現地参加せずに動画で新しいリリースなどを見ていたのですが、キーノートで紹介されていたMosaic AIという機械学習や生成AI関連のサービスのデモがとても印象に残り、いつか自分でも試してみたいな、と考えていました。

こちらの動画の約1時間くらいからデモを見ることが出来ます。

youtu.be

このデモではテーブルデータから必要な情報を抽出する関数、プロンプトから画像を生成する関数、SNSにメッセージを投稿する関数をベースLLMに与え、ユーザーからこういう商品のプロモーションをSNSで打ち出したい、といった指示を受け取ると商品の売上情報を集計し、ターゲットに適した画像を生成し、SNSのメッセージを生成し、最終的にSNSへの投稿を実現する、といったような一連の処理を実行するAgentがとてもスピーディーに作れる様子が紹介されています。

このAgentを構築するために重要な役割を果たしているDatabricksの機能が"Foundation Model API"です。様々なオープンソースのLLMがDatabricksで高速に動作するように最適化され、それらをAPIで利用することが出来ます。Foundation Model APIによってさまざまなLLMをチャット形式で使ってみたり、それらのLLMを推論エンジンにした処理をSQLと組み合わせて使うことが出来る"AI Functions"が利用できるようになったりします。

今回はFoundation Model APIとAI Functionsについて試したことをまとめてみたいと思います。

Foundation Model API

提供されたLLMの利用

Foundation Model APIはDatabricksの"Model Serving"というサーバレスサービスをベースに提供される、様々な基盤モデル(LLM)を環境を独自に作ることなくAPIを経由してすぐに利用できるサービスです。

learn.microsoft.com

Foundation Model APIは、現在Azure Databricksでは利用出来るリージョンが限られており、2024年10月1日時点では日本リージョンは東西含めてまだ提供されていません。今回は"eastus"リージョンで機能の検証を行いました。

さてDatabricksのWorkspaceにアクセスし、左ペインから"Machine Learning"-"Serving"と進むと、Model Servingを利用して提供される機械学習モデル一覧が表示され、その中にFoundation Model APIとして提供される複数のLLMも確認することが出来ます。

Model ServingのEndpoint一覧とFoundation Model API

これらは"par-per-token"という料金体系で利用することが出来ます。発生するコストは入力トークン数と出力トークン数に依存し、詳細はこちらのページに記載されています。

www.databricks.com

何か使いたいモデルを選択し、"Use"ボタンをクリックすると"Playground"というチャット形式の画面に遷移し、選択したLLMにテキストを入力し、出力結果を確認することが出来ます。

Playground

それ以外のLLMを使いたい場合

Foundation Model APIで提供されていないLLM、たとえば"gpt-4o"などもModel Servingを使ってFoundation Model APIと同様にDatabricks上で利用することが出来ます。

Model Servingの画面で"Create serving endpoint"をクリックします。

Create serving endpoint

例えばAzure OpenAI ServiceのLLMを使いたい場合は"Served entities"から"Source"を"External model","Provider"を"OpenAI"に。

SourceとProviderを指定する。

Entity detailsでは"OpenAI API type"で"azure"を選択。あとはAPI KeyやAPI base等の必要な情報を入力します。

OpenAI API typeをazureにする。

その他endpointの"Name"など必要な情報を入力後、"Create"をクリックすると、Model Servingの一覧に追加され、Foundation Model APIと同様、Playgroundで試すことが出来ました。

Model Serving Endpointにgpt-4oが追加されました。

AI Functions

さて、Foundation Model APIはPlaygroundで利用するだけでなく、AI FunctionsとしてSQLと組み合わせて利用することも出来ます

DatabricksにおけるAI Functionsは組み込み済みのSQL関数で、SQLクエリの中でAIを呼び出し、データ抽出と同時にAIによる推論も実行出来ます。

※Azure DatabricksではAI Functionsは2024/10/01現在Public Previewの状態です。

Foundation Model APIを利用するAI Functions

AI Functionsはいくつか用意されていますが、その中で以下のものについてはFoundation Model APIを利用して機能します。

  • ai_analyze_sentiment
  • ai_classify
  • ai_extract
  • ai_fix_grammar
  • ai_gen
  • ai_mask
  • ai_similarity
  • ai_summarize
  • ai_translate

呼び出されるLLMは"Mixtral-8x7B Instruct"で、これらのAI Functionsは現在USリージョンでのみ利用可能とのことです。

この中でai_analyze_sentimentを実際に試してみました。検証に使用したデータはHugging Faceで公開されているこちらの商品レビューデータセットです。

huggingface.co

Notebookを開き、次のコードを実行してテーブルを作成しました。

from datasets import load_dataset

dataset = load_dataset(
    'mteb/amazon_reviews_multi', 
    'ja', 
    split=['train','test']
)

test_df = spark.createDataFrame(dataset[1].to_pandas())
test_df.write.saveAsTable(
    "amazon_reviews_test",
    mode="overwrite"
)

テーブル作成後、次のようなSQLをspark.sqlで実行しました。

spark.sql(
    """
    SELECT 
        text, 
        ai_analyze_sentiment(text) as sentiment 
    FROM (
        SELECT 
            text 
        FROM 
            amazon_reviews_test
        TABLESAMPLE(10 ROWS)
    )
    """
).show()

このコードのai_analyze_sentiment(text) as sentimentai_analyze_sentiment関数を呼び出し、テーブル内のtext(レビュー本文が格納されているカラム)を受け取り、感情分析を実行し、その結果をsentimentというカラムに出力してくれます。

実行すると、以下の様に感情分析結果が付与されていることが確認出来ました。

&ai_analyze_sentimentの実行結果

Model Servingで提供されるLLMを利用したAI Functions

Model Servingで提供される外部のLLMや、"Mixtral-8x7B Instruct"以外のFoundation Modelを使いたい場合はai_queryというAI Functionを利用すると実現することが出来ます。

たとえば"databricks-meta-llama-3-1-70b-instruct"というFoundation Modelを利用して商品レビューからSNS投稿用のタグを作らせる、といったことが出来ます。

spark.sql(
    """
    SELECT 
        text,
        ai_query(
            "databricks-meta-llama-3-1-70b-instruct",\
                "この文章からSNS用のタグを作成して:" || text
        ) AS TAGS
    FROM (
        SELECT 
            text 
        FROM 
            amazon_reviews_test
        TABLESAMPLE(10 ROWS)
    )
    """
).show()

SQLの中にLLMへのプロンプトを書くってなかなか新鮮だなぁと感じました。結果の一部を抜き出したのがこちらです。

生成されたタグ

これはけっこう使える用途が多そうです!テキスト系のデータはこれまでDelta Tableとして格納するよりファイルでそのまま読み込むことが多かったのですが、AI Functionsを活用するためにDelta Tableとしてテキストデータも格納しておくのも良さそうだな、と感じました。

まとめ

今回はDatabricksで提供されているFoundation Model APIやAI Functionsについてまとめてみました。Azure Databricksではまだ日本リージョンでの利用が出来ないことが残念ですが、LLMとデータを掛け合わせたアプリケーションの開発を容易にしてくれる便利な機能だと思いました。テキスト系のデータ加工はAI Functionsでかなり対応出来そうで、たとえば人名や地名などをマスキングする、といった用途でも活用出来ます。

Databricksには生成AIに関連する機能が他にもたくさん追加されているので、そちらも追って試していきたいと思います。