cls_token_id = decoder_tokenizer.vocab["[CLS]"],
mask_token_id = decoder_tokenizer.vocab["[MASK]"],
bos_token_id = decoder_tokenizer.vocab["[BOS]"],
eos_token_id = decoder_tokenizer.vocab["[EOS]"],
)
# Initialize a brand new bert-based decoder.
decoder = BertGenerationDecoder(config=decoder_config)
# Setup enc-decoder mode.
bert2bert = EncoderDecoderModel(encoder=encoder, decoder=decoder)
bert2bert.config.decoder_start_token_id=decoder_tokenizer.vocab["[CLS]"]
bert2bert.config.pad_token_id=decoder_tokenizer.vocab["[PAD]"]
# Elementary Training.
optimizer = torch.optim.Adam(bert2bert.parameters(), lr=0.000001)
bert2bert.cuda()
for epoch in range(30):
print("*"*50, "Epoch", epoch, "*"*50)
if True:
for batch in tqdm(sierra_dl):
# tokenize commands and goals.
inputs = encoder_tokenizer(batch["command"], add_special_tokens=True, return_tensors="pt", padding=True, truncation=True)
labels = decoder_tokenizer(batch["symbolic_plan_processed"], return_tensors="pt", padding=True, max_length=sierra_ds.max_plan_length, truncation=True, add_special_tokens=True, )
# Move to GPU.
for key,item in inputs.items():
if type(item).__name__ == "Tensor":
inputs[key] = item.cuda()
for key, item in labels.items():
if type(item).__name__ == "Tensor":
decoder_tokenizer = BasicTokenizer("yz/vocab.txt")
config_decoder.update({
"vocab_size": len(decoder_tokenizer.vocab),
"num_hidden_layers":3,
"num_attention_heads":3
})
config = EncoderDecoderConfig.from_encoder_decoder_configs(config_encoder, config_decoder)
# 导入模型 BERT
model = EncoderDecoderModel(config=config)
model.encoder = BertModel.from_pretrained('bert-base-uncased')
if args.gpu and torch.cuda.is_available():
model = model.cuda()
loss_fun = nn.CrossEntropyLoss()
#loss_fun = nn.CrossEntropyLoss(ignore_index=0)
optimizer = torch.optim.Adam(model.decoder.parameters(), lr=args.lr)
# 记录时间
begin_time = datetime.now()
print("Start training BERT: ", begin_time)
# 开始训练
for epoch in range(args.epoch):
model.train()
train_loss_list = []
train_acc_list = []
