add modeling

modeling_flexbert.py

@@ -68,10 +68,10 @@ from transformers.modeling_outputs import (
 )
 from transformers.models.bert.modeling_bert import BertPreTrainedModel
 
-from bert_padding import index_put_first_axis
+from .bert_padding import index_put_first_axis
 
-from src.bert_layers.activation import get_act_fn
-from src.bert_layers.attention import (
+from .bert_layers.activation import get_act_fn
+from .bert_layers.attention import (
     FlexBertPaddedAttention,
     FlexBertPaddedParallelAttention,
     FlexBertPaddedRopeAttention,
@@ -81,15 +81,15 @@ from src.bert_layers.attention import (
     FlexBertUnpadRopeAttention,
     FlexBertUnpadRopeParallelAttention,
 )
-from src.bert_layers.configuration_bert import FlexBertConfig
-from src.bert_layers.embeddings import (
+from .bert_layers.configuration_bert import FlexBertConfig
+from .bert_layers.embeddings import (
     BertAlibiEmbeddings,
     FlexBertAbsoluteEmbeddings,
     FlexBertCompiledSansPositionEmbeddings,
     FlexBertSansPositionEmbeddings,
     get_embedding_layer,
 )
-from src.bert_layers.initialization import (
+from .bert_layers.initialization import (
     ModuleType,
     TileLinear,
     TileMode,
@@ -98,7 +98,7 @@ from src.bert_layers.initialization import (
     tile_linear,
     tile_norm,
 )
-from src.bert_layers.layers import (
+from .bert_layers.layers import (
     BertAlibiEncoder,
     BertPooler,
     BertPredictionHeadTransform,
@@ -113,10 +113,10 @@ from src.bert_layers.layers import (
     FlexBertUnpadPreNormLayer,
     get_encoder_layer,
 )
-from src.bert_layers.loss import get_loss_fn
-from src.bert_layers.mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
-from src.bert_layers.normalization import get_norm_layer
-from src.bert_layers.padding import pad_input, unpad_input
+from .bert_layers.loss import get_loss_fn
+from .bert_layers.mlp import FlexBertGLU, FlexBertMLP, FlexBertParallelGLU
+from .bert_layers.normalization import get_norm_layer
+from .bert_layers.padding import pad_input, unpad_input
 
 logger = logging.getLogger(__name__)
 
@@ -868,7 +868,7 @@ class FlexBertPreTrainedModel(BertPreTrainedModel):
 
     def _init_module_weights(self, module: nn.Module):
         """
-        Custom weight init of modules using src.bert_layers.initialization.init_weights
+        Custom weight init of modules using .bert_layers.initialization.init_weights
         Currently only supports init of embedding modules
         """
         assert isinstance(module, nn.Module)