tfm.nlp.models.T5Transformer

Transformer Encoder+Decoder for sequence to sequence.

tfm.nlp.models.T5Transformer(
    config: tfm.nlp.models.T5TransformerParams,
    compute_dtype: tf.DType = tf.float32,
    **kwargs
)

dtype the variable allocation dtype.
name a string for the module name.

checkpoint_items

Methods

create_variable

View source

create_variable(
    name: Text,
    shape: ShapeLike,
    initializer: Initializer,
    dtype: tf.DType = tf.float32,
    **kwargs
)

decode

View source

decode(
    encoded,
    decoder_target_tokens,
    encoder_input_tokens=None,
    encoder_dense_inputs=None,
    decoder_input_tokens=None,
    encoder_segment_ids=None,
    encoder_dense_segment_ids=None,
    decoder_segment_ids=None,
    decode_position=None,
    cache=None,
    max_decode_len=None,
    decode=False,
    training=False
) -> Dict[str, tf.Tensor]

encode

View source

encode(
    encoder_input_tokens=None,
    encoder_segment_ids=None,
    encoder_dense_inputs=None,
    encoder_dense_segment_ids=None,
    training=False
)

read_variable

View source

read_variable(
    variable: tf.Variable, as_dtype: Optional[tf.DType] = None
)

__call__

View source

__call__(
    encoder_input_tokens=None,
    decoder_target_tokens=None,
    encoder_dense_inputs=None,
    encoder_dense_segment_ids=None,
    decoder_input_tokens=None,
    encoder_segment_ids=None,
    decoder_segment_ids=None,
    training=False
)

Applies Transformer model on the inputs.

Args
encoder_input_tokens input tokens to the encoder.
decoder_target_tokens target tokens to the decoder.
encoder_dense_inputs input dense vectors to the encoder.
encoder_dense_segment_ids dense input segmentation info for packed
decoder_input_tokens input tokens to the decoder, only required for training.
encoder_segment_ids input segmentation info for packed examples. examples.
decoder_segment_ids target segmentation info for packed examples.
training whether it is training pass, affecting dropouts.

Returns
a dictionary of logits/cache.