SnowparkMLを使ってみよう～LLM関数を添えて～

概要

この記事は2024年のSnowflakeアドベントカレンダーとして執筆したものです。

SnowflakeのLLM関数やSnowpark MLの機能が充実してきたらしい、という噂を聞いて試してみたくなりました。
今回は、

• LLM関数を使って自然言語のベクトル化(埋め込み)
• 分類問題を解く（機械学習）
• モデルをSnowflakeオブジェクトとして保存

というステップをやってみたいと思います。

実行環境

本記事のPythonコード/SQLはすべてSnowflakeのノートブックで実行しています。
PythonもSQLも同じノートブックで実行できてハッピー。

• ウェアハウス：X-small
• Python：3.9

Pythonライブラリ

• Snowflake：1.0.0
• Snowflake-ml-python:1.7.1
• pandas：2.2.2
• scikit-learn：1.5.1

実装内容

データ準備

今回使うデータはKaggleから取得した「Restaurant Reviews」というデータです。
まずはテーブルを作成して、データをロードします。
ML_DBのTESTスキーマにテーブルを作っていきます。

create or replace TABLE ML_DB.TEST.RAW_RESTAURANT_REVIEWS (
    REVIEW VARCHAR(200)
    ,LIKED NUMBER(38,0)
);

Snowparkを使ってデータを読み込む。遅延評価がグッドポイント。
get_active_sessionは現在のセッションを取得する関数です。Snowsight上で実行する場合はこちらが利用できますが、ローカルの環境で実行する時はSessionを新たに作る必要があります。

import snowflake.snowpark as snowpark
from snowflake.snowpark.functions import col
from snowflake.snowpark.context import get_active_session


def fetch_data_from_table(session,tableName):
    spdf = session.table(tableName)
    spdf.show()
    return spdf
    
session = get_active_session()   
table_name = "ML_DB.TEST.RAW_RESTAURANT_REVIEWS"
fetch_data_from_table(session, tableName=table_name)

レストランの口コミと「イイネ」が押されたかどうかのフラグがあります。
全部で1000レコードあります。

テキストをベクトル化

早速Snowflake CoretexのLLM関数を使ってみます。
自然言語の埋め込みにはsnowflake-arctic-embed-m-v1.5というモデルを使用します。
ベクトル化したデータは一旦テーブルに格納します。
Snowflakeはテーブル列の型としてVECTOR型をサポートしています（便利）。

create or replace table ML_DB.TEST.VEC_RESTAURANT_REVIEWS (
    REVIEW VARCHAR(200)
    ,LIKED NUMBER(38,0)
    ,VECTOR VECTOR(FLOAT, 768)
);

insert into ML_DB.TEST.VEC_RESTAURANT_REVIEWS
select
    REVIEW
    ,LIKED
    ,SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m', REVIEW) as VECTOR
from
    ML_DB.TEST.RAW_RESTAURANT_REVIEWS
;

EMBED_TEXT_768は2024年12月現在、日本語に対応していません。
日本語の処理自体は可能ですが、精度がイマイチな感じでした(筆者の感覚)

VEC_RESTAURANT_REVIEWSの中身↓

ベクトルの次元圧縮

それでは、作成したベクトルを使って機械学習をしましょう！…​…​
といきたいところなのですが、今回はこの後で作成するロジスティックのモデルをモデルレジストリーに保存する予定です。
2024年12月現在、モデルレジストリーに登録出来るモデルは「特徴量が500個以下」という制約があります。先程作ったベクトルは768次元なので、このままでは登録できません。
そこでPCAによって次元削減をします。

from sklearn.decomposition import PCA
import pandas as pd

tableName = "SNOWFLAKE_TEST.XS_TEST.VEC_RESTAURANT_REVIEWS"
spdf = fetch_data_from_table(session,tableName) # Snowpark DataFrame

df = spdf.to_pandas() # テーブルをpandasデータフレームに変換する
feat = df['VECTOR'].tolist()
    
pca = PCA(n_components=500)
feat_decomposed = pca.fit_transform(feat)

column_names = ['v' + str(i) for i in range(1, len(feat_decomposed[0])+1)]
X = pd.DataFrame(feat_decomposed,columns=column_names)
X.head(10)

最終的には500列のpandas dfの形にしました。

Snowpark MLを使ってモデル作成

ここからはロジスティック回帰で二値分類を行います。コードは普通にscikit-Learnを使う時とほとんど一緒ですが、データの与え方が少し違うので注意しましょう。

from sklearn.model_selection import train_test_split
from snowflake.ml.modeling.linear_model import LogisticRegression as lr

train_indices, test_indices = train_test_split(df.index, test_size=0.2, random_state=42)
df = pd.concat([df,X],axis=1)


# モデル作成
# ここの書き方がscikit-learnとちょっと違う
clf = lr(input_cols=column_names, label_cols='LIKED')
clf.fit(df.loc[train_indices])

# 推論
# predict()の出力もscikit-learnとちょっと違う
pred = clf.predict(df.loc[test_indices])
pred.head(10)

通常のScikit Learnも使えるのですが、Snowpark MLを使っておくとモデル登録時に楽をできるのでオススメです。
LightGBMやXGBoostも利用できます。
SnowparkMLは入力にSnowpark DataFrameを使うことが出来ます。テーブルから取得したままの形で機械学習モデルの構築ができるのは、めちゃくちゃメリットだと感じます。

推論結果出力↓

初めて見た時は何が起こっているのかわかりませんでした。入力したデータフレームがそのまま出力されてるやんけ！！！

実はpredict()を使うと、入力のデータフレームに予測結果を一列足した形のデータフレームが出力されます。
このOUTPUT_LIKEDが今回のロジスティック回帰モデルが出力した予測値になります。

機械学習モデルを保存

最後に作成したモデルを登録します。

from snowflake.ml.registry import Registry

# 保存
reg = Registry(session=session, database_name='ML_DB', schema_name='TEST')
mv = reg.log_model(clf, model_name="RESTAURANT_REVIEWS_LR",comment="register ml model test")

registerの引数にはsample_input_dataというものがあり、SnowparkML以外のモデルを登録する時には必須の引数です。そのモデルでどんな入力データを与える必要があるかを明示する必要があります。
SnowparkMLを使っておくと勝手に処理してくれるので、ちょっと楽ができます。

登録したモデルはSnowflakeのオブジェクトとして保管されます。このモデルを呼び出して推論や再学習も可能です。バージョン管理もできます。

Snowsight上ではこのモデルのメタデータを表示できたり、使用できるメソッドを一覧してくれたりします。

補足

モデルレジストリ―の機能は2024年12月時点でプレビュー機能となっており、いくつかの制約があります。
そのうちの一つが、途中で述べた「モデルに使用できる特徴量は500個まで」というものです。
今回の実験でもう一つ制約を受けたのが「モデルの容量制限」です。モデル登録時に一度ステージに構成ファイルが置かれるのですが、この構成ファイルのサイズが250KBを超えるとレジストリ―に登録できません。上記の500次元のベクトルを使ったモデルは登録できませんでしたが、300次元のベクトルで試したら登録することができました。

おまけ

せっかくなので精度を測ってみました。

from sklearn.metrics import roc_auc_score
score = roc_auc_score(df.loc[test_indices]['LIKED'], pred['OUTPUT_LIKED'])
print (score)　#出力→ 0.90745

次元圧縮を行いましたが、結構良い結果となったのではないでしょうか。
LLM関数はかなり手軽に扱えるので、工夫次第ではとても有用な機能になりそうだなと思います。

感想

LLM関数は他のクラウドサービスのように、LLMのエンドポイントやAPIを管理することなく利用できるのがとても便利でした。
SnowparkMLはまだ制約はあるものの、Snowflakeの機能のみでMLパイプラインを構成することができるようになるため手軽ですし、ちょっとしたプロダクトであれば実用できそうだなと感じました。

今年も一年お疲れ様でした！