def create_and_check_megatron_bert_for_masked_lm( self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels ): model = MegatronBertForMaskedLM(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))
# hf-experiments # @author Loreto Parisi (loretoparisi at gmail dot com) # Copyright (c) 2021 Loreto Parisi (loretoparisi at gmail dot com) import os import torch from transformers import BertTokenizer, MegatronBertForMaskedLM # The tokenizer. Megatron was trained with standard tokenizer(s). tokenizer = BertTokenizer.from_pretrained('nvidia/megatron-bert-uncased-345m', cache_dir=os.getenv( "cache_dir", "../../models")) model = MegatronBertForMaskedLM.from_pretrained( os.path.join(os.getenv("cache_dir", "../../models"), 'nvidia/megatron-bert-cased-345m')) # Copy to the device and use FP16. assert torch.cuda.is_available() device = torch.device("cuda") model.to(device) model.eval() model.half() # Create inputs (from the BERT example page). input = tokenizer("The capital of France is [MASK]", return_tensors="pt").to(device) label = tokenizer("The capital of France is Paris", return_tensors="pt")["input_ids"].to(device) # Run the model.