import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_ranking as tfr

# Prep data
ds = tfds.load("mslr_web/10k_fold1", split="train")
ds = ds.map(lambda feature_map: {
    "_mask": tf.ones_like(feature_map["label"], dtype=tf.bool),
    **feature_map
})
ds = ds.shuffle(buffer_size=1000).padded_batch