This sounds very interesting! I’ve been working on a similar task of Question Generation using T5 and BERT and this would probably be an ideal extension to that. I’ve no prior experience with FLAX/JAX so also would be great to learn about them. I’m interested to join and contribute if there’s still any space. Looking forward to it 
I’m a Master’s student and live in Canada. PT/PDT time zone.