TabFM 是 Google Research 基金會用於表格資料的模型。 它使用上下文學習，將訓練資料列作為上下文輸入，並在單次前向傳播中完成預測，無需微調、超參數搜尋或針對特定資料集進行訓練。 它支援二元分類與多類別分類（最多 10 個類別），以及對混合數值與類別欄位的表格進行迴歸分析。
本筆記本在乳癌資料集上執行零劑量分類，在糖尿病資料集上執行零劑量迴歸分析。
Note
此範例需要 Databricks AI 環境版本 6 或以上。
連接到無伺服器的 GPU 運算
點選「 連接 」下拉選單，選擇 無伺服器 GPU。 打開 環境 側面板，將 加速器 設為 1xH100，然後選擇 AI v6。
需求
- 首次執行時，需要網際網路連線，才能從 Hugging Face Hub 下載模型權重。
- 儲存於 Databricks 密鑰中的 Hugging Face 讀取權杖。 在驗證步驟中，將
hf_secret_key和hf_secret_scope小工具設為你的密鑰範圍和金鑰。 - 模型權重依據 TabFM 非商業授權 v1.0 授權。
- 本筆記本包含 tabfm-1.0.0-pytorch 的原始碼，版權屬於 Google Research，並依照 Apache 2.0 授權。
TabFM 已預先安裝於 Databricks AI 環境第 6 版中，因此無需額外安裝。
匯入程式庫
匯入 PyTorch、scikit-learn 的資料集載入器和指標，以及來自 tabfm 套件的 TabFMClassifier / TabFMRegressor，然後確認 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_key 和 hf_secret_scope 小工具設定為儲存 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 個數值特徵）上執行零樣本分類。 新增一個由 radius_band 衍生的類別 mean radius 欄位，使輸入表同時包含數值和類別資料類型。 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 個或更少。