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