Is it possible to use a stratified sampling strategy for the train/test/validate splits that the automl library does? We are working in a context where we need to segregate certain groups from the training and test sets to see how our models generalize at level of these groups.

One hack that I considered was trying to to twist our classification model into a forecasting model by ordering our groups and then using the time_col parameter, but this is definitely the Wrong Way. Does anyone have any insight?

Thank you!