TabFM: modelo de fundação tabular zero-shot

Open in Databricks

TabFM é um foundation model do Google Research para dados tabulares. Ele usa in-context learning, em que as linhas de treinamento são passadas como contexto e as previsões são feitas em uma única passada direta, sem necessidade de ajuste fino, busca de hiperparâmetros ou treinamento específico de dataset. Ele oferece suporte a classificação binária e multiclasse (até 10 classes) e a regressão em tabelas com colunas numéricas e categóricas mistas.

Este notebook executa classificação zero-shot no dataset de câncer de mama e regressão zero-shot no dataset de diabetes.

nota

Este exemplo requer a versão 6 ou acima do ambiente Databricks AI.

Conectar-se ao compute de GPU serverless

Clique no dropdown Connect e selecione Serverless GPU . Abra o painel lateral Environment , defina Accelerator como 1xA10 e selecione AI v6 .

Requirements

• Acesso à internet para download os pesos do modelo do Hugging Face Hub na primeira execução.
• Um token de leitura do Hugging Face armazenado como um segredo do Databricks. Defina os widgets hf_secret_scope e hf_secret_key na etapa de autenticação para o escopo e a key do seu segredo.
• Os pesos do modelo estão licenciados sob a Licença Não Comercial do TabFM v1.0.
• Este notebook inclui o código-fonte de tabfm-1.0.0-pytorch, Copyright Google Research, que é licenciado sob a licença Apache 2.0.

O TabFM vem pré-instalado no ambiente Databricks AI versão 6, portanto, nenhuma instalação adicional é necessária.

Importar bibliotecas

Importe o PyTorch, os carregadores de dataset e métricas do scikit-learn e TabFMClassifier / TabFMRegressor do pacote tabfm, depois verifique a disponibilidade de GPU.

Python

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

Authenticate with Hugging Face

Configure os widgets hf_secret_scope e hf_secret_key para o secret scope e a key do Databricks que armazenam seu token de leitura do Hugging Face, e faça o log in para que o cliente do Hub possa autenticar os downloads.

Python

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 classification

Execute a classificação zero-shot no dataset de câncer de mama (569 amostras, 30 recursos numéricos). Uma coluna de radius_band categórica derivada de mean radius é adicionada para que a tabela de entrada misture tipos numéricos e categóricos. O TabFM passa as linhas de treinamento como contexto e prevê os rótulos de teste em uma única passagem direta.

Carregar e dividir o dataset de classificação

Carregue o dataset de câncer de mama, adicione um radius_band recurso categórico derivado e divida-o em 80% para treinamento e 20% para teste com estratificação no rótulo.

Python

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

Ajustar e prever

Carregue os pesos do modelo de classificação, passe todas as linhas de treinamento como exemplos in-context e preveja os rótulos de classe e as probabilidades para o conjunto de teste. Relate a acurácia e o ROC-AUC.

Python

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 = np.asarray(tabfm_clf.predict(X_test_clf)).reshape(-1)

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

Regressão zero-shot

Execute a regressão zero-shot no diabetes dataset (442 amostras, 10 recursos numéricos). Uma coluna bmi_band categórica é adicionada. O TabFM prevê um escore contínuo de progressão de doenças para cada amostra de teste.

Load and split the regression dataset

Load the diabetes dataset, add a derived categorical bmi_band recurso, and split it 80% ensinar / 20% test.

Python

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

Ajustar e prever

Carregue os pesos do modelo de regressão, passe todas as linhas de treinamento como exemplos em contexto e preveja pontuações contínuas para o conjunto de teste. Relate o RMSE e o R².

Python

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

Next steps

Para adaptar este notebook a outro dataset, carregue um DataFrame do Pandas, separe a coluna de destino, deixe as colunas categóricas como strings, divida em conjuntos de treinamento e teste e insira TabFMClassifier ou TabFMRegressor via swap. Como o TabFM passa linhas de treinamento como exemplos em contexto, o uso de memória escala com o tamanho do conjunto de dados de treinamento; portanto, comece com uma amostra representativa para tabelas grandes e mantenha os alvos de classificação em 10 classes ou menos.

• Cartão do modelo TabFM
• Melhores práticas para compute de GPU serverless
• Solucionar problemas no compute de GPU serverless
• Classic machine learning on serverless GPU
• Notebooks de exemplo do AI Runtime

Notebook de exemplo

TabFM: zero-shot tabular foundation model

Abrir notebook em uma nova abaOpen in Databricks