INNER CODE UNIT · Python
create_t5_encoder_decoder
Ki6an/fastT5 · fastT5/onnx_exporter.py:25
def create_t5_encoder_decoder(pretrained_version="t5-base"):
"""Generates an encoder and a decoder model with a language model head from a pretrained huggingface model
Args:
pretrained_version (str): Name of a pretrained model, or path to a pretrained / finetuned version of T5
Returns:
simplified_encoder: pytorch t5 encoder with a wrapper to output only the hidden states
decoder_with_lm_head: pytorch t5 decoder with a language modeling head
"""
if 'mt5' in pretrained_version:
model = MT5ForConditionalGeneration.from_pretrained(pretrained_version, use_auth_token=get_auth_token())
else:
model = T5ForConditionalGeneration.from_pretrained(pretrained_version, use_auth_token=get_auth_token())
return turn_model_into_encoder_decoder(model)