TabFM:ゼロショット表状基礎モデル

TabFM は、表形式データのGoogle Research財団モデルです。 このプログラムはコンテキスト内学習を採用しており、トレーニング行をコンテキストとして渡し、予測を単一のフォワードパスで行い、ファインチューニングやハイパーパラメータ検索、データセット固有のトレーニングを必要としません。 2進分類および多クラス分類(最大10クラス)と、数値列とカテゴリカラムが混在するテーブルでの回帰処理をサポートしています。

このノートブックは乳がんデータセットでゼロショット分類を、糖尿病データセットでゼロショット回帰を実行します。

Note

この例はDatabricks AI環境バージョン6以上が必要です。

サーバーレス GPU コンピューティングに接続する

[ 接続 ] ドロップダウンをクリックし、[ サーバーレス GPU] を選択します。 環境サイドパネルを開き、アクセラレーターを1xH100に設定し、AI v6を選択します。

Requirements

  • 初回の運行時にはHugging Face Hubからモデルの重さをダウンロードできるインターネットアクセス。
  • Hugging Faceの読み取りトークンは Databricksの秘密として保存されています。 認証ステップで hf_secret_scope と hf_secret_key ウィジェットをシークレットのスコープとキーに設定してください。
  • モデルの重さは TabFM Non-Commercial License v1.0の下でライセンスされています。
  • このノートブックには、Apache 2.0ライセンスの下でライセンスされているtabfm-1.0.0-pytorch(著作権はGoogle Research)のソースコードが含まれています。

TabFMはDatabricks AI環境バージョン6に事前インストールされているため、追加のインストールは不要です。

ライブラリをインポートする

PyTorch、scikit-learn のデータセットローダーと評価指標、および TabFMRegressor パッケージから tabfmTabFMClassifier / をインポートし、GPU が利用可能かどうかを確認します。

import numpy as np
import pandas as pd
import torch

from sklearn.datasets import load_breast_cancer, load_diabetes
from sklearn.metrics import accuracy_score, roc_auc_score, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split

from tabfm import TabFMClassifier, TabFMRegressor, tabfm_v1_0_0_pytorch as tabfm_v1_0_0

print(f"Torch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

ハグフェイスで認証

hf_secret_scopeとhf_secret_keyウィジェットを、Hugging Faceの読み取りトークンを保存するDatabricksの秘密範囲とキーに設定し、ログインしてHubクライアントがダウンロードを認証できるようにします。

from huggingface_hub import login

# Set these widgets to the Databricks secret scope and key that hold your Hugging Face read token.
dbutils.widgets.text("hf_secret_scope", "", "Hugging Face secret scope")
dbutils.widgets.text("hf_secret_key", "hf_token", "Hugging Face secret key")

hf_token = dbutils.secrets.get(
    scope=dbutils.widgets.get("hf_secret_scope"),
    key=dbutils.widgets.get("hf_secret_key"),
)
login(token=hf_token)

ゼロショット分類

乳がんデータセット(569サンプル、30数値特徴)でゼロショット分類を実施します。 mean radiusから導出されたカテゴリカルなradius_band列が追加され、入力テーブルは数値型とカテゴリ型を混在させます。 TabFMはトレーニング行をコンテキストとして渡し、テストラベルを単一のフォワードパスで予測します。

分類データセットの読み込みと分割

乳がんデータセットを読み込み、派生したカテゴリ radius_band 特徴を追加し、ターゲットに対して層化したうえで、訓練用 80% / テスト用 20% に分割します。

breast = load_breast_cancer(as_frame=True)
clf_df = breast.frame.copy()
clf_df["radius_band"] = pd.qcut(
    clf_df["mean radius"],
    q=4,
    labels=["small", "medium", "large", "xlarge"],
).astype(str)

X_clf = clf_df.drop(columns=["target"])
y_clf = clf_df["target"]

X_train_clf, X_test_clf, y_train_clf, y_test_clf = train_test_split(
    X_clf,
    y_clf,
    test_size=0.2,
    random_state=42,
    stratify=y_clf,
)

display(X_train_clf.head(5))
print({
    "train_rows": len(X_train_clf),
    "test_rows": len(X_test_clf),
    "feature_count": X_train_clf.shape[1],
})

当てはめて予測

分類モデルの重みを読み込み、すべての訓練行をコンテキスト内の例として渡し、テストセットのクラスラベルと確率を予測します。 報告の正確性とROC-AUC。

tabfm_clf_model = tabfm_v1_0_0.load(model_type="classification")
tabfm_clf = TabFMClassifier(model=tabfm_clf_model)
tabfm_clf.fit(X_train_clf, y_train_clf)
clf_pred_proba = np.asarray(tabfm_clf.predict_proba(X_test_clf))
clf_pred = (clf_pred_proba[:, 1] >= 0.5).astype(int)

clf_results = pd.DataFrame({
    "actual": y_test_clf.reset_index(drop=True),
    "predicted": clf_pred.astype(int),
    "positive_class_probability": clf_pred_proba[:, 1],
})

accuracy = accuracy_score(y_test_clf, clf_pred)
roc_auc = roc_auc_score(y_test_clf, clf_pred_proba[:, 1])

print({
    "accuracy": round(float(accuracy), 4),
    "roc_auc": round(float(roc_auc), 4),
})
display(clf_results.head(10))

ゼロショット回帰

糖尿病データセット(サンプル442、数値特徴10)でゼロショット回帰を実施します。 カテゴリカル bmi_band 列が追加されます。 TabFMは各検査サンプルに対して連続的な疾患進行スコアを予測します。

回帰データセットを読み込み、分割します

糖尿病データセットを読み込み、導出されたカテゴリカル bmi_band 特徴を追加し、80% トレーニング/20% テストに分割します。

diabetes = load_diabetes(as_frame=True)
reg_df = diabetes.frame.copy()
reg_df["bmi_band"] = pd.qcut(
    reg_df["bmi"],
    q=4,
    labels=["low", "mid_low", "mid_high", "high"],
).astype(str)

X_reg = reg_df.drop(columns=["target"])
y_reg = reg_df["target"]

X_train_reg, X_test_reg, y_train_reg, y_test_reg = train_test_split(
    X_reg,
    y_reg,
    test_size=0.2,
    random_state=42,
)

display(X_train_reg.head(5))
print({
    "train_rows": len(X_train_reg),
    "test_rows": len(X_test_reg),
    "feature_count": X_train_reg.shape[1],
})

当てはめて予測

回帰モデルの重みを読み込み、すべてのトレーニング行をコンテキスト内の例として渡し、テストセットの連続スコアを予測します。 RMSEとR²を報告してください。

tabfm_reg_model = tabfm_v1_0_0.load(model_type="regression")
tabfm_reg = TabFMRegressor(model=tabfm_reg_model)
tabfm_reg.fit(X_train_reg, y_train_reg)
reg_pred = np.asarray(tabfm_reg.predict(X_test_reg)).reshape(-1)

rmse = np.sqrt(mean_squared_error(y_test_reg, reg_pred))
r2 = r2_score(y_test_reg, reg_pred)

reg_results = pd.DataFrame({
    "actual": y_test_reg.reset_index(drop=True),
    "predicted": reg_pred,
})
reg_results["absolute_error"] = (reg_results["actual"] - reg_results["predicted"]).abs()

print({
    "rmse": round(float(rmse), 4),
    "r2": round(float(r2), 4),
})
display(reg_results.head(10))

次のステップ

このノートブックを別のデータセットに適応させるには、pandasのDataFrameを読み込み、ターゲットカラムを分離し、カテゴリカラムを文字列として残し、トレインセットとテストセットに分割し、 TabFMClassifier または TabFMRegressorで入れ替えます。 TabFMは訓練行をコンテキスト内の例として渡すため、メモリ使用量はトレーニングセットのサイズに応じてスケールするため、大きなテーブルの代表的なサンプルから始め、分類目標は10クラス以下に抑えましょう。

ノートブックの例

TabFM:ゼロショット表状基礎モデル

ノートブックを入手