Apply fit_transform() to each of the training folds, and transform to each of the validation folds in TabularPredictor

[alberto-jj]

Hi,
I am currently running Autogluon TabularPredictor on a dataset previously split into train and test sets.
I re-scaled my features using fit_transform() using an algorithm to correct batch effects (reComBat) before running the TabularPredictor.fit().
However, I would like to apply my re-scaling with fit_transform() to each of the training folds of TabularPredictor and transform each of the validation folds (AND afterward, run the rest of AutoGluon preprocessing as default).

I have seen the class autogluon.features.generators.AbstractFeatureGenerator, but I'm not sure if it will help me to get my desired output (or if I will end up just applying my "fit_transform()" re-scaling without the remaining Autogluon default preprocessing).

Here is the code I have with my re-scaling method.

from sklearn.model_selection import train_test_split
from reComBat import reComBat
from autogluon.tabular import TabularPredictor

#Splitting train and test sets
X_train, X_test, y_train, y_test = train_test_split(df.iloc[:, :], df["group"], test_size=0.3, random_state=1997, stratify=df["group"])

#Define my re-scaling model (reComBat) to be used before TabularPredictor.fit()
model = reComBat(parametric=True,
                 model='elastic_net',
                 config={'alpha': 1e-5},
                 n_jobs=7,
                 verbose=True)

#Defining adjustment covariates in X_train to perform my desired re-scaling
covars = X_train.iloc[:, 1:5].copy()
covars.rename(columns={'age': 'age_numerical'}, inplace=True)
covars.center = covars.center.astype(str)
covars.group = covars.group.astype(str)
covars.gender = covars.gender.astype(str)


#Fitting the re-scaling (reComBat) model on the train data
model.fit(data=X_train.iloc[:,5:], batches=covars.center,X=covars.drop(['center'],axis=1))

#Re-scaling the train data using the fitted model
transformed_train_data = model.transform(data=X_train.iloc[:,5:], batches=covars.center,X=covars.drop(['center'],axis=1))

#Re-scaling the test data using the fitted model
transformed_test_data = model.transform(data=X_test.iloc[:,5:], batches=covars.center,X=covars.drop(['center'],axis=1))

#Fit tabular predictor
label = 'group'
time = 3600 * 4

save_path = 'automl/results/batch_corrected'  
predictor = TabularPredictor(label=label, path=save_path, eval_metric='balanced_accuracy',
                             sample_weight='auto_weight').fit(transformed_train_data, presets='best_quality', auto_stack=True, time_limit=time)  `

Unfortunately, I cannot share the data due to ethical constraints, but its all continuous numerical features without missing values to handle, nor categories to encode.

Thanks a lot in advance for your help.

edit:format