ds, feature_key="targets", batch_size=128)

num_shards = len(input_lengths) + int(use_prefix_lm_task)
if shard_ds:

ds_shards = [ds.shard(num_shards, i) for i in range(num_shards)]

else:

ds_shards = [ds for _ in range(num_shards)]

processed_ds = []
hyperparams = zip(input_lengths, hyperparams, range(num_shards))
for input_length, (noise_span_length, noise_density), i in hyperparams:

ds = ds_shards[i]
ds = t5.data.preprocessors.split_tokens(

ds,
feature_key="targets",
min_tokens_per_segment=None,
max_tokens_per_segment=input_length)

ds = t5.data.preprocessors.denoise(

ds,
output_features,
inputs_fn=t5.data.preprocessors.noise_span_to_unique_sentinel,
targets_fn=t5.data.preprocessors.nonnoise_span_to_unique_sentinel,
noise_density=noise_density,
noise_mask_fn=functools.partial(

t5.data.preprocessors.random_spans_noise_mask,
mean_noise_span_length=noise_span_length),

input_feature_key=input_feature_key)

if optional_task_prefixes:

ds = prepend_prompt(
UL2- Unifying Language Learning Paradigms