Extract embeddings from tabular model

[renat-abbyazov]

Hi everyone!

Is there any way to extract tabular embeddings for categorical features from model in flash-tabular? Here is the code example for adult census dataset

import torch
import pandas as pd
import flash
from flash.core.data.utils import download_data
from flash.tabular import TabularClassificationData, TabularClassifier
from sklearn.model_selection import cross_validate, train_test_split
from sklearn.compose import make_column_selector as selector
from sklearn.metrics import classification_report, f1_score
import numpy as np

adult_census = pd.read_csv("https://raw.githubusercontent.com/INRIA/scikit-learn-mooc/main/datasets/adult-census.csv")
target_name = "class"
target = adult_census[target_name]
data = adult_census.drop(columns=[target_name, "education-num"])

X_train, X_test, y_train, y_test = train_test_split(data, target, random_state=0)

categorical_columns_selector = selector(dtype_include=object)
categorical_columns = categorical_columns_selector(data)

numerical_columns_selector = selector(dtype_exclude=object)
numerical_columns = numerical_columns_selector(data)

datamodule = TabularClassificationData.from_data_frame(
    categorical_fields=categorical_columns,
    numerical_fields=numerical_columns,
    target_fields=target_name,
    train_data_frame=pd.concat([X_train, y_train], axis=1),
    val_split=0.1,
    batch_size=64,
)

model = TabularClassifier.from_data(datamodule, backbone="fttransformer")
trainer = flash.Trainer(max_epochs=1)
trainer.fit(model, datamodule=datamodule)

for col in categorical_columns:
    print(col, data[col].nunique())
    
embeddings = trainer.model.backbone.extract_embedding()
for col, emb in zip(categorical_columns, embeddings):
    print(col, emb.weight.shape)

So the question is how to correctly map the categorical feature value to the corresponding index in embedding matrix and
what is the index of nan or unk category?

Thanks in advance.