TabFM : modèle de fondation tabulaire « zero-shot »

Ouvrir la version du bloc-notes de cette page

TabFM est un modèle de base de Google Research pour les données tabulaires. Il utilise l’apprentissage en contexte, où les lignes d’entraînement sont transmises en contexte et les prédictions sont faites en un seul passage avant, sans recours à l’ajustement fin, à la recherche d’hyperparamètres ou à l’entraînement spécifique à chaque ensemble de données. Il supporte la classification binaire et multi-classe (jusqu’à 10 classes) ainsi que la régression sur des tables à colonnes mixtes numériques et catégorielles.

Ce notebook effectue une classification zero-shot sur le jeu de données du cancer du sein et une régression zero-shot sur le jeu de données du diabète.

Note

Cet exemple nécessite l’environnement IA Databricks version 6 ou supérieure.

Connecter au calcul GPU sans serveur

Cliquez sur la liste déroulante Se connecter , puis sélectionnez GPU serverless. Ouvrez le panneau latéral Environnement , réglez Accélérateur à 1xH100, et sélectionnez AI v6.

Spécifications

  • Accès Internet pour télécharger les poids des modèles depuis le Hugging Face Hub lors de la première tentative.
  • Un jeton de lecture Hugging Face stocké sous forme de secret Databricks. Définissez les widgets hf_secret_scope et hf_secret_key dans l’étape d’authentification de façon à leur attribuer la portée et la clé de votre secret.
  • Les poids des modèles sont concédés sous la licence TabFM Non-Commercial License v1.0.
  • Ce notebook inclut le code source de tabfm-1.0.0-pytorch, Copyright Google Research, qui est sous licence Apache 2.0.

TabFM est préinstallé dans l’environnement Databricks AI version 6, donc aucune installation supplémentaire n’est nécessaire.

Importer des bibliothèques

Importez PyTorch, les chargeurs de jeux de données et les métriques de scikit-learn, ainsi que TabFMClassifier / TabFMRegressor depuis le package tabfm, puis vérifiez la disponibilité du 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)}")

S’authentifier avec Hugging Face

Définissez les widgets hf_secret_scope et hf_secret_key sur la portée de secret et la clé Databricks où est stocké votre jeton de lecture de Hugging Face, puis connectez-vous afin que le client du Hub puisse authentifier les téléchargements.

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)

Classification zero-shot

Effectuer une classification zero-shot sur l’ensemble de données sur le cancer du sein (569 échantillons, 30 caractéristiques numériques). Une colonne catégorielle radius_band dérivée de mean radius est ajoutée afin que la table d’entrée mélange les types numériques et catégoriels. TabFM passe les lignes de formation en contexte et prédit les étiquettes test lors d’une seule passe.

Chargez et divisez l’ensemble de données de classification

Chargez le jeu de données de cancer du sein, ajoutez une variable catégorielle dérivée radius_band, puis divisez-le en 80 % pour l’entraînement / 20 % pour le test, en stratifiant selon la variable cible.

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

Ajustement et prédiction

Chargez les pondérations du modèle de classification, utilisez toutes les lignes d’entraînement comme exemples en contexte et prédisez les étiquettes de classe ainsi que les probabilités pour l’ensemble de test. Indiquez l’exactitude et le 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))

Régression zero-shot

Effectuez une régression zero-shot sur le jeu de données diabetes (442 échantillons, 10 caractéristiques numériques). Une colonne catégorielle bmi_band est ajoutée. TabFM prédit un score de progression continue de la maladie pour chaque échantillon de test.

Chargez et séparez le jeu de données de régression

Chargez le jeu de données sur le diabète, ajoutez une variable catégorielle dérivée bmi_band, puis divisez-le en 80 % pour l’entraînement et 20 % pour le 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],
})

Ajustement et prédiction

Chargez les poids du modèle de régression, passez toutes les lignes d’entraînement comme exemples contextuels, et prédisez les scores continus pour l’ensemble de test. Indiquez le RMSE et le 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))

Étapes suivantes

Pour adapter ce notebook à un autre jeu de données, chargez un DataFrame Pandas, séparez la colonne cible, laissez les colonnes catégorielles sous forme de chaînes de caractères, divisez les données en jeu de formation et de test, puis remplacez par TabFMClassifier ou TabFMRegressor. Comme TabFM utilise les lignes de formation comme exemples en contexte, l’utilisation de la mémoire augmente avec la taille de l’ensemble de formation ; pour les grandes tables, commencez donc par un échantillon représentatif et limitez les cibles de classification à 10 classes ou moins.

Exemple de bloc-notes

TabFM : modèle de fondation tabulaire « zero-shot »

Obtenir un ordinateur portable

  • Dernière mise à jour :