adapter that handles tf data dataset def initself dataset classweightnone distributionnone from keras utils moduleutils import tensorflow as tf if not isinstance dataset tf data dataset tf distribute distributeddataset raise valueerror expected argument dataset to be a tf data dataset freceived dataset if classweight is not none dataset dataset map makeclassweightmapfnclassweight prefetchtf data autotune if distribution is not none dataset distribution distributedatasetdataset self dataset dataset def getnumpyiteratorself from keras utils moduleutils import tensorflow as tf def converttonumpyx if isinstancex tf sparsetensor x tf sparse todensex return x numpy for batch in self dataset yield tree mapstructureconverttonumpy batch def gettfdatasetself return self dataset property def numbatchesself cardinality self dataset cardinality if callablecardinality dataset cardinality is normally expected to be a callable cardinality intself dataset cardinality else however in the case of distributeddataset it s a np int64 cardinality intcardinality return none for unknown and infiite cardinality datasets if cardinality 0 return none return cardinality property def batchsizeself firstelementspec tree flattenself dataset elementspec0 return firstelementspec shape0 property def haspartialbatchself return none property def partialbatchsizeself return none def makeclassweightmapfnclassweight from keras utils moduleutils import tensorflow as tf classweighttensor tf converttotensor classweight getintc 1 0 for c in rangemaxclassweight keys 1 def classweightsmapfndata special casing for rank 1 where we can guarantee sparse encoding adapter that handles tf data dataset iniitialize the tfdatasetadapter args dataset the input tf data dataset instance class_weight a map where the keys are integer class ids and values are the class weights e g 0 0 2 1 0 6 2 0 3 distribution a keras distribution distribution instance used to shard the input dataset into per worker process dataset instance dataset cardinality is normally expected to be a callable however in the case of distributeddataset it s a np int64 return none for unknown and infiite cardinality datasets applies class weighting to a dataset the dataset is assumed to be in format x y or x y sw where y must be a single tensor args class_weight a map where the keys are integer class ids and values are the class weights e g 0 0 2 1 0 6 2 0 3 returns a function that can be used with tf data dataset map to apply class weighting convert class_weight to sample_weight special casing for rank 1 where we can guarantee sparse encoding
import tree

from keras.trainers.data_adapters import data_adapter_utils
from keras.trainers.data_adapters.data_adapter import DataAdapter


class TFDatasetAdapter(DataAdapter):
    

    def __init__(self, dataset, class_weight=None, distribution=None):
        
        from keras.utils.module_utils import tensorflow as tf

        if not isinstance(
            dataset, (tf.data.Dataset, tf.distribute.DistributedDataset)
        ):
            raise ValueError(
                "Expected argument `dataset` to be a tf.data.Dataset. "
                f"Received: {dataset}"
            )
        if class_weight is not None:
            dataset = dataset.map(
                make_class_weight_map_fn(class_weight)
            ).prefetch(tf.data.AUTOTUNE)
        if distribution is not None:
            dataset = distribution.distribute_dataset(dataset)
        self._dataset = dataset

    def get_numpy_iterator(self):
        from keras.utils.module_utils import tensorflow as tf

        def convert_to_numpy(x):
            if isinstance(x, tf.SparseTensor):
                x = tf.sparse.to_dense(x)
            return x.numpy()

        for batch in self._dataset:
            yield tree.map_structure(convert_to_numpy, batch)

    def get_tf_dataset(self):
        return self._dataset

    @property
    def num_batches(self):
        cardinality = self._dataset.cardinality
        if callable(cardinality):
            
            cardinality = int(self._dataset.cardinality())
        else:
            
            cardinality = int(cardinality)
        
        if cardinality < 0:
            return None
        return cardinality

    @property
    def batch_size(self):
        first_element_spec = tree.flatten(self._dataset.element_spec)[0]
        return first_element_spec.shape[0]

    @property
    def has_partial_batch(self):
        return None

    @property
    def partial_batch_size(self):
        return None


def make_class_weight_map_fn(class_weight):
    
    from keras.utils.module_utils import tensorflow as tf

    class_weight_tensor = tf.convert_to_tensor(
        [
            class_weight.get(int(c), 1.0)
            for c in range(max(class_weight.keys()) + 1)
        ]
    )

    def class_weights_map_fn(*data):
        
        x, y, sw = data_adapter_utils.unpack_x_y_sample_weight(data)
        if sw is not None:
            raise ValueError(
                "You cannot `class_weight` and `sample_weight` "
                "at the same time."
            )
        if tree.is_nested(y):
            raise ValueError(
                "`class_weight` is only supported for Models with a single "
                "output."
            )

        if y.shape.rank >= 2:
            y_classes = tf.__internal__.smart_cond.smart_cond(
                tf.shape(y)[-1] > 1,
                lambda: tf.argmax(y, axis=-1),
                lambda: tf.cast(tf.round(tf.squeeze(y, axis=-1)), tf.int32),
            )
        else:
            
            y_classes = tf.cast(tf.round(y), tf.int32)

        cw = tf.gather(class_weight_tensor, y_classes)
        return x, y, cw

    return class_weights_map_fn