SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Model output test (#6155)

Browse Source 
* Use return_dict=True in all tests

* Formatting
This commit is contained in:
Sylvain Gugger
2020-07-31 09:44:37 -04:00
committed by GitHub
parent 86caab1e0b
commit d951c14ae4
26 changed files with 320 additions and 765 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/test_modeling_mbart.py
View File

@@ -114,13 +114,14 @@ class MBartEnroIntegrationTest(AbstractMBartIntegrationTest):
decoder_ffn_dim=32,
max_position_embeddings=48,
add_final_layer_norm=True,
return_dict=True,
)
lm_model = BartForConditionalGeneration(config).to(torch_device)
context = torch.Tensor([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]).long().to(torch_device)
summary = torch.Tensor([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]).long().to(torch_device)
loss, logits, enc_features = lm_model(input_ids=context, decoder_input_ids=summary, labels=summary)
result = lm_model(input_ids=context, decoder_input_ids=summary, labels=summary)
expected_shape = (*summary.shape, config.vocab_size)
self.assertEqual(logits.shape, expected_shape)
self.assertEqual(result["logits"].shape, expected_shape)
@require_torch