INNER CODE UNIT · Python

turn_model_into_encoder_decoder

Ki6an/fastT5 · fastT5/onnx_exporter.py:44

def turn_model_into_encoder_decoder(model):
    encoder = model.encoder
    decoder = model.decoder
    lm_head = model.lm_head

    decoder_with_lm_head = DecoderWithLMhead(decoder, lm_head, model.config)
    simplified_encoder = T5Encoder(encoder)
    decoder_with_lm_head_init = DecoderWithLMheadInitial(decoder, lm_head, model.config)

    return simplified_encoder, decoder_with_lm_head, decoder_with_lm_head_init


def generate_onnx_representation(
    pretrained_version=None,
    model=None,
    output_path=None,
    input_sequence_length=256,
    onnx_opset_version=12,  # no other opset versions are tested, change at your own risk

View source record →