[Python] How to create a DiCE counterfactual model with a Tensorflow backend?

asked this question the other day and just saw it was closed. I've made some changes and had another attempt at making a TF model:

import dice_ml
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.models import Sequential

data_url = "https://archive.ics.uci.edu/ml/machine-learning-databases/statlog/german/german.data"
columns = ["Status", "Duration", "CreditHistory", "Purpose", "CreditAmount", "Savings", "Employment", "InstallmentRate", "PersonalStatus", "OtherDebtors", "Residence", "Property", "Age", "OtherInstallment", "Housing", "ExistingCredits", "Job", "NumLiable", "Telephone", "ForeignWorker", "Target"]
dataset = pd.read_csv(data_url, delim_whitespace=True, names=columns)

target = dataset['Target']
train_dataset, test_dataset, y_train, y_test = train_test_split(dataset, target, test_size=0.2, random_state=42, stratify=target)
x_train = train_dataset.drop(columns=['Target'], axis=1)
x_test = test_dataset.drop(columns=['Target'], axis=1)

numerical = ['Duration', 'CreditAmount', 'InstallmentRate', 'Residence', 'Age', 'ExistingCredits', 'NumLiable']
categorical = x_train.columns.difference(numerical)

categorical_transformer = Pipeline(steps=[
('onehot', OneHotEncoder(handle_unknown='ignore'))
])

numerical_transformer = Pipeline(steps=[
('scaler', StandardScaler())
])

transformations = ColumnTransformer(
transformers=[
('num', numerical_transformer, numerical),
('cat', categorical_transformer, categorical)
]
)

x_train_transformed = transformations.fit_transform(x_train)
x_test_transformed = transformations.transform(x_test)

x_train_transformed_dense = pd.DataFrame(x_train_transformed, columns=transformations.get_feature_names_out())
x_test_transformed_dense = pd.DataFrame(x_test_transformed, columns=transformations.get_feature_names_out())

model = Sequential([
Input(shape=(x_train_transformed_dense.shape[1],)),
Dense(64, activation='relu'),
Dense(32, activation='relu'),
Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

model.fit(x_train_transformed_dense, y_train, epochs=10, batch_size=32, validation_split=0.2)

m = dice_ml.Model(model=model, backend="TF2")

d = dice_ml.Data(dataframe=dataset, continuous_features=numerical, outcome_name='Target')

exp = dice_ml.Dice(d, m, method="random")

instance_index = dataset[dataset["Target"] == 2].index[0]
query_instance = x_test.loc[[instance_index]]
query_instance_transformed = transformations.transform(query_instance)

query_instance_transformed_dense = pd.DataFrame(query_instance_transformed, columns=transformations.get_feature_names_out())

cf = exp.generate_counterfactuals(query_instance_transformed_dense, total_CFs=10, desired_range=None,
desired_class="opposite",
permitted_range=None, features_to_vary="all")


The issue Im having is this but im not sure how to go about fixing this:

Traceback (most recent call last):
File "/Users/user/PycharmProjects/CounterfactualExplanations/TensorFlow.py", line 66, in <module>
cf = exp.generate_counterfactuals(query_instance_transformed_dense, total_CFs=10, desired_range=None,
File "/Users/user/PycharmProjects/CounterfactualExplanations/.venv/lib/python3.9/site-packages/dice_ml/explainer_interfaces/explainer_base.py", line 185, in generate_counterfactuals
self.data_interface.set_continuous_feature_indexes(query_instance)
File "/Users/user/PycharmProjects/CounterfactualExplanations/.venv/lib/python3.9/site-packages/dice_ml/data_interfaces/base_data_interface.py", line 32, in set_continuous_feature_indexes
self.continuous_feature_indexes = [query_instance.columns.get_loc(name) for name in
File "/Users/user/PycharmProjects/CounterfactualExplanations/.venv/lib/python3.9/site-packages/dice_ml/data_interfaces/base_data_interface.py", line 32, in <listcomp>
self.continuous_feature_indexes = [query_instance.columns.get_loc(name) for name in
File "/Users/user/PycharmProjects/CounterfactualExplanations/.venv/lib/python3.9/site-packages/pandas/core/indexes/base.py", line 3804, in get_loc
raise KeyError(key) from err
KeyError: 'Duration'

Continue reading...