TabFM: zero-shot tabellarisch basismodel

TabFM is een Google Research foundation-model voor tabelgegevens. Het maakt gebruik van in-context learning, waarbij trainingsrijen als context worden doorgegeven en voorspellingen in één enkele voorwaartse doorgang worden gedaan, zonder dat finetuning, het zoeken naar hyperparameters of dataset-specifieke training vereist is. Het ondersteunt binaire en multi-klasse classificatie (tot 10 klassen) en regressie op tabellen met gemengde numerieke en categorische kolommen.

Dit notitieboekje past zero-shot classificatie toe op de borstkankerdataset en zero-shot regressie op de diabetesdataset.

Opmerking

Dit voorbeeld vereist de Databricks AI-omgeving versie 6 of hoger.

Verbinding maken met serverloze GPU-rekenkracht

Klik op de vervolgkeuzelijst Verbinding maken en selecteer Serverloze GPU. Open het Environment-zijpaneel , zet Accelerator op 1xH100 en selecteer AI v6.

Requirements

  • Internettoegang is vereist om de modelgewichten van de Hugging Face Hub te downloaden bij het eerste gebruik.
  • Een Hugging Face-read-token opgeslagen als een Databricks-geheim. Stel de hf_secret_scope- en hf_secret_key-widgets in de authenticatiestap in op de scope en sleutel van je secret.
  • De modelgewichten zijn gelicentieerd onder de TabFM Non-Commercial License v1.0.
  • Dit notitieboekje bevat broncode van tabfm-1.0.0-pytorch, Copyright Google Research, dat is gelicentieerd onder de Apache 2.0-licentie.

TabFM is vooraf geïnstalleerd in de Databricks AI-omgeving versie 6, dus er is geen extra installatie nodig.

Bibliotheken importeren

Importeer PyTorch, de scikit-learn-datasetladers en -metriek, en TabFMClassifier / TabFMRegressor uit het pakket tabfm, en controleer vervolgens of er een GPU beschikbaar is.

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

Authenticatie met Hugging Face

Stel de hf_secret_scope en hf_secret_key widgets in op de geheime scope en key van Databricks die je Hugging Face-leestoken opslaan, en log dan in zodat de Hub-client downloads kan authenticeren.

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 classificatie

Voer een zero-shot classificatie uit op de borstkankerdataset (569 steekproeven, 30 numerieke kenmerken). Een categorische radius_band kolom afgeleid van mean radius wordt toegevoegd zodat de invoertabel numerieke en categorische types mengt. TabFM gebruikt trainingsrijen als context en voorspelt testlabels in één voorwaartse stap.

Laad en splits de classificatiedataset

Laad de borstkankerdataset, voeg een afgeleide categorische radius_band eigenschap toe en splits deze in 80% training / 20% test met stratificatie op het doel.

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

Pas en voorspel

Laad de gewichten van het classificatiemodel, gebruik alle trainingsrijen als contextuele voorbeelden en voorspel klasselabels en waarschijnlijkheden voor de testset. Rapportagenauwkeurigheid en 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 regressie

Voer zero-shot regressie uit op de diabetesdataset (442 steekproeven, 10 numerieke kenmerken). Er wordt een categorische bmi_band kolom toegevoegd. TabFM voorspelt een continue ziekteprogressiescore voor elke testmonster.

Laad en splits de regressiedataset

Laad de diabetes-gegevensset, voeg een afgeleid categorisch bmi_band kenmerk toe en splits de gegevensset in 80% training en 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],
})

Pas en voorspel

Laad de gewichten van het regressiemodel, gebruik alle trainingsrijen als contextvoorbeelden en voorspel continue scores voor de testset. Meld RMSE en 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))

Volgende stappen 

Om deze notebook aan te passen aan een andere dataset, laad je een pandas-DataFrame, scheid je de doelkolom af, laat categorische kolommen als strings staan, splits je de gegevens in trainings- en testsets en vervang je dit door TabFMClassifier of TabFMRegressor. Omdat TabFM trainingsrijen als in-context voorbeelden doorgeeft, schaalt het geheugengebruik mee met de grootte van de trainingsset, dus begin met een representatieve steekproef voor grote tabellen en houd classificatiedoelen op 10 of minder klassen.

Voorbeeld van notebook

TabFM: zero-shot tabellarisch basismodel

Notebook krijgen