TabFM: базовая модель для табличных данных без дообучения

Откройте версию страницы для записной книжки

TabFM — это базовая модель Google Research для табличных данных. Он использует контекстное обучение, при котором строки обучающей выборки передаются в качестве контекста, а предсказания формируются за один прямой проход, без необходимости в тонкой настройке, подборе гиперпараметров или обучении под конкретный набор данных. Он поддерживает бинарную и многоклассовую классификацию (до 10 классов) и регрессию в таблицах с смешанными числовыми и категориальными столбцами.

В этом блокноте выполняются классификация без предварительного обучения на наборе данных по раку молочной железы и регрессия без предварительного обучения на наборе данных о диабете.

Note

Для этого примера требуется среда AI Databricks версии 6 или выше.

Подключитесь к беспсерверным вычислительным ресурсам на базе GPU

Щелкните раскрывающийся список "Подключиться" и выберите бессерверный GPU. Откройте боковую панель окружения , установите Accelerator на 1xH100 и выберите AI v6.

Requirements

  • Доступ в интернет для скачивания весов моделей с Hugging Face Hub при первом запуске.
  • Токен доступа на чтение Hugging Face, сохранённый в виде секрета Databricks. В шаге аутентификации задайте для виджетов hf_secret_scope и hf_secret_key область действия и ключ вашего секрета.
  • Веса моделей лицензированы по некоммерческой лицензии TabFM версии 1.0.
  • Этот блокнот содержит исходный код из tabfm-1.0.0-pytorch, © Google Research, распространяемый по лицензии Apache 2.0.

TabFM предварительно установлен в AI-среде Databricks версии 6, поэтому дополнительная установка не требуется.

Импорт библиотек

Импортируйте PyTorch, загрузчики наборов данных и метрики scikit-learn, а также TabFMClassifier / TabFMRegressor из пакета tabfm, затем проверьте доступность 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)}")

Аутентифицироваться с помощью Hugging Face

Установите для виджетов hf_secret_scope и hf_secret_key область секретов Databricks и ключ, в которых хранятся ваш токен чтения Hugging Face, затем выполните вход, чтобы клиент Hub мог проходить аутентификацию при загрузке.

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)

Классификация нулевого выстрела

Проведите классификацию без предварительного обучения на наборе данных о раке молочной железы (569 образцов, 30 числовых признаков). Добавляется категориальный radius_band столбец, полученный из mean radius, так что во входной таблице сочетаются числовые и категориальные типы данных. TabFM передаёт строки обучающей выборки в качестве контекста и предсказывает метки тестовой выборки за один проход вперёд.

Загрузите и разделите классификационный набор данных

Загрузите набор данных по раку молочной железы, добавьте производный категориальный признак radius_band и разделите его на обучающую (80 %) и тестовую (20 %) выборки со стратификацией по целевой переменной.

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

Подгонка и предсказание

Загрузите веса модели классификации, передайте все строки обучающего набора в качестве контекстных примеров и спрогнозируйте метки классов и вероятности для тестового набора. Укажите точность и 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))

Регрессия с нулевой выстрелом

Проведите нулевую регрессию по набору диабета (442 образца, 10 числовых признаков). Добавляется категориальный bmi_band столбец. TabFM предсказывает непрерывный показатель прогрессирования заболевания для каждого тестового образца.

Загрузите и разделите регрессионный набор данных

Загрузите набор данных по диабету, добавьте производный bmi_band категориальный признак и разделите его на 80% тренировка / 20% тест.

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

Подгонка и предсказание

Загрузите веса регрессионной модели, передайте все строки обучающего набора в качестве контекстных примеров и спрогнозируйте непрерывные значения для тестовой выборки. Укажите RMSE и 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))

Дальнейшие действия

Чтобы адаптировать этот блокнот к другому набору данных, загрузите DataFrame pandas, отделите целевой столбец, оставьте категориальные столбцы в виде строк, разделите данные на обучающую и тестовую выборки и подставьте TabFMClassifier или TabFMRegressor. Поскольку TabFM передаёт обучающие строки в виде контекстных примеров, использование памяти масштабируется с размером обучающего набора, поэтому начинайте с репрезентативной выборки для больших таблиц и держите цели классификации на уровне 10 или менее классов.

Пример записной книжки

TabFM: базовая модель для табличных данных без дообучения

Получите ноутбук

  • Last updated on