Converting a tensorflow dataset of tensors to a dataset of dicts

I have a dataset (dataset_1) with a signature as seen below:

<TensorSliceDataset element_spec={'duration': TensorSpec(shape=(), dtype=tf.float32, name=None), 'protocol_type': TensorSpec(shape=(), ..., 'flag': TensorSpec(shape=(), dtype=tf.float32, name=None) }>

with 42 features. What i would like it to convert it to one with a dict where x is a tensor with 41 features and y is a tensor with the label as seen below.
It is also noted that 'class' is my target variable.

<PrefetchDataset element_spec=OrderedDict([('x', TensorSpec(shape=(None, 41), dtype=tf.float32, name=None)), ('y', TensorSpec(shape=(None, 1), dtype=tf.float32, name=None))])

What i have done so far is pass dataset_1 to the function :

def preprocess(dataset):

  def batch_format_fn(element):
   
    return collections.OrderedDict(
        
        x =  tf.reshape(tf.concat([element[i] for i in element.keys() if i!='class'],0), [-1, 41]) ,
        y = tf.reshape(element['class'], [-1, 1])
        
         )
    
  return dataset.repeat(5).shuffle(100, seed=1).batch(  # Repeat = number of rounds
      12) .map(batch_format_fn).prefetch(10)

as seen here :

dataset_2 = preprocess(dataset_1)

This seems to produce results of the appropriate shape but i noticed some wrong values.
What i mean is after inspecting a batch of dataset_2 i may see all values of an array in 'x' be 0 and there no such case in dataset_1.

Could the above process somehow produce something like that?