def create_and_check_for_masked_lm(self, config, input_ids, token_type_ids,
input_mask, sequence_labels,
token_labels, choice_labels):
model = YosoForMaskedLM(config=config)
model.to(torch_device)
model.eval()
result = model(input_ids,
attention_mask=input_mask,
token_type_ids=token_type_ids,
labels=token_labels)
self.parent.assertEqual(
result.logits.shape,
(self.batch_size, self.seq_length, self.vocab_size))
def convert_yoso_checkpoint(checkpoint_path, yoso_config_file,
pytorch_dump_path):
orig_state_dict = torch.load(checkpoint_path,
map_location="cpu")["model_state_dict"]
config = YosoConfig.from_json_file(yoso_config_file)
model = YosoForMaskedLM(config)
new_state_dict = convert_checkpoint_helper(config.max_position_embeddings,
orig_state_dict)
print(model.load_state_dict(new_state_dict))
model.eval()
model.save_pretrained(pytorch_dump_path)
print(
f"Checkpoint successfuly converted. Model saved at {pytorch_dump_path}"
)