| | |
| | |
| |
|
| | |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from einops import rearrange |
| | from torch import Tensor |
| |
|
| | from transformers.models.xlm_roberta.modeling_xlm_roberta import create_position_ids_from_input_ids |
| |
|
| |
|
| | class XLMRobertaEmbeddings(nn.Module): |
| | def __init__( |
| | self, |
| | embed_dim, |
| | vocab_size, |
| | max_position_embeddings, |
| | type_vocab_size, |
| | padding_idx=None, |
| | device=None, |
| | dtype=None, |
| | ): |
| | """ |
| | If max_position_embeddings <= 0, there's no position embeddings |
| | If type_vocab_size <= 0, there's no token type embeddings |
| | """ |
| | factory_kwargs = {"device": device, "dtype": dtype} |
| | super().__init__() |
| | self.word_embeddings = nn.Embedding( |
| | vocab_size, embed_dim, padding_idx=padding_idx, **factory_kwargs |
| | ) |
| | self.max_position_embeddings = max_position_embeddings |
| | self.type_vocab_size = type_vocab_size |
| | if self.max_position_embeddings > 0: |
| | self.position_embeddings = nn.Embedding( |
| | 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) |
| |
|
| | 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) |
| | """ |
| | batch_size, seqlen = input_ids.shape |
| | embeddings = self.word_embeddings(input_ids) |
| | if self.max_position_embeddings > 0: |
| | if position_ids is None: |
| | position_ids = create_position_ids_from_input_ids(input_ids, padding_idx=self.word_embeddings.padding_idx).to(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_embeddings = self.token_type_embeddings(token_type_ids) |
| | embeddings = embeddings + token_type_embeddings |
| | return embeddings |
| |
|
