nvidia/Kumo-Tabular

Kumo Tabular

Kumo Tabular is NVIDIA's pretrained tabular foundation model for classification and regression.

Getting Started

Install structured-data-models for inference:

pip install structured-data-models

Use labeled examples as context to predict class probabilities for new data:

from sklearn.datasets import load_breast_cancer
import sdm

df = load_breast_cancer(as_frame=True).frame

table = sdm.TableTensor.from_pandas(
    df=df,
    stypes=sdm.infer_stypes(df, overrides={"target": "categorical"}),
    device="cuda",
)
model = sdm.models.KumoTabular(task="classification", device="cuda")

probs = model(
    x_context=table[:300].drop_columns("target"),
    y_context=table[:300, "target"],
    x_query=table[300:].drop_columns("target"),
    num_estimators=8,
)

print(probs)

To learn more, visit structured-data-models.

License

Kumo Tabular weights are released under OpenMDW 1.1.