FlexBert / __init__.py
NohTow's picture
Using dict as input
ce9aa51
from .attention import (
BertAlibiUnpadAttention,
BertAlibiUnpadSelfAttention,
BertSelfOutput,
FlexBertPaddedAttention,
FlexBertUnpadAttention,
)
from .embeddings import (
BertAlibiEmbeddings,
FlexBertAbsoluteEmbeddings,
FlexBertSansPositionEmbeddings,
)
from .layers import (
BertAlibiEncoder,
BertAlibiLayer,
BertResidualGLU,
FlexBertPaddedPreNormLayer,
FlexBertPaddedPostNormLayer,
FlexBertUnpadPostNormLayer,
FlexBertUnpadPreNormLayer,
)
from .modeling_flexbert import (
BertLMPredictionHead,
BertModel,
BertForMaskedLM,
BertForSequenceClassification,
BertForMultipleChoice,
BertOnlyMLMHead,
BertOnlyNSPHead,
BertPooler,
BertPredictionHeadTransform,
FlexBertModel,
FlexBertForMaskedLM,
FlexBertForSequenceClassification,
FlexBertForMultipleChoice,
)
from .bert_padding import(
IndexFirstAxis,
IndexPutFirstAxis
)
__all__ = [
"BertAlibiEmbeddings",
"BertAlibiEncoder",
"BertForMaskedLM",
"BertForSequenceClassification",
"BertForMultipleChoice",
"BertResidualGLU",
"BertAlibiLayer",
"BertLMPredictionHead",
"BertModel",
"BertOnlyMLMHead",
"BertOnlyNSPHead",
"BertPooler",
"BertPredictionHeadTransform",
"BertSelfOutput",
"BertAlibiUnpadAttention",
"BertAlibiUnpadSelfAttention",
"FlexBertPaddedAttention",
"FlexBertUnpadAttention",
"FlexBertAbsoluteEmbeddings",
"FlexBertSansPositionEmbeddings",
"FlexBertPaddedPreNormLayer",
"FlexBertPaddedPostNormLayer",
"FlexBertUnpadPostNormLayer",
"FlexBertUnpadPreNormLayer",
"FlexBertModel",
"FlexBertForMaskedLM",
"FlexBertForSequenceClassification",
"FlexBertForMultipleChoice",
"IndexFirstAxis",
"IndexPutFirstAxis"
]