TabTransformerTF

Implementation of TabTransformer in TensorFlow and Keras.

Installation

The package can be install using

pip install tabtransformertf

TabTransformer Usage

from tabtransformertf.models.tabtransformer import TabTransformer
from tabtransformertf.utils.preprocessing import df_to_dataset, build_categorical_prep

# Category encoding layers
category_prep_layers = build_categorical_prep(train_data, CATEGORICAL_FEATURES)

# Preprocess dataset
train_dataset = df_to_dataset(train_data[FEATURES + [LABEL]], LABEL)

# Initialise model
tabtransformer = TabTransformer(
    numerical_features = NUMERIC_FEATURES,  # list with numerical features names
    categorical_features = CATEGORICAL_FEATURES,  # list with categorical features names
    categorical_lookup=category_prep_layers,  # dictionary with encoding layers
    embedding_dim=32,  
    out_dim=1, 
    out_activation='sigmoid',
    depth=4,
    heads=8,
    attn_dropout=0.2,
    ff_dropout=0.2,
    mlp_hidden_factors=[2, 4],
    use_column_embedding=True,  # flag to use fixed positional column embeddings
)

preds = tabtransformer.predict(train_dataset)

FTTransformer Usage

from tabtransformertf.models.fttransformer import FTTransformerEncoder, FTTransformer

# Encoder is specified separately in case we decide to pre-train the model
ft_linear_encoder = FTTransformerEncoder(
    numerical_features = NUMERIC_FEATURES, # list of numeric features
    categorical_features = CATEGORICAL_FEATURES, # list of categorical features
    numerical_data = X_train[NUMERIC_FEATURES].values, # train array of numerical features
    categorical_data = X_train[CATEGORICAL_FEATURES].values, # train array of categorical features
    y = None, # not needed for linear
    numerical_embedding_type='linear',
    embedding_dim=16,  # Embedding dimension (for categorical, numerical, and contextual)
    depth=3,  # Number of Transformer Blocks (layers)
    heads=6,  # Number of attention heads in a Transofrmer Block
    attn_dropout=0.2,  # Dropout for attention layers
    ff_dropout=0.2,  # Dropout in Dense layers
    use_column_embedding=True,  # Fixed column embeddings
    explainable=True  # Whether we want to output attention importances or not
)

# Pass the encoder to the model
ft_linear_transformer = FTTransformer(
    encoder=ft_linear_encoder,  # Encoder from above
    out_dim=1,  # Number of outputs in final layer
    out_activation='sigmoid',  # Activation function for final layer
)

preds = ft_linear_transformer.predict(train_dataset)

Credits

As a reference, I've combined this implementation with Keras guide.