clean up embeddings file

embedding.py

@@ -7,10 +7,9 @@ https://github.com/Dao-AILab/flash-attention/blob/43950dda456e095969d842fca7a73c
 
 import torch
 import torch.nn as nn
-from torch import Tensor
 
 
-class BertEmbeddings(nn.Module):
+class JinaBertEmbeddings(nn.Module):
     def __init__(
         self,
         embed_dim,
@@ -37,24 +36,32 @@ class BertEmbeddings(nn.Module):
                 max_position_embeddings, embed_dim, **factory_kwargs
             )
         if self.type_vocab_size > 0:
-            self.token_type_embeddings = nn.Embedding(type_vocab_size, embed_dim, **factory_kwargs)
+            self.token_type_embeddings = nn.Embedding(
+                type_vocab_size, embed_dim, **factory_kwargs
+            )
 
     def forward(self, input_ids, position_ids=None, token_type_ids=None):
         """
         input_ids: (batch, seqlen)
         position_ids: (batch, seqlen)
         token_type_ids: (batch, seqlen)
+        ..note: layer norm and dropout has been taken out from Embeddings forward, but in `moddeling_bert.py`.
+                This is different from jina_bert_implementation.
         """
-        batch_size, seqlen = input_ids.shape
+        _, seqlen = input_ids.shape
         embeddings = self.word_embeddings(input_ids)
         if self.max_position_embeddings > 0:
             if position_ids is None:
-                position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
+                position_ids = torch.arange(
+                    seqlen, dtype=torch.long, device=input_ids.device
+                )
             position_embeddings = self.position_embeddings(position_ids)
             embeddings = embeddings + position_embeddings
         if self.type_vocab_size > 0:
             if token_type_ids is None:
-                token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=input_ids.device)
+                token_type_ids = torch.zeros(
+                    seqlen, dtype=torch.long, device=input_ids.device
+                )
             token_type_embeddings = self.token_type_embeddings(token_type_ids)
             embeddings = embeddings + token_type_embeddings
         return embeddings

modeling_bert.py

@@ -37,7 +37,7 @@ from .bert_padding import (
 )
 
 from .block import Block
-from .embedding import BertEmbeddings
+from .embedding import JinaBertEmbeddings
 from .mha import MHA
 from .mlp import FusedMLP, Mlp