Selecting according to labels in a TensorFlow generator

I have a very large dataset (VoxCeleb) and each datum has a label (multi-class, can assume that the label is a number between 1 to 5000) and an audio recording. Since it is too large to load entirely in one time, my strategy is to use a generator (in TensorFlow 1, it means using fit_generator instead of fit in the training. I, however, use TensorFlow 2.8 and Keras).

Usually, a generator selects the batch randomly (by shuffelling the indices). I want the batch to be selected only semi-randomly in the following sense:

1. In Each batch there are n_s total samples.
2. The batch contains n_c distinct labels/classes (they are chosen randomly).
3. Each label/class in the batch has n_p samples (utterances).
4. n_s = n_c * n_p
5. An epoch is running on all the labels, such that every label is seen at least once.

This a general DataGenerator class I modified:

from os import path

import numpy as np
from keras.utils import Sequence
from keras.preprocessing.sequence import pad_sequences

from pre_processing import load_data # customize function


class DataGenerator(Sequence):
    """Generates data for Keras
    Sequence based data generator. Suitable for building data generator for training and prediction.
    """

    def __init__(self, list_IDs, labels, n_classes, input_path, target_path,
                 to_fit=True, batch_size=n_s, shuffle=True):
        """Initialization
        :param list_IDs: list of all 'label' ids to use in the generator
        :param to_fit: True to return X and y, False to return X only
        :param batch_size: batch size at each iteration
        :param shuffle: True to shuffle label indexes after every epoch
        """
        self.input_path = input_path
        self.target_path = target_path
        self.list_IDs = list_IDs
        self.labels = labels
        self.n_classes = n_classes
        self.to_fit = to_fit
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.on_epoch_end()

    def __len__(self):
        """Denotes the number of batches per epoch
        :return: number of batches per epoch
        """
        return int(np.floor(len(self.list_IDs) / self.batch_size))

    def __getitem__(self, index):
        """Generate one batch of data
        :param index: index of the batch
        :return: X and y when fitting. X only when predicting
        """

        # Generate indexes of the batch
        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]

        # Find list of IDs
        list_IDs_temp = [self.list_IDs[k] for k in indexes]

        # Generate data
        X = self._generate_X(list_IDs_temp)

        if self.to_fit:
            y = self._generate_y(list_IDs_temp)
            return [X], y
        else:
            return [X]

    def on_epoch_end(self):
        """
        Updates indexes after each epoch
        """
        self.indexes = np.arange(len(self.list_IDs))
        if self.shuffle:
            np.random.shuffle(self.indexes)

    def _generate_X(self, list_IDs_temp):
        """Generates data containing batch_size images
        :param list_IDs_temp: list of label ids to load
        :return: batch of images
        """
        # Initialization
        X = []

        # Generate data
        for i, ID in enumerate(list_IDs_temp):
            # Store sample
            temp = self._load_input(self.input_path, ID)
            X.append(temp)

        X = pad_sequences(X, value=0, padding='post')

        return X

    def _generate_y(self, list_IDs_temp):
        """Generates data containing batch_size masks
        :param list_IDs_temp: list of label ids to load
        :return: batch if masks
        """
        # TODO: modify
        y = []

        # Generate data
        for i, ID in enumerate(list_IDs_temp):
            # Store sample
            y.append(self._load_target(self.target_path, ID))

        # y = pad_sequences(y, value=0, padding='post')

        return y

        def _load_input(self, input_path, ID):
            feats = load_data(path.join(input_path, ID))
            return feats

        def _load_target(self, target_path, ID):
            return self.labels[ID]

Here I assume input_file is the directory where the audio files are saved, and ID is the file name. The function load_data is a customized function which reads the audio file and extracts some features (returns a Tensor).

How to write a such a generator that selects n_s samples in each batch according to the specifications above? Just shuffling indices and choosing randomly won't work here.

Added in Edit:

One approach is to sample randomly many samples n_big >> n_s and then filter out samples until we have at least n_c labels with n_p samples each. However, this is not guaranteed to work, so it may be computationally expansive (for each batch try many random subsets).

Added in Edit: I found something similar, but not what I need. This is a code in PyTorch (not TensorFlow), which creates a data generator for VoxCeleb: https://github.com/clovaai/voxceleb_trainer/blob/master/DatasetLoader.py