Upload 10 files

config.gin

@@ -0,0 +1,145 @@
+from __gin__ import dynamic_registration
+import __main__ as train_script
+import seqio
+from t5x import adafactor
+from t5x.examples.t5 import network
+from t5x import gin_utils
+from t5x import models
+from t5x import partitioning
+from t5x import trainer
+from t5x import utils
+import ul2_tasks
+
+# Macros:
+# ==============================================================================
+BATCH_SIZE = 256
+DROPOUT_RATE = 0.0
+LABEL_SMOOTHING = 0.0
+LOSS_NORMALIZING_FACTOR = None
+MIXTURE_OR_TASK_MODULE = None
+MIXTURE_OR_TASK_NAME = 'pretrain_biological_ul2'
+MODEL = @models.EncoderDecoderModel()
+MODEL_DIR = '/models/bioul2-mini-nl8'
+OPTIMIZER = @adafactor.Adafactor()
+RANDOM_SEED = None
+SHUFFLE_TRAIN_EXAMPLES = True
+TASK_FEATURE_LENGTHS = {'inputs': 512, 'targets': 512}
+TRAIN_STEPS = 500000
+USE_CACHED_TASKS = False
+USE_HARDWARE_RNG = False
+VOCABULARY = @seqio.SentencePieceVocabulary()
+Z_LOSS = 0.0001
+
+# Parameters for adafactor.Adafactor:
+# ==============================================================================
+adafactor.Adafactor.decay_rate = 0.8
+adafactor.Adafactor.logical_factor_rules = \
+    @adafactor.standard_logical_factor_rules()
+adafactor.Adafactor.step_offset = 0
+
+# Parameters for utils.CheckpointConfig:
+# ==============================================================================
+utils.CheckpointConfig.restore = @utils.RestoreCheckpointConfig()
+utils.CheckpointConfig.save = @utils.SaveCheckpointConfig()
+
+# Parameters for utils.create_learning_rate_scheduler:
+# ==============================================================================
+utils.create_learning_rate_scheduler.base_learning_rate = 1.0
+utils.create_learning_rate_scheduler.factors = 'constant * rsqrt_decay'
+utils.create_learning_rate_scheduler.warmup_steps = 10000
+
+# Parameters for train/utils.DatasetConfig:
+# ==============================================================================
+train/utils.DatasetConfig.batch_size = %BATCH_SIZE
+train/utils.DatasetConfig.mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+train/utils.DatasetConfig.module = %MIXTURE_OR_TASK_MODULE
+train/utils.DatasetConfig.pack = True
+train/utils.DatasetConfig.seed = None
+train/utils.DatasetConfig.shuffle = %SHUFFLE_TRAIN_EXAMPLES
+train/utils.DatasetConfig.split = 'train'
+train/utils.DatasetConfig.task_feature_lengths = %TASK_FEATURE_LENGTHS
+train/utils.DatasetConfig.use_cached = %USE_CACHED_TASKS
+
+# Parameters for train_eval/utils.DatasetConfig:
+# ==============================================================================
+train_eval/utils.DatasetConfig.batch_size = %BATCH_SIZE
+train_eval/utils.DatasetConfig.mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+train_eval/utils.DatasetConfig.module = %MIXTURE_OR_TASK_MODULE
+train_eval/utils.DatasetConfig.pack = True
+train_eval/utils.DatasetConfig.seed = 42
+train_eval/utils.DatasetConfig.shuffle = False
+train_eval/utils.DatasetConfig.split = 'validation'
+train_eval/utils.DatasetConfig.task_feature_lengths = %TASK_FEATURE_LENGTHS
+train_eval/utils.DatasetConfig.use_cached = %USE_CACHED_TASKS
+
+# Parameters for models.EncoderDecoderModel:
+# ==============================================================================
+models.EncoderDecoderModel.input_vocabulary = %VOCABULARY
+models.EncoderDecoderModel.label_smoothing = %LABEL_SMOOTHING
+models.EncoderDecoderModel.loss_normalizing_factor = %LOSS_NORMALIZING_FACTOR
+models.EncoderDecoderModel.module = @network.Transformer()
+models.EncoderDecoderModel.optimizer_def = %OPTIMIZER
+models.EncoderDecoderModel.output_vocabulary = %VOCABULARY
+models.EncoderDecoderModel.z_loss = %Z_LOSS
+
+# Parameters for partitioning.PjitPartitioner:
+# ==============================================================================
+partitioning.PjitPartitioner.logical_axis_rules = \
+    @partitioning.standard_logical_axis_rules()
+partitioning.PjitPartitioner.model_parallel_submesh = None
+partitioning.PjitPartitioner.num_partitions = 1
+
+# Parameters for utils.RestoreCheckpointConfig:
+# ==============================================================================
+utils.RestoreCheckpointConfig.path = []
+
+# Parameters for utils.SaveCheckpointConfig:
+# ==============================================================================
+utils.SaveCheckpointConfig.dtype = 'float32'
+utils.SaveCheckpointConfig.keep = 5
+utils.SaveCheckpointConfig.period = 10000
+utils.SaveCheckpointConfig.save_dataset = False
+
+# Parameters for seqio.SentencePieceVocabulary:
+# ==============================================================================
+seqio.SentencePieceVocabulary.sentencepiece_model_file = 'spiece.model'
+
+# Parameters for network.T5Config:
+# ==============================================================================
+network.T5Config.dropout_rate = %DROPOUT_RATE
+network.T5Config.dtype = 'bfloat16'
+network.T5Config.emb_dim = 384
+network.T5Config.head_dim = 64
+network.T5Config.logits_via_embedding = False
+network.T5Config.mlp_activations = ('gelu', 'linear')
+network.T5Config.mlp_dim = 1536
+network.T5Config.num_decoder_layers = 8
+network.T5Config.num_encoder_layers = 8
+network.T5Config.num_heads = 8
+network.T5Config.vocab_size = 32128
+
+# Parameters for train_script.train:
+# ==============================================================================
+train_script.train.checkpoint_cfg = @utils.CheckpointConfig()
+train_script.train.eval_period = 10000
+train_script.train.eval_steps = 20
+train_script.train.infer_eval_dataset_cfg = None
+train_script.train.model = %MODEL
+train_script.train.model_dir = %MODEL_DIR
+train_script.train.partitioner = @partitioning.PjitPartitioner()
+train_script.train.random_seed = %RANDOM_SEED
+train_script.train.summarize_config_fn = @gin_utils.summarize_gin_config
+train_script.train.total_steps = %TRAIN_STEPS
+train_script.train.train_dataset_cfg = @train/utils.DatasetConfig()
+train_script.train.train_eval_dataset_cfg = @train_eval/utils.DatasetConfig()
+train_script.train.trainer_cls = @trainer.Trainer
+train_script.train.use_hardware_rng = %USE_HARDWARE_RNG
+
+# Parameters for trainer.Trainer:
+# ==============================================================================
+trainer.Trainer.learning_rate_fn = @utils.create_learning_rate_scheduler()
+trainer.Trainer.num_microbatches = None
+
+# Parameters for network.Transformer:
+# ==============================================================================
+network.Transformer.config = @network.T5Config()

config.json

@@ -0,0 +1,32 @@
+{
+    "_name_or_path": "./",
+    "architectures": [
+      "T5ForConditionalGeneration"
+    ],
+    "d_ff": 1536,
+    "d_kv": 64,
+    "d_model": 384,
+    "decoder_start_token_id": 0,
+    "dense_act_fn": "gelu_new",
+    "dropout_rate": 0.1,
+    "eos_token_id": 1,
+    "feed_forward_proj": "gated-gelu",
+    "initializer_factor": 1.0,
+    "is_encoder_decoder": true,
+    "is_gated_act": true,
+    "layer_norm_epsilon": 1e-06,
+    "model_type": "t5",
+    "n_positions": 512,
+    "num_decoder_layers": 8,
+    "num_heads": 8,
+    "num_layers": 8,
+    "output_past": true,
+    "pad_token_id": 0,
+    "relative_attention_max_distance": 128,
+    "relative_attention_num_buckets": 32,
+    "tie_word_embeddings": false,
+    "torch_dtype": "float32",
+    "transformers_version": "4.22.1",
+    "use_cache": true,
+    "vocab_size": 32128
+  }

convert_t5x_checkpoint_to_flax.py

@@ -0,0 +1,162 @@
+# https://gist.github.com/stefan-it/30e4998ef159f33696e377a46f699d9f
+
+import argparse
+
+from t5x import checkpoints
+from transformers import T5Config, FlaxT5ForConditionalGeneration, AutoModelForSeq2SeqLM
+import torch
+
+
+def convert_t5x_checkpoint_to_flax(t5x_checkpoint_path, config_name, flax_dump_folder_path):
+    config = T5Config.from_pretrained(config_name)
+    flax_model = FlaxT5ForConditionalGeneration(config=config)
+    t5x_model = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path)
+
+    split_mlp_wi = "wi_0" in t5x_model["target"]["encoder"]["layers_0"]["mlp"]
+
+    # Encoder
+    for layer_index in range(config.num_layers):
+        layer_name = f"layers_{str(layer_index)}"
+        
+        # Self-Attention
+        t5x_attention_key = t5x_model["target"]["encoder"][layer_name]["attention"]["key"]["kernel"]
+        t5x_attention_out = t5x_model["target"]["encoder"][layer_name]["attention"]["out"]["kernel"]
+        t5x_attention_query = t5x_model["target"]["encoder"][layer_name]["attention"]["query"]["kernel"]
+        t5x_attention_value = t5x_model["target"]["encoder"][layer_name]["attention"]["value"]["kernel"]
+
+        ## Layer Normalization
+        t5x_attention_layer_norm = t5x_model["target"]["encoder"][layer_name]["pre_attention_layer_norm"]["scale"]
+        
+        if split_mlp_wi:
+            t5x_mlp_wi_0 = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi_0"]["kernel"]
+            t5x_mlp_wi_1 = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi_1"]["kernel"]
+        else:
+            t5x_mlp_wi = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi"]["kernel"]
+        
+        t5x_mlp_wo = t5x_model["target"]["encoder"][layer_name]["mlp"]["wo"]["kernel"]
+        
+        ## Layer Normalization
+        t5x_mlp_layer_norm = t5x_model["target"]["encoder"][layer_name]["pre_mlp_layer_norm"]["scale"]
+        
+        # Assigning
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value
+        
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["layer_norm"]["weight"] = t5x_attention_layer_norm
+        
+        if split_mlp_wi:
+            flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wi_0"]["kernel"] = t5x_mlp_wi_0
+            flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wi_1"]["kernel"] = t5x_mlp_wi_1
+        else:
+            flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wi"]["kernel"] = t5x_mlp_wi
+
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wo"]["kernel"] = t5x_mlp_wo
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["layer_norm"]["weight"] = t5x_mlp_layer_norm
+
+    # Only for layer 0:
+    t5x_encoder_rel_embedding = t5x_model["target"]["encoder"]["relpos_bias"]["rel_embedding"].T
+    flax_model.params["encoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_encoder_rel_embedding
+
+    # Assigning
+    t5x_encoder_norm = t5x_model["target"]["encoder"]["encoder_norm"]["scale"]
+    flax_model.params["encoder"]["final_layer_norm"]["weight"] = t5x_encoder_norm
+
+    # Decoder
+    for layer_index in range(config.num_decoder_layers):
+        layer_name = f"layers_{str(layer_index)}"
+
+        # Self-Attention
+        t5x_attention_key = t5x_model["target"]["decoder"][layer_name]["self_attention"]["key"]["kernel"]
+        t5x_attention_out = t5x_model["target"]["decoder"][layer_name]["self_attention"]["out"]["kernel"]
+        t5x_attention_query = t5x_model["target"]["decoder"][layer_name]["self_attention"]["query"]["kernel"]
+        t5x_attention_value = t5x_model["target"]["decoder"][layer_name]["self_attention"]["value"]["kernel"]
+
+        ## Layer Normalization
+        t5x_pre_attention_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_self_attention_layer_norm"]["scale"]
+
+        # Encoder-Decoder-Attention
+        t5x_enc_dec_attention_key = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["key"]["kernel"]
+        t5x_enc_dec_attention_out = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["out"]["kernel"]
+        t5x_enc_dec_attention_query = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["query"]["kernel"]
+        t5x_enc_dec_attention_value = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["value"]["kernel"]
+
+        ## Layer Normalization
+        t5x_cross_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_cross_attention_layer_norm"]["scale"]
+
+        # MLP
+        if split_mlp_wi:
+            t5x_mlp_wi_0 = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi_0"]["kernel"]
+            t5x_mlp_wi_1 = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi_1"]["kernel"]
+        else:
+            t5x_mlp_wi = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi"]["kernel"]
+
+        t5x_mlp_wo = t5x_model["target"]["decoder"][layer_name]["mlp"]["wo"]["kernel"]
+
+        ## Layer Normalization
+        tx5_mlp_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_mlp_layer_norm"]["scale"]
+        
+        # Assigning
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value
+
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["layer_norm"]["weight"] = t5x_pre_attention_layer_norm
+        
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["k"]["kernel"] = t5x_enc_dec_attention_key
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["o"]["kernel"] = t5x_enc_dec_attention_out
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["q"]["kernel"] = t5x_enc_dec_attention_query
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["v"]["kernel"] = t5x_enc_dec_attention_value
+        
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["layer_norm"]["weight"] = t5x_cross_layer_norm
+        
+        if split_mlp_wi:
+            flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_0"]["kernel"] = t5x_mlp_wi_0
+            flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_1"]["kernel"] = t5x_mlp_wi_1
+        else:
+            flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi"]["kernel"] = t5x_mlp_wi
+
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wo"]["kernel"] = t5x_mlp_wo
+        
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["layer_norm"]["weight"] = tx5_mlp_layer_norm
+
+    # Decoder Normalization
+    tx5_decoder_norm = t5x_model["target"]["decoder"]["decoder_norm"]["scale"]
+    flax_model.params["decoder"]["final_layer_norm"]["weight"] = tx5_decoder_norm
+
+    # Only for layer 0:
+    t5x_decoder_rel_embedding = t5x_model["target"]["decoder"]["relpos_bias"]["rel_embedding"].T
+    flax_model.params["decoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_decoder_rel_embedding
+
+    # Token Embeddings
+    tx5_token_embeddings = t5x_model["target"]["token_embedder"]["embedding"]
+    flax_model.params["shared"]["embedding"] = tx5_token_embeddings
+
+    # LM Head
+    flax_model.params["lm_head"]["kernel"] = t5x_model["target"]["decoder"]["logits_dense"]["kernel"]
+
+    flax_model.save_pretrained(flax_dump_folder_path)
+    print("T5X Model was sucessfully converted!")
+
+def convert_flax_to_pytorch(flax_dump_folder_path, pytorch_dump_folder_path):
+    model = AutoModelForSeq2SeqLM.from_pretrained(flax_dump_folder_path, from_flax=True, torch_dtype=torch.float32)
+    model.save_pretrained(pytorch_dump_folder_path)
+    print("Flax model was sucessfully converted to Pytorch!")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Required parameters
+    parser.add_argument(
+        "--t5x_checkpoint_path", default=None, type=str, required=True, help="Path the TX5 checkpoint."
+    )
+    parser.add_argument(
+        "--config_name", default=None, type=str, required=True, help="Config name of T5 model."
+    )
+    parser.add_argument(
+        "--flax_dump_folder_path", default=None, type=str, required=True, help="Path to the output FLAX model."
+    )
+    args = parser.parse_args()
+    convert_t5x_checkpoint_to_flax(args.t5x_checkpoint_path, args.config_name, args.flax_dump_folder_path)
+    convert_flax_to_pytorch(args.flax_dump_folder_path, args.flax_dump_folder_path)

export_checkpoint.py

@@ -0,0 +1,13 @@
+import argparse
+from transformers import T5ForConditionalGeneration, TFT5ForConditionalGeneration
+
+def main(args):
+    pt_model = T5ForConditionalGeneration.from_pretrained(args.model_dir, from_flax=True)
+    pt_model.save_pretrained(args.model_dir)
+    tf_model = TFT5ForConditionalGeneration.from_pretrained(args.model_dir, from_pt=True)
+    tf_model.save_pretrained(args.model_dir)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_dir', type=str, default='.')

mini_nl8.gin

@@ -0,0 +1,23 @@
+# T5.1.1 Efficient mini nl8 model.
+
+import seqio
+include 't5x/examples/t5/t5_1_1/base.gin'  # imports vocab, optimizer and model.
+
+# ------------------- Network specification overrides --------------------------
+network.Transformer.config = @network.T5Config()
+network.T5Config:
+  emb_dim = 384
+  num_heads = 8
+  num_encoder_layers = 8
+  num_decoder_layers = 8
+  head_dim = 64
+  mlp_dim = 1536
+
+# ------------------- Model specification overrides --------------------------
+VOCABULARY = @seqio.SentencePieceVocabulary()
+seqio.SentencePieceVocabulary.sentencepiece_model_file = "spiece.model"
+
+MODEL = @models.EncoderDecoderModel()
+models.EncoderDecoderModel:
+  input_vocabulary = %VOCABULARY
+  output_vocabulary = %VOCABULARY

mini_nl8_pretrain.gin

@@ -0,0 +1,26 @@
+# Register necessary SeqIO Tasks/Mixtures.
+from __gin__ import dynamic_registration
+from t5x import utils
+import ul2_tasks
+import __main__ as train_script
+
+include 'mini_nl8.gin'
+include 't5x/configs/runs/pretrain.gin'
+
+
+# ------------------- Training specification overrides --------------------------
+train_script.train:
+  eval_period = 2000
+  stats_period = 100
+
+utils.SaveCheckpointConfig:
+    period = 50000
+    keep = 4
+    use_gda = False
+
+MIXTURE_OR_TASK_NAME = "pretrain_biological_ul2"
+USE_CACHED_TASKS = False
+TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
+TRAIN_STEPS = 500000
+DROPOUT_RATE = 0.0
+BATCH_SIZE = 256

requirements.txt

@@ -0,0 +1,13 @@
+datasets>=1.16.1
+transformers>=4.13.0
+flax>=0.3.5
+optax>=0.1.0
+tqdm>=4.61.1
+numpy>=1.19.5
+tokenizers>=0.10.3
+sentencepiece>=0.1.96
+protobuf>=3.17.3,<=3.20.99
+tensorboard>=2.7.0
+torch>=1.9.0
+tensorflow>=2.7.0
+jax[tpu]>=0.2.28

ul2_objective.py

@@ -0,0 +1,195 @@
+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

ul2_tasks.py

@@ -0,0 +1,147 @@
+import functools
+from typing import Dict
+
+import seqio
+import tensorflow as tf
+from datasets import load_dataset, load_from_disk
+from t5.evaluation import metrics
+from seqio import utils, FunctionDataSource
+import t5.data
+from datasets import load_dataset, load_from_disk
+from t5.data import postprocessors
+from t5.data import preprocessors
+
+
+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
+
+vocabulary = seqio.SentencePieceVocabulary("spiece.model")
+
+DEFAULT_OUTPUT_FEATURES = {
+    "inputs": seqio.Feature(vocabulary=vocabulary, add_eos=True, required=False),
+    "targets": seqio.Feature(vocabulary=vocabulary, add_eos=True),
+}
+
+def gen_dataset(split, shuffle=False, seed=None, column="text", path=None, name=None):
+    dataset = load_dataset(path, name, streaming=True, use_auth_token=True)
+    if shuffle:
+        if seed:
+            dataset = dataset.shuffle(seed=seed)
+        else:
+            dataset = dataset.shuffle()
+    while True:
+        for item in dataset[str(split)]:
+            yield item[column]
+
+
+def dataset_fn(split, shuffle_files, seed=None, path=None, name=None):
+    return tf.data.Dataset.from_generator(
+        functools.partial(
+            gen_dataset, split, shuffle_files, seed, path=path, name=name
+        ),
+        output_signature=tf.TensorSpec(shape=(), dtype=tf.string, name=path),
+    )
+
+
+@utils.map_over_dataset
+def target_to_key(x, key_map, target_key):
+    """Assign the value from the dataset to target_key in key_map"""
+    return {**key_map, target_key: x}
+
+## First way to add to task registry
+dataset_name = 'Siddharth63/biological_dataset'
+dataset = load_dataset(dataset_name)
+
+dataset_shapes = {"train": dataset["train"].num_rows,
+                  "validation": dataset["validation"].num_rows}
+
+TaskRegistry.add(
+    "pretrain_biological_ul2",
+    source=seqio.FunctionDataSource(
+        dataset_fn=functools.partial(dataset_fn, dataset=dataset),
+        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": DEFAULT_OUTPUT_FEATURES["targets"]},
+    metric_fns=[metrics.accuracy]
+)
+
+
+## Second way to add to task registry
+# TaskRegistry.add(
+#     "pretrain_biological_ul2",
+#     source=seqio.FunctionDataSource(
+#         dataset_fn=functools.partial(
+#             dataset_fn, path="Siddharth63/biological_dataset", name="full"
+#         ),
+#         splits=("train", "validation"),
+#         caching_permitted=False,
+#     ),
+#     preprocessors=[
+#         functools.partial(
+#             target_to_key,
+#             key_map={
+#                 "inputs": "text",
+#                 "targets": "text",
+#             },
+#             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": DEFAULT_OUTPUT_FEATURES["targets"],
+#         "inputs": seqio.Feature(vocabulary=vocabulary, add_eos=True),
+#     },
+#     metric_fns=[metrics.accuracy],
+# )