How to implement custom batch sizes

I am trying to do a Listwise LTR NN, where I am predicting a rank based on the QID.

I currently have this custom datset but now need to make sure each batch is only relevant to the QID and doesn’t have other QID in same batch


Any pointers would be great please