TabFM: Zero-Shot-Basismodell für tabellarische Daten

TabFM ist ein Google Research Foundation-Modell für tabellarische Daten. Es verwendet In-Context-Learning, bei dem Trainingszeilen als Kontext weitergegeben und Vorhersagen in einem einzigen Vorwärtsdurchlauf getroffen werden, ohne dass Feinabstimmung, Hyperparametersuche oder datensatzspezifisches Training erforderlich sind. Es unterstützt binäre und mehrklassige Klassifikationen (bis zu 10 Klassen) sowie Regressionen auf Tabellen mit gemischten numerischen und kategorialen Spalten.

Dieses Notizbuch führt eine Zero-Shot-Klassifikation im Brustkrebsdatensatz und eine Zero-Shot-Regression im Diabetes-Datensatz durch.

Note

Dieses Beispiel benötigt die Databricks-KI-Umgebung Version 6 oder höher.

Verbindung zu Serverless GPU-Compute herstellen

Klicken Sie auf die Dropdownliste "Verbinden" , und wählen Sie "Serverlose GPU" aus. Öffne das Environment-Seitenpanel , stelle Accelerator auf 1xH100 und wähle AI v6 aus.

Anforderungen

  • Internetzugang, um die Modellgewichte vom Hugging Face Hub beim ersten Durchlauf herunterzuladen.
  • Ein Hugging Face-Lesetoken, der als Databricks-Geheimnis gespeichert wird. Stellen Sie im Authentifizierungsschritt die Widgets hf_secret_scope und hf_secret_key auf den Bereich und den Schlüssel Ihres Geheimnisses ein.
  • Die Modellgewichte sind unter der TabFM Non-Commercial License v1.0 lizenziert.
  • Dieses Notizbuch enthält Quellcode von tabfm-1.0.0-pytorch, Copyright Google Research, das unter der Apache 2.0-Lizenz lizenziert ist.

TabFM ist in der Databricks KI-Umgebung Version 6 vorinstalliert, sodass keine zusätzliche Installation erforderlich ist.

Importieren von Bibliotheken

Importiere PyTorch, die scikit-learn-Datenset-Loader und -Metriken sowie TabFMClassifier / TabFMRegressor aus dem tabfm Paket, um dann die GPU-Verfügbarkeit zu überprüfen.

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

Authentifizieren Sie sich mit Hugging Face

Lege für die Widgets hf_secret_scope und hf_secret_key den Databricks-Geheimnisbereich und den Schlüssel fest, in denen dein Hugging Face-Lese-Token gespeichert ist, und melde dich dann an, damit der Hub-Client Downloads authentifizieren kann.

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)

Zero-Shot-Klassifikation

Führen Sie eine Zero-Shot-Klassifikation im Brustkrebsdatensatz durch (569 Stichproben, 30 numerische Merkmale). Eine kategorische Spalte radius_band , abgeleitet von mean radius , wird hinzugefügt, sodass die Eingabetabelle numerische und kategorische Typen mischt. TabFM übergibt Trainingszeilen als Kontext und sagt Testbezeichnungen in einem einzigen Vorwärtsdurchgang voraus.

Laden und teilen Sie den Klassifikationsdatensatz

Laden Sie den Brustkrebs-Datensatz, fügen Sie ein abgeleitetes kategoriales radius_band Merkmal hinzu und teilen Sie ihn in 80 % Training und 20 % Test mit Stratifizierung anhand der Zielvariable.

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

Anpassen und vorhersagen

Laden Sie die Gewichte des Klassifikationsmodells, übergeben Sie alle Trainingszeilen als In-Context-Beispiele, und sagen Sie Klassenlabels und Wahrscheinlichkeiten für den Testdatensatz voraus. Berichtsgenauigkeit und 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

Führen Sie eine Zero-Shot-Regression auf dem Diabetes-Datensatz durch (442 Stichproben, 10 numerische Merkmale). Eine kategoriale Spalte bmi_band wird hinzugefügt. TabFM sagt für jede Testprobe einen kontinuierlichen Krankheitsprogressionswert voraus.

Laden und teilen Sie den Regressionsdatensatz auf

Lade den Diabetes-Datensatz, füge ein abgeleitetes kategorisches bmi_band-Merkmal hinzu und teile ihn in 80 % Training und 20 % Test auf.

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

Anpassen und vorhersagen

Laden Sie die Gewichte des Regressionsmodells, verwenden Sie alle Trainingszeilen als In-Context-Beispiele und prognostizieren Sie kontinuierliche Scores für den Testsatz. Melden Sie RMSE und 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ächste Schritte

Um dieses Notizbuch an einen anderen Datensatz anzupassen, laden Sie einen pandas-DataFrame, trennen Sie die Zielspalte ab, belassen Sie kategorische Spalten als Strings, teilen Sie die Daten in Trainings- und Testsätze auf und setzen Sie TabFMClassifier oder TabFMRegressor ein. Da TabFM Trainingszeilen als In-Context-Beispiele übergibt, skaliert der Speicherverbrauch mit der Größe des Trainingssatzes, daher beginnen Sie mit einer repräsentativen Stichprobe für große Tabellen und halten Sie die Klassifikationsziele bei 10 oder weniger Klassen.

Beispiel-Notebook

TabFM: Zero-Shot-Basismodell für tabellarische Daten

Notebook abrufen