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.