examples/tensorflow/summarization/run_summarization.py/dataset_to_tf
def dataset_to_tf(dataset, model, tokenizer, total_batch_size, num_epochs, shuffle):
    if dataset is None:
        return None
    train_generator = partial(sample_generator, dataset, model, tokenizer, shuffle=shuffle)
    train_signature = {
        feature: tf.TensorSpec(shape=(None,), dtype=tf.int32)
        for feature in dataset.features
        if feature != "special_tokens_mask"
    }
    if (
        model is not None
        and "decoder_input_ids" not in train_signature
        and hasattr(model, "prepare_decoder_input_ids_from_labels")
    ):
        train_signature["decoder_input_ids"] = train_signature["labels"]
    # This may need to be changed depending on your particular model or tokenizer!
    padding_values = {
        key: tf.convert_to_tensor(tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0, dtype=tf.int32)
        for key in train_signature.keys()
    }
    padding_values["labels"] = tf.convert_to_tensor(-100, dtype=tf.int32)
    train_signature["labels"] = train_signature["input_ids"]
    train_signature = (train_signature, train_signature["labels"])
    options = tf.data.Options()
    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
    tf_dataset = (
        tf.data.Dataset.from_generator(train_generator, output_signature=train_signature)
        .with_options(options)
        .padded_batch(
            batch_size=total_batch_size,
            drop_remainder=True,
            padding_values=(padding_values, np.array(-100, dtype=np.int32)),
        )
        .repeat(int(num_epochs))
    )
    return tf_dataset