TabFM: modelo tabular de base zero-shot

TabFM es un modelo de base de Google Research para datos tabulares. Utiliza aprendizaje en contexto, en el que las filas del conjunto de entrenamiento se proporcionan como contexto y las predicciones se realizan en una sola pasada hacia delante, sin necesidad de ajuste fino, búsqueda de hiperparámetros ni entrenamiento específico para cada conjunto de datos. Soporta clasificación binaria y multiclase (hasta 10 clases) y regresión en tablas con columnas numéricas y categóricas mixtas.

Este notebook ejecuta la clasificación sin entrenamiento previo en el conjunto de datos de cáncer de mama y la regresión sin entrenamiento previo en el conjunto de datos de diabetes.

Nota:

Este ejemplo requiere el entorno de IA de Databricks versión 6 o superior.

Conectar al cómputo de GPU sin servidor

Haga clic en la lista desplegable Conectar y seleccione GPU sin servidor. Abre el panel lateral de Entorno , pon el acelerador en 1xH100 y selecciona AI v6.

Requirements

  • Acceso a Internet para descargar los pesos de los modelos desde el Hugging Face Hub en la primera tirada.
  • Un token de lectura de Hugging Face almacenado como secreto de Databricks. Configura los widgets hf_secret_scope y hf_secret_key en el paso de autenticación con el alcance y la clave de tu secreto.
  • Los pesos de los modelos están licenciados bajo la TabFM Non-Commercial License v1.0.
  • Este cuaderno incluye código fuente de tabfm-1.0.0-pytorch, Copyright Google Research, que está licenciado bajo la licencia Apache 2.0.

TabFM está preinstalado en el entorno de IA de Databricks versión 6, por lo que no se requiere instalación adicional.

Importar bibliotecas

Importa PyTorch, los cargadores de conjuntos de datos y las métricas de scikit-learn, y TabFMClassifier / TabFMRegressor del paquete tabfm y, a continuación, verifica que la GPU esté disponible.

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

Autenticarse con Hugging Face

Configura los widgets hf_secret_scope y hf_secret_key con el ámbito de secretos y la clave de Databricks que almacenan tu token de lectura de Hugging Face, y luego inicia sesión para que el cliente del Hub pueda autenticar las descargas.

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)

Clasificación de cero disparos

Ejecuta la clasificación zero-shot en el conjunto de datos de cáncer de mama (569 muestras, 30 características numéricas). Se añade una columna categórica radius_band derivada de mean radius para que la tabla de entrada mezcle tipos numéricos y categóricos. TabFM pasa las filas de entrenamiento como contexto y predice las etiquetas de prueba en un solo pase hacia adelante.

Carga y divide el conjunto de datos de clasificación

Carga el conjunto de datos de cáncer de mama, añade una característica categórica derivada radius_band y divídelo en un 80 % para entrenamiento y un 20 % para prueba, con estratificación respecto de la variable objetivo.

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

Ajuste y predicción

Carga los pesos del modelo de clasificación, pasa todas las filas de entrenamiento como ejemplos en contexto y predice las etiquetas de clase y probabilidades para el conjunto de prueba. Informe la precisión y el 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))

Regresión zero-shot

Ejecuta una regresión zero-shot en el conjunto de datos sobre diabetes (442 muestras, 10 características numéricas). Se añade una columna categórica bmi_band . TabFM predice una puntuación continua de progresión de la enfermedad para cada muestra de prueba.

Carga y divide el conjunto de datos de regresión

Carga el conjunto de datos de diabetes, añade una característica categórica bmi_band derivada y divídela 80% de entrenamiento / 20% de prueba.

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

Ajuste y predicción

Carga los pesos del modelo de regresión, pasa todas las filas de entrenamiento como ejemplos en contexto y predice las puntuaciones continuas para el conjunto de pruebas. Indica el RMSE y el 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))

Pasos siguientes

Para adaptar este cuaderno a otro conjunto de datos, carga un DataFrame de pandas, separa la columna de destino, deja las columnas categóricas como cadenas, divide el conjunto en conjuntos de entrenamiento y de prueba, y sustituye TabFMClassifier o TabFMRegressor. Como TabFM pasa las filas de entrenamiento como ejemplos en contexto, el uso de memoria aumenta en función del tamaño del conjunto de entrenamiento, así que, para tablas grandes, empieza con una muestra representativa y limita los objetivos de clasificación a 10 clases o menos.

Cuaderno de ejemplo

TabFM: modelo tabular de base zero-shot

Obtener el portátil