Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Flax RoFormer #15005

Merged
merged 7 commits into from Jan 4, 2022
Merged

Conversation

stancld
Copy link
Contributor

@stancld stancld commented Jan 2, 2022

What does this PR do?

This PR adds the flax implementation of RoFormer model.

Fixes #14605

Before submitting

Who can review?

@patrickvonplaten @patil-suraj

@stancld stancld changed the title [WIP] Add Flax RoFormer Add Flax RoFormer Jan 2, 2022
@stancld stancld marked this pull request as ready for review January 2, 2022 17:54
Copy link
Contributor

@patil-suraj patil-suraj left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks a lot, @stancld for adding this model!

Comment on lines 283 to 285
return jnp.einsum("bslh,...sh->bslh", layer, cos_pos) + jnp.einsum(
"bslh,...sh->bslh", rotate_half_layer, sin_pos
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(nit) maybe split this into two lines, would be simpler to read.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1

src/transformers/models/roformer/modeling_flax_roformer.py Outdated Show resolved Hide resolved
src/transformers/models/roformer/modeling_flax_roformer.py Outdated Show resolved Hide resolved
tests/test_modeling_flax_roformer.py Outdated Show resolved Hide resolved
@patrickvonplaten
Copy link
Contributor

Waiting for the CI to be green and then we can merge

@patrickvonplaten
Copy link
Contributor

Merging - thanks a lot for adding this model @stancld !

@patrickvonplaten patrickvonplaten merged commit 21aecc0 into huggingface:master Jan 4, 2022
stevhliu pushed a commit to stevhliu/transformers that referenced this pull request Jan 6, 2022
* Add FlaxRoFormer

* Clean code + make quality

* Fix output pooling for FlaxRoFormerForMultipleChoiceModule

* Apply suggestions from code review

* add flax model to repos

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
Albertobegue pushed a commit to Albertobegue/transformers that referenced this pull request Jan 27, 2022
* Add FlaxRoFormer

* Clean code + make quality

* Fix output pooling for FlaxRoFormerForMultipleChoiceModule

* Apply suggestions from code review

* add flax model to repos

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[Flax] Add Flax implementation of RoFormer
4 participants