updates

modeling_flexbert.py

@@ -64,14 +64,13 @@ from transformers.modeling_outputs import (
     ModelOutput,
     MultipleChoiceModelOutput,
     SequenceClassifierOutput,
-    CausalLMOutput,
 )
 from transformers.models.bert.modeling_bert import BertPreTrainedModel
-
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
 from .bert_padding import index_put_first_axis
 
-from .bert_layers.activation import get_act_fn
-from .bert_layers.attention import (
+from .activation import get_act_fn
+from .attention import (
     FlexBertPaddedAttention,
     FlexBertPaddedParallelAttention,
     FlexBertPaddedRopeAttention,
@@ -81,15 +80,15 @@ from .bert_layers.attention import (
     FlexBertUnpadRopeAttention,
     FlexBertUnpadRopeParallelAttention,
 )
-from .bert_layers.configuration_bert import FlexBertConfig
-from .bert_layers.embeddings import (
+from .configuration_bert import FlexBertConfig
+from .embeddings import (
     BertAlibiEmbeddings,
     FlexBertAbsoluteEmbeddings,
     FlexBertCompiledSansPositionEmbeddings,
     FlexBertSansPositionEmbeddings,
     get_embedding_layer,
 )
-from .bert_layers.initialization import (
+from .initialization import (
     ModuleType,
     TileLinear,
     TileMode,
@@ -98,7 +97,7 @@ from .bert_layers.initialization import (
     tile_linear,
     tile_norm,
 )
-from .bert_layers.layers import (
+from .layers import (
     BertAlibiEncoder,
     BertPooler,
     BertPredictionHeadTransform,
@@ -113,7 +112,6 @@ from .bert_layers.layers import (
     FlexBertUnpadPreNormLayer,
     get_encoder_layer,
 )
-from .bert_layers.loss import get_loss_fn
 from .mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
 from .normalization import get_norm_layer
 from .padding import pad_input, unpad_input