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.

Sign up for GitHub

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

[Flax] Add Electra models #11426

Merged
merged 14 commits into from
May 4, 2021
Merged

[Flax] Add Electra models #11426

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
5b61735
add electra model to flax
CoderPat Apr 25, 2021
343d61e
Remove Electra Next Sentence Prediction model added by mistake
CoderPat Apr 25, 2021
1fc055b
fix parameter sharing and loosen equality threshold
CoderPat Apr 25, 2021
53d1835
fix styling issues
CoderPat Apr 25, 2021
8d4621e
add mistaken removen imports
CoderPat Apr 25, 2021
f536d77
fix electra table
CoderPat Apr 25, 2021
f41e4db
Add FlaxElectra to automodels and fixe docs
CoderPat Apr 25, 2021
9f96419
fix issues pointed out the PR
CoderPat Apr 26, 2021
283425c
Merge branch 'master' into add-flax-electra
CoderPat Apr 26, 2021
7e7e424
Merge remote-tracking branch 'upstream/master' into add-flax-electra
CoderPat May 3, 2021
2d19a8c
fix flax electra to comply with latest changes
CoderPat May 4, 2021
4fad31d
remove stale class
CoderPat May 4, 2021
ef9b4cf
fix conflict and merge to master
patrickvonplaten May 4, 2021
b60c2ff
add copied from
patrickvonplaten May 4, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions src/transformers/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1390,6 +1390,19 @@
"FlaxBertPreTrainedModel",
]
)
_import_structure["models.electra"].extend(
[
"FlaxElectraForMaskedLM",
"FlaxElectraForMultipleChoice",
"FlaxElectraForPreTraining",
"FlaxElectraForQuestionAnswering",
"FlaxElectraForSequenceClassification",
"FlaxElectraForTokenClassification",
"FlaxElectraModel",
"FlaxElectraPreTrainedModel",
]
)

_import_structure["models.roberta"].append("FlaxRobertaModel")
else:
from .utils import dummy_flax_objects
Expand Down Expand Up @@ -2551,6 +2564,16 @@
FlaxBertModel,
FlaxBertPreTrainedModel,
)
from .models.electra import (
FlaxElectraForMaskedLM,
FlaxElectraForMultipleChoice,
FlaxElectraForPreTraining,
FlaxElectraForQuestionAnswering,
FlaxElectraForSequenceClassification,
FlaxElectraForTokenClassification,
FlaxElectraModel,
FlaxElectraPreTrainedModel,
)
from .models.roberta import FlaxRobertaModel
else:
# Import the same objects as dummies to get them in the namespace.
Expand Down
26 changes: 25 additions & 1 deletion src/transformers/models/electra/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@

from typing import TYPE_CHECKING

from ...file_utils import _BaseLazyModule, is_tf_available, is_tokenizers_available, is_torch_available
from ...file_utils import _BaseLazyModule, is_tf_available, is_tokenizers_available, is_torch_available, is_flax_available


_import_structure = {
Expand Down Expand Up @@ -56,6 +56,18 @@
"TFElectraPreTrainedModel",
]

if is_flax_available():
_import_structure["modeling_flax_electra"] = [
"FlaxElectraForMaskedLM",
"FlaxElectraForMultipleChoice",
"FlaxElectraForPreTraining",
"FlaxElectraForQuestionAnswering",
"FlaxElectraForSequenceClassification",
"FlaxElectraForTokenClassification",
"FlaxElectraModel",
"FlaxElectraPreTrainedModel",
]


if TYPE_CHECKING:
from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig
Expand Down Expand Up @@ -91,6 +103,18 @@
TFElectraPreTrainedModel,
)

if is_flax_available():
from .modeling_flax_electra import (
FlaxElectraForMaskedLM,
FlaxElectraForMultipleChoice,
FlaxElectraForPreTraining,
FlaxElectraForQuestionAnswering,
FlaxElectraForSequenceClassification,
FlaxElectraForTokenClassification,
FlaxElectraModel,
FlaxElectraPreTrainedModel,
)

else:
import importlib
import os
Expand Down
Loading