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

Fix #7284 (#7289)

Browse Source
This commit is contained in:
Sylvain Gugger
2020-09-21 10:31:26 -04:00
committed by GitHub
parent 8d464374ba
commit 63276b76d4
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -434,13 +434,15 @@ class DataCollatorForNextSentencePrediction:
else: else:
input_ids = self._tensorize_batch(input_ids) input_ids = self._tensorize_batch(input_ids)
return { result = {
"input_ids": input_ids, "input_ids": input_ids,
"attention_mask": self._tensorize_batch(attention_masks), "attention_mask": self._tensorize_batch(attention_masks),
"token_type_ids": self._tensorize_batch(segment_ids), "token_type_ids": self._tensorize_batch(segment_ids),
"masked_lm_labels": mlm_labels if self.mlm else None,
"next_sentence_label": torch.tensor(nsp_labels), "next_sentence_label": torch.tensor(nsp_labels),
} }
if self.mlm:
result["masked_lm_labels"] = mlm_labels
return result
def _tensorize_batch(self, examples: List[torch.Tensor]) -> torch.Tensor: def _tensorize_batch(self, examples: List[torch.Tensor]) -> torch.Tensor:
length_of_first = examples[0].size(0) length_of_first = examples[0].size(0)