Construct the embeddings from word, position and token_type embeddings.
| 146 | |
| 147 | |
| 148 | class BertEmbeddings(nn.Module): |
| 149 | """Construct the embeddings from word, position and token_type embeddings. |
| 150 | """ |
| 151 | |
| 152 | def __init__(self, config): |
| 153 | super().__init__() |
| 154 | self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) |
| 155 | self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) |
| 156 | self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) |
| 157 | |
| 158 | # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load |
| 159 | # any TensorFlow checkpoint file |
| 160 | self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
| 161 | self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| 162 | |
| 163 | def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None): |
| 164 | if input_ids is not None: |
| 165 | input_shape = input_ids.size() |
| 166 | else: |
| 167 | input_shape = inputs_embeds.size()[:-1] |
| 168 | |
| 169 | seq_length = input_shape[1] |
| 170 | device = input_ids.device if input_ids is not None else inputs_embeds.device |
| 171 | if position_ids is None: |
| 172 | position_ids = torch.arange(seq_length, dtype=torch.long, device=device) |
| 173 | position_ids = position_ids.unsqueeze(0).expand(input_shape) |
| 174 | if token_type_ids is None: |
| 175 | token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) |
| 176 | |
| 177 | if inputs_embeds is None: |
| 178 | inputs_embeds = self.word_embeddings(input_ids) |
| 179 | position_embeddings = self.position_embeddings(position_ids) |
| 180 | token_type_embeddings = self.token_type_embeddings(token_type_ids) |
| 181 | |
| 182 | embeddings = inputs_embeds + position_embeddings + token_type_embeddings |
| 183 | embeddings = self.LayerNorm(embeddings) |
| 184 | embeddings = self.dropout(embeddings) |
| 185 | return embeddings |
| 186 | |
| 187 | |
| 188 | class BertSelfAttention(nn.Module): |