def test(informal):
if torch.cuda.is_available():
device = torch.device('cuda:3')
print(f'Using GPU device: {device}')
else:
device = torch.device('cpu')
print(f'GPU is not available, using CPU device {device}')
test_config = {'batch_size': 5, 'epoch': 29, 'save_dir': './checkpoints/'}
test_dataset = FormalDataset(informal)
dataloader = DataLoader(test_dataset,
batch_size=test_config['batch_size'],
shuffle=False,
num_workers=4,
drop_last=False)
config = DistilBertConfig()
model = DistilBertForMaskedLM(config)
load_model(test_config['epoch'], model, test_config['save_dir'])
model.to(device)
model.eval()
with torch.no_grad():
for i, batch in tqdm(enumerate(dataloader)):
inp = batch['input_ids'].to(device)
attn = batch['attention_mask'].to(device)
logits = model(input_ids=inp, attention_mask=attn)[0]
preds = decode_text(test_dataset.tokenizer, logits)
for seq in preds:
with open('test_pred.txt', 'a') as res_file:
res_file.writelines(seq + '\n')
def create_and_check_distilbert_for_masked_lm(
self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
):
model = DistilBertForMaskedLM(config=config)
model.to(torch_device)
model.eval()
result = model(input_ids, attention_mask=input_mask, labels=token_labels)
self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
def create_and_check_distilbert_for_masked_lm(self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels):
model = DistilBertForMaskedLM(config=config)
model.eval()
loss, prediction_scores = model(input_ids, attention_mask=input_mask, masked_lm_labels=token_labels)
result = {
"loss": loss,
"prediction_scores": prediction_scores,
}
self.parent.assertListEqual(
list(result["prediction_scores"].size()),
[self.batch_size, self.seq_length, self.vocab_size])
self.check_loss_output(result)