Pass an embedding layer to the constructor of the BertModel class

examples/BERT/model.py

@@ -43,7 +43,8 @@ def __init__(self, ntoken, ninp, dropout=0.5):
  self.norm = LayerNorm(ninp)
  self.dropout = Dropout(dropout)
 
- def forward(self, src, token_type_input):
+ def forward(self, seq_inputs):
+ src, token_type_input = seq_inputs
  src = self.embed(src) + self.pos_embed(src) \
  + self.tok_type_embed(src, token_type_input)
  return self.dropout(self.norm(src))
@@ -99,16 +100,16 @@ def forward(self, src, src_mask=None, src_key_padding_mask=None):
 class BertModel(nn.Module):
  """Contain a transformer encoder."""
 
- def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
+ def __init__(self, ntoken, ninp, nhead, nhid, nlayers, embed_layer, dropout=0.5):
  super(BertModel, self).__init__()
  self.model_type = 'Transformer'
- self.bert_embed = BertEmbedding(ntoken, ninp)
+ self.bert_embed = embed_layer
  encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
  self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
  self.ninp = ninp
 
- def forward(self, src, token_type_input):
- src = self.bert_embed(src, token_type_input)
+ def forward(self, seq_inputs):
+ src = self.bert_embed(seq_inputs)
  output = self.transformer_encoder(src)
  return output
 
@@ -118,15 +119,16 @@ class MLMTask(nn.Module):
 
  def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
  super(MLMTask, self).__init__()
- self.bert_model = BertModel(ntoken, ninp, nhead, nhid, nlayers, dropout=0.5)
+ embed_layer = BertEmbedding(ntoken, ninp)
+ self.bert_model = BertModel(ntoken, ninp, nhead, nhid, nlayers, embed_layer, dropout=0.5)
  self.mlm_span = Linear(ninp, ninp)
  self.activation = F.gelu
  self.norm_layer = LayerNorm(ninp, eps=1e-12)
  self.mlm_head = Linear(ninp, ntoken)
 
  def forward(self, src, token_type_input=None):
  src = src.transpose(0, 1) # Wrap up by nn.DataParallel
- output = self.bert_model(src, token_type_input)
+ output = self.bert_model((src, token_type_input))
  output = self.mlm_span(output)
  output = self.activation(output)
  output = self.norm_layer(output)
@@ -147,7 +149,7 @@ def __init__(self, bert_model):
 
  def forward(self, src, token_type_input):
  src = src.transpose(0, 1) # Wrap up by nn.DataParallel
- output = self.bert_model(src, token_type_input)
+ output = self.bert_model((src, token_type_input))
  # Send the first <'cls'> seq to a classifier
  output = self.activation(self.linear_layer(output[0]))
  output = self.ns_span(output)
@@ -164,7 +166,7 @@ def __init__(self, bert_model):
  self.qa_span = Linear(bert_model.ninp, 2)
 
  def forward(self, src, token_type_input):
- output = self.bert_model(src, token_type_input)
+ output = self.bert_model((src, token_type_input))
  # transpose output (S, N, E) to (N, S, E)
  output = output.transpose(0, 1)
  output = self.activation(output)

examples/BERT/ns_task.py

@@ -5,7 +5,7 @@
 import torch.nn as nn
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
-from model import NextSentenceTask, BertModel
+from model import NextSentenceTask, BertModel, BertEmbedding
 from utils import run_demo, run_ddp, wrap_up
 
 
@@ -149,7 +149,8 @@ def run_main(args, rank=None):
  if args.checkpoint != 'None':
  model = torch.load(args.checkpoint)
  else:
- pretrained_bert = BertModel(len(vocab), args.emsize, args.nhead, args.nhid, args.nlayers, args.dropout)
+ embed_layer = BertEmbedding(len(vocab), args.emsize)
+ pretrained_bert = BertModel(len(vocab), args.emsize, args.nhead, args.nhid, args.nlayers, embed_layer, args.dropout)
  pretrained_bert.load_state_dict(torch.load(args.bert_model))
  model = NextSentenceTask(pretrained_bert)
 

examples/BERT/qa_task.py

@@ -9,7 +9,7 @@
 from model import QuestionAnswerTask
 from metrics import compute_qa_exact, compute_qa_f1
 from utils import print_loss_log
-from model import BertModel
+from model import BertModel, BertEmbedding
 
 
 def process_raw_data(data):
@@ -174,7 +174,8 @@ def train():
  train_dataset = process_raw_data(train_dataset)
  dev_dataset = process_raw_data(dev_dataset)
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
- pretrained_bert = BertModel(len(vocab), args.emsize, args.nhead, args.nhid, args.nlayers, args.dropout)
+ embed_layer = BertEmbedding(len(vocab), args.emsize)
+ pretrained_bert = BertModel(len(vocab), args.emsize, args.nhead, args.nhid, args.nlayers, embed_layer, args.dropout)
  pretrained_bert.load_state_dict(torch.load(args.bert_model))
  model = QuestionAnswerTask(pretrained_bert).to(device)
  criterion = nn.CrossEntropyLoss()