Using T5 pre-trained weight for Text style transfer

Hi, I am trying to create a similar model to Riley et al., 2020 ([2010.03802] TextSETTR: Label-Free Text Style Extraction and Tunable Targeted Restyling), their model uses the the pre-trained weights for T5 model.

My approach is similar to theirs as I am trying create a model with a encoder decoder structure , which both are initialized with the T5 weights. Similar to Riley´s model mine will also include a “style extractor” which has same structure as the encoder which also needs to be initialized with the T5 weights.

I am able to access the weights using the from_pretrained() and state_dict() functions. The problem I am stuck with is loading/initializing my model with the weights. Since the model need to have same structure as the T5 model (from my understanding) to be able to load the weights. Any tips on this front ?

Were you able to figure this out? I am trying to implement a similar model and I am having a hard time understanding how the pieces connect.