How to sample from datasets with equal weights in tf.data API when the number of classification is too big

Question

I used tf.compat.v1.data.experimental.sample_from_datasets to create datasets for estimator as follows:

The first trial

def add_training_neg_s