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

Fix argument label (#4792)

Browse Source 
* Fix argument label

* Fix test
This commit is contained in:
Sylvain Gugger
2020-06-05 15:20:29 -04:00
committed by GitHub
parent 3723f30a18
commit 4dd5cf2207
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/data/data_collator.py
View File

@@ -91,7 +91,7 @@ class DataCollatorForLanguageModeling(DataCollator):
batch = self._tensorize_batch(examples)
if self.mlm:
inputs, labels = self.mask_tokens(batch)
return {"input_ids": inputs, "masked_lm_labels": labels}
return {"input_ids": inputs, "labels": labels}
else:
return {"input_ids": batch, "labels": batch}

4
tests/test_trainer.py
View File

@@ -74,14 +74,14 @@ class DataCollatorIntegrationTest(unittest.TestCase):
batch = data_collator.collate_batch(examples)
self.assertIsInstance(batch, dict)
self.assertEqual(batch["input_ids"].shape, torch.Size((31, 107)))
self.assertEqual(batch["masked_lm_labels"].shape, torch.Size((31, 107)))
self.assertEqual(batch["labels"].shape, torch.Size((31, 107)))
dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True)
examples = [dataset[i] for i in range(len(dataset))]
batch = data_collator.collate_batch(examples)
self.assertIsInstance(batch, dict)
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
self.assertEqual(batch["masked_lm_labels"].shape, torch.Size((2, 512)))
self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
@require_torch