How to design custom Seq2seq model by evojax?

Dear developer:
i am developing a large seq2seq model through evojax. However, i found it is inflexible to develop my custom model with custom vocabulary based on the example seq2seq, which means i need to revise the source code, not import functions or revise parameters. I would appreciate it if you could give me some guidance. 
By the way, thanks for your awesome work.