TabFM: modello di fondazione tabulare zero-shot

Apri la versione notebook di questa pagina

TabFM è un modello di base di Google Research per dati tabulari. Utilizza l'apprendimento in contesto, dove le righe di addestramento vengono passate come contesto e le previsioni vengono fatte in un unico passaggio in avanti, senza bisogno di fine-tuning, ricerca di iperparametri o addestramento specifico per dataset. Supporta la classificazione binaria e multi-classe (fino a 10 classi) e la regressione su tabelle con colonne numeriche e categoriche miste.

Questo notebook esegue la classificazione zero-shot sul dataset del tumore al seno e la regressione zero-shot sul dataset del diabete.

Note

Questo esempio richiede l'ambiente di IA Databricks versione 6 o superiore.

Connessione alla computazione GPU senza server

Fare clic sull'elenco a discesa Connetti e selezionare GPU serverless. Apri il pannello laterale Ambiente , imposta Accelerator su 1xH100 e seleziona AI v6.

Requisiti

  • Accesso a Internet per scaricare i pesi dei modelli dall'Hugging Face Hub al primo tentativo.
  • Un token di lettura di Hugging Face archiviato come segreto Databricks. Imposta, nella fase di autenticazione, i widget hf_secret_scope e hf_secret_key sull'ambito e sulla chiave del tuo segreto.
  • I pesi dei modelli sono concessi in licenza secondo la TabFM Non-Commercial License v1.0.
  • Questo notebook include il codice sorgente di tabfm-1.0.0-pytorch, Copyright Google Research, che è concesso in licenza Apache 2.0.

TabFM è preinstallato nell'ambiente Databricks AI versione 6, quindi non è necessaria un'installazione aggiuntiva.

Importare librerie

Importa PyTorch, i loader di dataset e le metriche di scikit-learn e TabFMClassifier / TabFMRegressor dal pacchetto tabfm, quindi verifica la disponibilità della 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)}")

Autenticati con Hugging Face

Imposta i widget hf_secret_scope e hf_secret_key sull'ambito dei segreti e sulla chiave di Databricks che memorizzano il tuo token di accesso in lettura di Hugging Face, quindi accedi affinché il client Hub possa autenticare i download.

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)

Classificazione a zero esempi

Esegui una classificazione zero-shot sul dataset relativo al cancro al seno (569 campioni, 30 caratteristiche numeriche). Si aggiunge una colonna categorica radius_band derivata da mean radius così che la tabella di input mescoli tipi numerici e categorici. TabFM passa le righe di addestramento come contesto e prevede le etichette di test in un singolo passaggio in avanti.

Carica e suddividi il dataset di classificazione

Carica il dataset del tumore al seno, aggiungi una caratteristica categorica derivata radius_band e suddividilo in 80% per l'addestramento e 20% per il test, con stratificazione rispetto al target.

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

Addestra e predici

Caricare i pesi del modello di classificazione, fornire tutte le righe di addestramento come esempi in contesto e predire le etichette di classe e le probabilità per il set di test. Riporta l'accuratezza e il 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))

Regressione a zero colpi

Esegui una regressione zero-shot sul dataset diabetes (442 campioni, 10 caratteristiche numeriche). Viene aggiunta una colonna categoriale bmi_band . TabFM prevede un punteggio continuo di progressione della malattia per ogni campione di test.

Carica e suddivide il dataset di regressione

Carica il dataset del diabete, aggiungi una variabile categorica derivata bmi_band e suddividilo in 80% per l’addestramento e 20% per il 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],
})

Addestra e predici

Caricare i pesi del modello di regressione, superare tutte le righe di addestramento come esempi contestuali e prevedere i punteggi continui per il set di test. Riporta RMSE e 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))

Passaggi successivi

Per adattare questo notebook a un altro dataset, carica un DataFrame pandas, separa la colonna target, lascia le colonne categoriali come stringhe, suddividi nei set di addestramento e di test e sostituisci con TabFMClassifier o TabFMRegressor. Poiché TabFM fornisce le righe di addestramento come esempi nel contesto, l'utilizzo della memoria aumenta in funzione della dimensione del set di addestramento; quindi, per tabelle di grandi dimensioni, inizia con un campione rappresentativo e mantieni i target di classificazione a 10 classi o meno.

Notebook di esempio

TabFM: modello di fondazione tabulare zero-shot

Ottieni il notebook

  • Ultimo aggiornamento: