[Flax] Add other BERT classes#10977
Conversation
LysandreJik
left a comment
There was a problem hiding this comment.
Cool, LGTM! Great work implementing all of these. Will RoBERTa follow?
It is so similar to the PyTorch implementation it seems a script could take care of the implementation by copying the PyTorch one and replacing a few strings!
sgugger
left a comment
There was a problem hiding this comment.
Looks great! I don't know what happened to your formatter but we can definitely have most of submodule definitions fit on one line ;-)
@marcvanzee and I were also wondering about this in general -- is there a 80/20 solution that requires user input in some cases? It would have to not introduce silent errors (e.g. a model that seems to run the same but differs in some hard-to-find way). |
* add first code structures * add all bert models * add to init and docs * correct docs * make style
What does this PR do?
This PR adds the other BERT model classes for Flax.
Also the following checkpoints have been uploaded for Flax:
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors which may be interested in your PR.