test

train_exp1_base_engvoc.sh → :q

@@ -1,8 +1,8 @@
-PROJECT_DIR=${HOME}"/models/long-t5x"
+PROJECT_DIR=${HOME}"/models/ul2-t5x"
 T5X_DIR="../../t5x"  # directory where the t5x is cloned.
-MODEL_DIR="gs://nb-t5x-us-central2/exp1-t5-base-engvoc"
+MODEL_DIR="gs://nb-t5x-us-central2/exp3-t5-base-ul2-engvoc"
 export PYTHONPATH=${PROJECT_DIR}
-MIXTURE_OR_TASK_NAME="ncc_scandinavian_span_corruption_stream_engvoc"
+MIXTURE_OR_TASK_NAME="scandinavian_ul2_engvoc"
 
 python3 ${T5X_DIR}/t5x/train.py \
   --gin_search_paths=${PROJECT_DIR} \

base.gin

@@ -3,14 +3,13 @@ include 'pretrain_cont.gin'
 #include 't5x/configs/runs/pretrain.gin'
 #iinclude 't5x/configs/runs/finetune.gin'
 
-
 # Register necessary SeqIO Tasks/Mixtures.
 import t5.data.mixtures
 import tasks
 
 MIXTURE_OR_TASK_NAME = %gin.REQUIRED
 TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
-TRAIN_STEPS = 1_700_000
+TRAIN_STEPS = 1_300_000
 DROPOUT_RATE = 0.0 # Changed from the default since T5-1.1 recomments this.
 #INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000"
 #INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/t5_1_1_base/checkpoint_1000000"

base_modified_scand.gin

@@ -0,0 +1,56 @@
+# T5.1.1 Base model.
+from __gin__ import dynamic_registration
+
+import seqio
+from t5x import adafactor
+from t5x import models
+from t5x.examples.t5 import network
+
+# ------------------- Loss HParam ----------------------------------------------
+Z_LOSS = 0.0001
+LABEL_SMOOTHING = 0.0
+# NOTE: When fine-tuning the public T5 checkpoints (trained in T5 MeshTF)
+# the loss normalizing factor should be set to pretraining batch_size *
+# target_token_length.
+LOSS_NORMALIZING_FACTOR = None
+# Dropout should be specified in the "run" files
+DROPOUT_RATE = %gin.REQUIRED
+
+# Vocabulary (shared by encoder and decoder)
+VOCABULARY = @seqio.SentencePieceVocabulary()
+seqio.SentencePieceVocabulary.sentencepiece_model_file = "gs://nb-t5/t5/vocabs/wikipedia/no-da-en-sv-nn-is_32000_unigram.sp.model"
+seqio.SentencePieceVocabulary.extra_ids = 100
+
+# ------------------- Optimizer ------------------------------------------------
+# `learning_rate` is set by `Trainer.learning_rate_fn`.
+OPTIMIZER = @adafactor.Adafactor()
+adafactor.Adafactor:
+  decay_rate = 0.8
+  step_offset = 0
+  logical_factor_rules = @adafactor.standard_logical_factor_rules()
+
+# ------------------- Model ----------------------------------------------------
+MODEL = @models.EncoderDecoderModel()
+models.EncoderDecoderModel:
+  module = @network.Transformer()
+  input_vocabulary = %VOCABULARY
+  output_vocabulary = %VOCABULARY
+  optimizer_def = %OPTIMIZER
+  z_loss = %Z_LOSS
+  label_smoothing = %LABEL_SMOOTHING
+  loss_normalizing_factor = %LOSS_NORMALIZING_FACTOR
+
+# ------------------- Network specification ------------------------------------
+network.Transformer.config = @network.T5Config()
+network.T5Config:
+  vocab_size = 32128  # vocab size rounded to a multiple of 128 for TPU efficiency
+  dtype = 'bfloat16'
+  emb_dim = 768
+  num_heads = 12
+  num_encoder_layers = 12
+  num_decoder_layers = 12
+  head_dim = 64
+  mlp_dim = 2048
+  mlp_activations = ('gelu', 'linear')
+  dropout_rate = %DROPOUT_RATE
+  logits_via_embedding = False

base_scandvoc.gin

@@ -0,0 +1,21 @@
+include 'base_modified_scand.gin'
+include 'pretrain_cont.gin'
+#include 't5x/configs/runs/pretrain.gin'
+#iinclude 't5x/configs/runs/finetune.gin'
+
+# Register necessary SeqIO Tasks/Mixtures.
+import t5.data.mixtures
+import tasks
+
+MIXTURE_OR_TASK_NAME = %gin.REQUIRED
+TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
+TRAIN_STEPS = 1_300_000
+DROPOUT_RATE = 0.0 # Changed from the default since T5-1.1 recomments this.
+#INITIAL_CHECKPOINT_PATH = "gs://nb-t5x-us-central2/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000"
+#INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/t5_1_1_base/checkpoint_1000000"
+INITIAL_CHECKPOINT_PATH = "gs://t5-data/pretrained_models/t5x/flan_t5_base/checkpoint_1184000"
+
+
+PjitPartitioner.num_partitions = 1
+
+

batch_nynorsk_NCC_base.sh

@@ -1,11 +0,0 @@
-PROJECT_DIR=${HOME}"/models/t5-nynorsk-oversetter"
-export PYTHONPATH=${PROJECT_DIR}
-INITIAL_CHECKPOINT_PATH=\"gs://nb-t5x-us-central2/norwegian_NCC_plus_English_t5x_base/checkpoint_1500000\"
-TRAIN_STEPS=1505000
-
-python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v1\" &&
-python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v2\" &&
-python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v3\" &&
-python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v4\" &&
-python3 ../../t5x/t5x/train.py --gin_search_paths="./" --gin.TRAIN_STEPS=${TRAIN_STEPS} --gin_file="finetune_translate_base.gin" --gin.INITIAL_CHECKPOINT_PATH=${INITIAL_CHECKPOINT_PATH}  --gin.MIXTURE_OR_TASK_NAME=\"translate\" --gin.MODEL_DIR=\"gs://nb-t5x-us-central2/finetuned/nynorsk_NCC_base_v5\"
-

pretrain_cont.gin

@@ -96,7 +96,7 @@ utils.RestoreCheckpointConfig:
   mode = 'specific'
   dtype = 'float32'
 utils.SaveCheckpointConfig:
-  period = 1000
+  period = 10000
   dtype = 'float32'
   keep = None  # keep all checkpoints
   save_dataset = False  # don't checkpoint dataset state

tasks.py

@@ -8,6 +8,18 @@ from t5.data import preprocessors
 from t5.evaluation import metrics
 from seqio import FunctionDataSource, utils
 
+from ul2_objective import ul2_objective
+# values from UL2 paper https://arxiv.org/pdf/2205.05131.pdf chapter 3.1.2 table 1
+R_DENOISER_SPAN_LENGTHS = [3.0, 8.0]
+X_DENOISER_SPAN_LENGTHS = [3.0, 8.0, 64.0, 64.0]
+R_DENOISER_CORRUPT_RATES = [0.15, 0.15]
+X_DENOISER_CORRUPT_RATES = [0.5, 0.5, 0.15, 0.5]
+
+R_DENOISER_TOKEN_PREFIX = '[NLU]'
+X_DENOISER_TOKEN_PREFIX = '[NLG]'
+S_DENOISER_TOKEN_PREFIX = '[S2S]'
+
+
 TaskRegistry = seqio.TaskRegistry
 scand_vocabulary=seqio.SentencePieceVocabulary('gs://nb-t5/t5/vocabs/wikipedia/no-da-en-sv-nn-is_32000_unigram.sp.model', extra_ids=100)
 eng_vocabulary=seqio.SentencePieceVocabulary('gs://t5-data/vocabs/cc_all.32000.100extra/sentencepiece.model', extra_ids=0)
@@ -60,7 +72,7 @@ dataset_name = 'NbAiLab/scandinavian'
 dataset_params = {"path": dataset_name, "use_auth_token": True, "streaming": True}
 dataset_shapes = None
 TaskRegistry.add(
-    "ncc_scandinavian_span_corruption_stream_engvoc",
+    "scandinavian_span_engvoc",
     source=seqio.FunctionDataSource(
         dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
         splits=("train", "validation"),
@@ -83,12 +95,12 @@ TaskRegistry.add(
 )
 
 
-# Final pretraining task used in Raffel et al., 2019 adaptated to NCC
+# Final pretraining task used in Tay et al., 2022 adaptated by @beats
 dataset_name = 'NbAiLab/scandinavian'
 dataset_params = {"path": dataset_name, "use_auth_token": True, "streaming": True}
 dataset_shapes = None
 TaskRegistry.add(
-    "ncc_scandinavian_span_corruption_stream_scandvoc",
+    "scandinavian_span_scandvoc",
     source=seqio.FunctionDataSource(
         dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
         splits=("train", "validation"),
@@ -111,4 +123,80 @@ TaskRegistry.add(
 )
 
 
+# Final pretraining task used in Tay et al., 2022 adaptated by @beats
+dataset_name = 'NbAiLab/scandinavian'
+dataset_params = {"path": dataset_name, "use_auth_token": True, "streaming": True}
+dataset_shapes = None
+TaskRegistry.add(
+    "scandinavian_ul2_engvoc",
+    source=seqio.FunctionDataSource(
+        dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
+        splits=("train", "validation"),
+        caching_permitted=False,
+        num_input_examples=dataset_shapes,
+        ),
+    preprocessors=[
+        functools.partial(
+            target_to_key, key_map={
+                "inputs": None,
+                "targets": None,
+            }, target_key="targets"),
+        seqio.preprocessors.tokenize,
+        functools.partial(
+            ul2_objective,
+            shard_ds=False,
+            use_prefix_lm_task=True,  # use S-denoising
+            rates=[0.4 / len(R_DENOISER_SPAN_LENGTHS)]*len(R_DENOISER_SPAN_LENGTHS) + [
+                0.4 / len(X_DENOISER_SPAN_LENGTHS)]*len(X_DENOISER_SPAN_LENGTHS) + [0.2],  # equal total 40% rate for both R- and X-denoisers + 20% for S-denoising (suggested at the paper chapter 4.5)
+            mean_noise_span_lengths=R_DENOISER_SPAN_LENGTHS + X_DENOISER_SPAN_LENGTHS,
+            noise_densities=R_DENOISER_CORRUPT_RATES + X_DENOISER_CORRUPT_RATES,
+            optional_task_prefixes=[R_DENOISER_TOKEN_PREFIX]*len(R_DENOISER_SPAN_LENGTHS) + [
+                X_DENOISER_TOKEN_PREFIX]*len(X_DENOISER_SPAN_LENGTHS) + [S_DENOISER_TOKEN_PREFIX],
+            reserved_for_packing=1,  # make room for task prefix token
+        ),
+        seqio.preprocessors.append_eos_after_trim,
+    ],
+    output_features={"targets": ENG_OUTPUT_FEATURES["targets"]},
+    metric_fns=[]
+)
+
+
+# Final pretraining task used in Raffel et al., 2019 adaptated to NCC
+dataset_name = 'NbAiLab/scandinavian'
+dataset_params = {"path": dataset_name, "use_auth_token": True, "streaming": True}
+dataset_shapes = None
+TaskRegistry.add(
+    "scandinavian_ul2_scandvoc",
+    source=seqio.FunctionDataSource(
+        dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
+        splits=("train", "validation"),
+        caching_permitted=False,
+        num_input_examples=dataset_shapes,
+        ),
+    preprocessors=[
+        functools.partial(
+            target_to_key, key_map={
+                "inputs": None,
+                "targets": None,
+            }, target_key="targets"),
+        seqio.preprocessors.tokenize,
+        functools.partial(
+            ul2_objective,
+            shard_ds=False,
+            use_prefix_lm_task=True,  # use S-denoising
+            rates=[0.4 / len(R_DENOISER_SPAN_LENGTHS)]*len(R_DENOISER_SPAN_LENGTHS) + [
+                0.4 / len(X_DENOISER_SPAN_LENGTHS)]*len(X_DENOISER_SPAN_LENGTHS) + [0.2],  # equal total 40% rate for both R- and X-denoisers + 20% for S-denoising (suggested at the paper chapter 4.5)
+            mean_noise_span_lengths=R_DENOISER_SPAN_LENGTHS + X_DENOISER_SPAN_LENGTHS,
+            noise_densities=R_DENOISER_CORRUPT_RATES + X_DENOISER_CORRUPT_RATES,
+            optional_task_prefixes=[R_DENOISER_TOKEN_PREFIX]*len(R_DENOISER_SPAN_LENGTHS) + [
+                X_DENOISER_TOKEN_PREFIX]*len(X_DENOISER_SPAN_LENGTHS) + [S_DENOISER_TOKEN_PREFIX],
+            reserved_for_packing=1,  # make room for task prefix token
+        ),
+        seqio.preprocessors.append_eos_after_trim,
+    ],
+    output_features={"targets": SCAND_OUTPUT_FEATURES["targets"]},
+    metric_fns=[]
+)
+
+
 

train_exp1_base_scandvoc.sh → train_exp1_base_ul2_engvoc.sh

@@ -1,8 +1,8 @@
-PROJECT_DIR=${HOME}"/models/long-t5x"
+PROJECT_DIR=${HOME}"/models/ul2-t5x"
 T5X_DIR="../../t5x"  # directory where the t5x is cloned.
-MODEL_DIR="gs://nb-t5x-us-central2/exp1-t5-base-scandvoc"
+MODEL_DIR="gs://nb-t5x-us-central2/exp3-t5-base-ul2-engvoc"
 export PYTHONPATH=${PROJECT_DIR}
-MIXTURE_OR_TASK_NAME="ncc_scandinavian_span_corruption_stream_scandvoc"
+MIXTURE_OR_TASK_NAME="scandinavian_ul2_engvoc"
 
 python3 ${T5X_DIR}/t5x/train.py \
   --gin_search_paths=${PROJECT_DIR} \

train_exp2_base_ul2_scandvoc.sh

@@ -0,0 +1,12 @@
+PROJECT_DIR=${HOME}"/models/ul2-t5x"
+T5X_DIR="../../t5x"  # directory where the t5x is cloned.
+MODEL_DIR="gs://nb-t5x-us-central2/exp2-t5-base-ul2-scandvoc"
+export PYTHONPATH=${PROJECT_DIR}
+MIXTURE_OR_TASK_NAME="scandinavian_ul2_scandvoc"
+
+python3 ${T5X_DIR}/t5x/train.py \
+  --gin_search_paths=${PROJECT_DIR} \
+  --gin_file="base_scandvoc.gin" \
+  --gin.MODEL_DIR="'${MODEL_DIR}'" \
+  --gin.MIXTURE_OR_TASK_NAME="'${MIXTURE_OR_TASK_NAME}'" \
+

train_exp3_base_span_engvoc.sh

@@ -0,0 +1,12 @@
+PROJECT_DIR=${HOME}"/models/ul2-t5x"
+T5X_DIR="../../t5x"  # directory where the t5x is cloned.
+MODEL_DIR="gs://nb-t5x-us-central2/exp3-t5-base-span-engvoc"
+export PYTHONPATH=${PROJECT_DIR}
+MIXTURE_OR_TASK_NAME="scandinavian_span_engvoc"
+
+python3 ${T5X_DIR}/t5x/train.py \
+  --gin_search_paths=${PROJECT_DIR} \
+  --gin_file="base.gin" \
+  --gin.MODEL_DIR="'${MODEL_DIR}'" \
+  --gin.MIXTURE_OR_TASK_NAME="'${MIXTURE_OR_TASK_NAME}'" \
+

train_exp4_base_span_scandvoc.sh

@@ -0,0 +1,12 @@
+PROJECT_DIR=${HOME}"/models/ul2-t5x"
+T5X_DIR="../../t5x"  # directory where the t5x is cloned.
+MODEL_DIR="gs://nb-t5x-us-central2/exp4-t5-base-span-scandvoc"
+export PYTHONPATH=${PROJECT_DIR}
+MIXTURE_OR_TASK_NAME="scandinavian_span_scandvoc"
+
+python3 ${T5X_DIR}/t5x/train.py \
+  --gin_search_paths=${PROJECT_DIR} \
+  --gin_file="base_scandvoc.gin" \
+  --gin.MODEL_DIR="'${MODEL_DIR}'" \
+  --gin.MIXTURE_OR_TASK_NAME="'${MIXTURE_OR_TASK_NAME}'" \
+

ul2_objective.py

@@ -0,0 +1,184 @@
+import functools
+import tensorflow as tf
+import seqio
+import t5.data
+from typing import Optional, Sequence
+
+# found this function and modified from https://github.com/GoogleCloudPlatform/t5x-on-vertex-ai/blob/main/tasks/custom_tasks.py#L78
+# UL2 paper appendix code missed this function
+def prepend_prompt(dataset: tf.data.Dataset,
+                   output_features: seqio.preprocessors.OutputFeaturesType,
+                   sequence_length: Optional[
+                       seqio.preprocessors.SequenceLengthType] = None,
+                   prompt_mode: str = "",
+                   key: str = "inputs",
+                   mode: str = "") -> tf.data.Dataset:
+    """Prepends a prompt at the beginning of an input sequence."""
+    del sequence_length
+    if prompt_mode and mode:
+        # output_features may not have inputs key
+        out_keys = list(output_features.keys())
+        prompt_tokens = output_features[out_keys[0]
+                                        ].vocabulary.encode_tf(prompt_mode)
+
+        def add_to_inputs(x):
+            x[key] = tf.concat([prompt_tokens, x[key]], axis=0)
+            return x
+
+        dataset = dataset.map(add_to_inputs)
+    return dataset
+
+# modified from t5.data.preprocessors because output_features may not have inputs key
+def split_tokens_to_inputs_length(dataset, sequence_length,
+                                  output_features, **kwargs):
+    max_tokens = sequence_length['inputs']
+    # output_features may not have inputs key
+    out_keys = list(output_features.keys())
+    if output_features[out_keys[0]].add_eos:
+        # Leave room to insert an EOS token.
+        max_tokens -= 1
+
+    return t5.data.preprocessors.split_tokens(dataset, max_tokens_per_segment=max_tokens, **kwargs)
+
+# modified from t5.data.preprocessors because output_features may not have inputs key
+def prefix_lm(dataset, sequence_length, output_features):
+    """Prefix language modeling objective used in Raffel et al. 2019."""
+    ds = dataset
+    ds = t5.data.preprocessors.select_random_chunk(ds, output_features=output_features,
+                                                   feature_key='targets', max_length=65536)
+    ds = split_tokens_to_inputs_length(ds, output_features=output_features,
+                                       sequence_length=sequence_length)
+    ds = t5.data.preprocessors.denoise(
+        ds,
+        output_features,
+        inputs_fn=t5.data.preprocessors.drop_nonnoise_tokens,
+        targets_fn=t5.data.preprocessors.drop_noise_tokens,
+        noise_density=0.5,
+        noise_mask_fn=t5.data.preprocessors.random_prefix_noise_mask,
+    )
+    return ds
+
+# copied from UL2 paper https://arxiv.org/pdf/2205.05131.pdf appendix chapter 9.2
+# note: modified to use the prefix_lm() from above instead of the default t5.data.preprocessors.prefix_lm() because output_features may not have inputs key
+def ul2_objective(dataset: tf.data.Dataset,
+                  sequence_length: seqio.preprocessors.SequenceLengthType,
+                  output_features: seqio.preprocessors.OutputFeaturesType,
+                  use_prefix_lm_task: bool = False,
+                  rates: Optional[Sequence[float]] = None,
+                  mean_noise_span_lengths: Sequence[float] = (3.0,),
+                  noise_densities: Sequence[float] = (0.15,),
+                  shard_ds: bool = True,
+                  optional_task_prefixes: Optional[Sequence[str]] = None,
+                  input_feature_key: str = "inputs",
+                  merge_examples_to_reduce_padding: bool = True,
+                  reserved_for_packing: bool = None,
+                  seed: int = 7) -> tf.data.Dataset:
+    """UL2-like pre-training objectives.
+    This preprocessor amounts to calling the 'span_corruption' function several
+    times with different values of 'noise_density' and 'mean_noise_span_length'.
+    We either shard or copy the dataset, then apply each function to each shard.
+    Add S-denoising (prefixLM) using use_prefix_lm_task.
+    Args:
+        dataset: A tf.data.Dataset with dictionaries containing the key 'input_feature_key'.
+        sequence_length: dict mapping of feature key to int length for that feature.
+        output_features: mapping of keys to features.
+        use_prefix_lm_task: <bool> If True, include PrefixLM in the task mix.
+        rates: <Optional<List<float>> List of rates per task. If None, tasks are sampled uniformly.
+        mean_noise_span_lengths: List of mean number of tokens per masked span per example.
+        noise_densities: List of what fraction of the tokens to mask.
+        shard_ds: <bool> If True, shard dataset per objective.
+        optional_task_prefixes: <Optional<list<str>> Strings to prepend for each corruption scheme. NOTE: If including prefixLM task, it must be the last prefix.
+        input_feature_key: which feature to use from the dataset as the input text tokens.
+        merge_examples_to_reduce_padding: if True, combines multiple input examples to reduce padding.
+        reserved_for_packing: if specified, reduces the desired inputs length by the specified amount to enable multiple examples to be packed together downstream.
+        seed: tf.int64 for controlling the random choice of spans.
+    Returns:
+        a dataset
+    """
+
+    if optional_task_prefixes:  # Ensure each task has a prefix.
+        num_tasks = len(noise_densities) + int(use_prefix_lm_task)
+        valid_number_of_prefixes = num_tasks == len(optional_task_prefixes)
+        if not valid_number_of_prefixes:
+            raise ValueError(
+                "Number of task prefixes must match number of tasks.")
+    inputs_length = sequence_length[input_feature_key]
+    input_lengths, targets_lengths = [], []
+    sequence_lengths = {x: y for x, y in sequence_length.items()}
+    if reserved_for_packing:
+        inputs_length -= reserved_for_packing
+        for x, y in sequence_length.items():
+            sequence_lengths[x] = y - reserved_for_packing
+    hyperparams = list(zip(mean_noise_span_lengths, noise_densities))
+    for mean_noise_span_length, noise_density in hyperparams:
+        input_length, targets_length = t5.data.preprocessors.random_spans_helper(
+            extra_tokens_per_span_inputs=1,
+            extra_tokens_per_span_targets=1,
+            inputs_length=inputs_length,
+            mean_noise_span_length=mean_noise_span_length,
+            noise_density=noise_density)
+        input_lengths.append(input_length)
+        targets_lengths.append(targets_length)
+
+        if sequence_length["targets"] < targets_length:
+            upper_bound = max(targets_lengths)
+            raise ValueError(
+                f'Expected max targets length for span corruption ({upper_bound}) is '
+                f'greater than configured targets length '
+                f"({sequence_length['targets']})")
+    ds = dataset
+    ds = t5.data.preprocessors.select_random_chunk(
+        ds,
+        output_features=output_features,
+        feature_key="targets",
+        max_length=65536)
+    if merge_examples_to_reduce_padding:
+        ds = t5.data.preprocessors.reduce_concat_tokens(
+            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(
+                ds,
+                output_features,
+                prompt_mode=optional_task_prefixes[i],
+                mode=optional_task_prefixes[i],
+                key=input_feature_key)
+        processed_ds.append(ds)
+    if use_prefix_lm_task:
+        ds = ds_shards[-1]
+        ds = prefix_lm(
+            ds, sequence_lengths, output_features)
+        if optional_task_prefixes:
+            ds = prepend_prompt(
+                ds,
+                output_features,
+                prompt_mode=optional_task_prefixes[-1],
+                mode=optional_task_prefixes[-1],
+                key=input_feature_key)
+        processed_ds.append(ds)
+    ds = tf.data.experimental.sample_from_datasets(processed_ds, rates, seed)
+    return ds
+