TabFM: grundmodell för tabelldata utan tidigare träning

Öppna notebook-versionen av den här sidan

TabFM är en Google Research Foundation-modell för tabelldata. Den använder kontextinlärning, där rader i träningsdata ges som kontext och förutsägelser görs i ett enda framåtpass, utan att finjustering, sökning efter hyperparametrar eller datamängdsspecifik träning krävs. Den stödjer binär och multiklassklassifikation (upp till 10 klasser) samt regression på tabeller med blandade numeriska och kategoriska kolumner.

Denna notebook kör zero-shot-klassificering på bröstcancerdatasetet och zero-shot-regression på diabetesdatasetet.

Note

Detta exempel kräver Databricks AI-miljö version 6 eller högre.

Ansluta till serverlös GPU-beräkning

Klicka på listrutan Anslut och välj Serverlös GPU. Öppna sidopanelen Environment, ställ in Accelerator på 1xH100 och välj AI v6.

Kravspecifikation

  • Internetåtkomst för att ladda ner modellvikterna från Hugging Face Hub vid första körningen.
  • En Hugging Face-lästoken lagrad som en Databricks-hemlighet. Ställ in widgetarna hf_secret_key och hf_secret_scope i autentiseringssteget till omfånget och nyckeln för din hemlighet.
  • Modellvikterna är licensierade under TabFM Non-Commercial License v1.0.
  • Denna anteckningsbok innehåller källkod från tabfm-1.0.0-pytorch, Copyright Google Research, som är licensierad under Apache 2.0-licensen.

TabFM är förinstallerat i Databricks AI-miljö version 6, så ingen ytterligare installation krävs.

Importera bibliotek

Importera PyTorch, datasetladdare och mått för scikit-learn samt TabFMClassifier / TabFMRegressor från paketet tabfm, och kontrollera sedan att en GPU är tillgänglig.

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)}")

Autentisera med Hugging Face

Ställ in widgetarna hf_secret_key och hf_secret_scope till det Databricks-hemlighetsomfång och den nyckel som lagrar din Hugging Face-läs-token och logga sedan in så att Hub-klienten kan autentisera nedladdningar.

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)

Nollskottsklassificering

Kör zero-shot-klassificering på datasetet för bröstcancer (569 poster, 30 numeriska egenskaper). En kategorisk radius_band kolumn härledd från mean radius läggs till så att indatatabellen blandar numeriska och kategoriska typer. TabFM skickar träningsrader som kontext och förutsäger testetiketter i ett enda framåtpass.

Ladda och dela klassificeringsdatasetet

Ladda datauppsättningen för bröstcancer, lägg till en härledd kategorisk radius_band variabel och dela upp den i 80 % träning / 20 % test med stratifiering utifrån målvariabeln.

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],
})

Passa och förutsäga

Ladda klassificeringsmodellens vikter, skicka in alla träningsrader som kontextuella exempel och förutsäg klassetiketter och sannolikheter för testmängden. Rapportens noggrannhet och 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))

zero-shot-regression

Kör zero-shot-regression på diabetesdatasetet (442 prover, 10 numeriska egenskaper). En kategorisk bmi_band kolumn läggs till. TabFM förutsäger en kontinuerlig sjukdomsprogressionspoäng för varje testprov.

Ladda och dela regressionsdatasetet

Ladda diabetesdatamängden, lägg till en härledd kategorisk bmi_band-variabel och dela upp den i 80 % träning / 20 % test.

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],
})

Passa och förutsäga

Ladda regressionsmodellens vikter, mata in alla träningsrader som in-context-exempel och prediktera kontinuerliga värden för testmängden. Rapportera RMSE och 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))

Nästa steg

För att anpassa denna anteckningsbok till en annan datauppsättning, läs in en pandas-DataFrame, separera målkolumnen, låt kategoriska kolumner vara strängar, dela upp i tränings- och testuppsättningar och ersätt med TabFMClassifier eller TabFMRegressor. Eftersom TabFM skickar träningsrader som exempel i kontexten, skalar minnesanvändningen med träningsmängdsstorleken, så börja med ett representativt urval för stora tabeller och håll klassificeringsmålen på 10 eller färre klasser.

Exempelanteckningsbok

TabFM: grundmodell för tabelldata utan tidigare träning

Hämta anteckningsbok

  • Uppdaterades senast