[Flax] Add FlaxMBart #12236
[Flax] Add FlaxMBart #12236
Conversation
There was a problem hiding this comment.
Thanks a lot for adding this @stancld!
It's looking great overall, I left a few comments. Specifically
- The order of layer norm and attention and
- Could we add as much copied from statements as possible?
* Fix shift_tokens_right method according to MBart implementation * Update shift_tokens_right in tests accordingly * Fix the import issue and update docs file * make style quality
* Change the order of normalization layer and attention * Add some copu statementes
|
@patil-suraj Thank you a lot for your suggestions. I fixed the order of attention and normalization layers and some other minor bugs. Also added some additional copy statements. I also changed the |
Besides, add `lang_code_to_id` to MBartTokenizeFast
|
@stancld I pushed a couple of commits to add the |
| [prev_output_tokens[i, eos_idx] for i, eos_idx in enumerate(index_of_eos)] | ||
| ).squeeze() | ||
| # for loop basically does jax-compatible version of prev_output_tokens[:, 1:] = prev_output_tokens[:, :-1].clone() | ||
| for i in range(prev_output_tokens.shape[1], 0, -1): |
What does this PR do?
This PR adds flax implementation of MBart.
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
@patrickvonplaten @patil-suraj