We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
data_collator.py does not allow NoneType labels for test set predictions on Glue
# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): Distilbert

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [x] the official example scripts: (give details below) 
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)


## To reproduce

Steps to reproduce the behavior:

1. Open the  example Colab for text-classification from here https://huggingface.co/transformers/examples.html
2. Try to run the prediction function found in run_glue.py to predict on the official Glue test set. 
3. The error is as shown below. 

Earlier, this has worked with the exact same program, ever since an update recently, this error shows up. 

`TypeError                                 Traceback (most recent call last)
<ipython-input-16-9eecdd4d48b1> in <module>()
      2 output_mode = "classification"
      3 
----> 4 predictions = trainer.predict(test_dataset=test_dataset).predictions
      5 if output_mode == "classification":
      6     predictions = np.argmax(predictions, axis=1)

7 frames
/usr/local/lib/python3.6/dist-packages/transformers/data/data_collator.py in default_data_collator(features)
     45     if "label" in first:
     46         dtype = torch.long if type(first["label"]) is int else torch.float
---> 47         batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
     48     elif "label_ids" in first:
     49         if isinstance(first["label_ids"], torch.Tensor):

TypeError: must be real number, not NoneType`

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
The error can be seen in the Colab notebook here https://colab.research.google.com/drive/1H_92qdsOOql2hS210qNrfMEEMRcAoHD_?usp=sharing

<!-- A clear and concise description of what you would expect to happen. -->

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version:
- Platform: Colab
- Python version: NA
- PyTorch version (GPU?):
- Tensorflow version (GPU?):
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?:

data_collator.py does not allow NoneType labels for test set predictions on Glue
# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): Distilbert

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [x] the official example scripts: (give details below) 
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)


## To reproduce

Steps to reproduce the behavior:

1. Open the  example Colab for text-classification from here https://huggingface.co/transformers/examples.html
2. Try to run the prediction function found in run_glue.py to predict on the official Glue test set. 
3. The error is as shown below. 

Earlier, this has worked with the exact same program, ever since an update recently, this error shows up. 

`TypeError                                 Traceback (most recent call last)
<ipython-input-16-9eecdd4d48b1> in <module>()
      2 output_mode = "classification"
      3 
----> 4 predictions = trainer.predict(test_dataset=test_dataset).predictions
      5 if output_mode == "classification":
      6     predictions = np.argmax(predictions, axis=1)

7 frames
/usr/local/lib/python3.6/dist-packages/transformers/data/data_collator.py in default_data_collator(features)
     45     if "label" in first:
     46         dtype = torch.long if type(first["label"]) is int else torch.float
---> 47         batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
     48     elif "label_ids" in first:
     49         if isinstance(first["label_ids"], torch.Tensor):

TypeError: must be real number, not NoneType`

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
The error can be seen in the Colab notebook here https://colab.research.google.com/drive/1H_92qdsOOql2hS210qNrfMEEMRcAoHD_?usp=sharing

<!-- A clear and concise description of what you would expect to happen. -->

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version:
- Platform: Colab
- Python version: NA
- PyTorch version (GPU?):
- Tensorflow version (GPU?):
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?:


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/transformers/data/data_collator.py`
Content:
```
1 from dataclasses import dataclass
2 from typing import Any, Callable, Dict, List, NewType, Tuple
3 
4 import torch
5 from torch.nn.utils.rnn import pad_sequence
6 
7 from ..tokenization_utils import PreTrainedTokenizer
8 
9 
10 InputDataClass = NewType("InputDataClass", Any)
11 
12 """
13 A DataCollator is a function that takes a list of samples from a Dataset
14 and collate them into a batch, as a dictionary of Tensors.
15 """
16 DataCollator = NewType("DataCollator", Callable[[List[InputDataClass]], Dict[str, torch.Tensor]])
17 
18 
19 def default_data_collator(features: List[InputDataClass]) -> Dict[str, torch.Tensor]:
20     """
21     Very simple data collator that:
22     - simply collates batches of dict-like objects
23     - Performs special handling for potential keys named:
24         - `label`: handles a single value (int or float) per object
25         - `label_ids`: handles a list of values per object
26     - does not do any additional preprocessing
27 
28     i.e., Property names of the input object will be used as corresponding inputs to the model.
29     See glue and ner for example of how it's useful.
30     """
31 
32     # In this function we'll make the assumption that all `features` in the batch
33     # have the same attributes.
34     # So we will look at the first element as a proxy for what attributes exist
35     # on the whole batch.
36     if not isinstance(features[0], dict):
37         features = [vars(f) for f in features]
38 
39     first = features[0]
40     batch = {}
41 
42     # Special handling for labels.
43     # Ensure that tensor is created with the correct type
44     # (it should be automatically the case, but let's make sure of it.)
45     if "label" in first:
46         dtype = torch.long if type(first["label"]) is int else torch.float
47         batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
48     elif "label_ids" in first:
49         if isinstance(first["label_ids"], torch.Tensor):
50             batch["labels"] = torch.stack([f["label_ids"] for f in features])
51         else:
52             dtype = torch.long if type(first["label_ids"][0]) is int else torch.float
53             batch["labels"] = torch.tensor([f["label_ids"] for f in features], dtype=dtype)
54 
55     # Handling of all other possible keys.
56     # Again, we will use the first element to figure out which key/values are not None for this model.
57     for k, v in first.items():
58         if k not in ("label", "label_ids") and v is not None and not isinstance(v, str):
59             if isinstance(v, torch.Tensor):
60                 batch[k] = torch.stack([f[k] for f in features])
61             else:
62                 batch[k] = torch.tensor([f[k] for f in features], dtype=torch.long)
63 
64     return batch
65 
66 
67 @dataclass
68 class DataCollatorForLanguageModeling:
69     """
70     Data collator used for language modeling.
71     - collates batches of tensors, honoring their tokenizer's pad_token
72     - preprocesses batches for masked language modeling
73     """
74 
75     tokenizer: PreTrainedTokenizer
76     mlm: bool = True
77     mlm_probability: float = 0.15
78 
79     def __call__(self, examples: List[torch.Tensor]) -> Dict[str, torch.Tensor]:
80         batch = self._tensorize_batch(examples)
81         if self.mlm:
82             inputs, labels = self.mask_tokens(batch)
83             return {"input_ids": inputs, "labels": labels}
84         else:
85             return {"input_ids": batch, "labels": batch}
86 
87     def _tensorize_batch(self, examples: List[torch.Tensor]) -> torch.Tensor:
88         length_of_first = examples[0].size(0)
89         are_tensors_same_length = all(x.size(0) == length_of_first for x in examples)
90         if are_tensors_same_length:
91             return torch.stack(examples, dim=0)
92         else:
93             if self.tokenizer._pad_token is None:
94                 raise ValueError(
95                     "You are attempting to pad samples but the tokenizer you are using"
96                     f" ({self.tokenizer.__class__.__name__}) does not have one."
97                 )
98             return pad_sequence(examples, batch_first=True, padding_value=self.tokenizer.pad_token_id)
99 
100     def mask_tokens(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
101         """
102         Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
103         """
104 
105         if self.tokenizer.mask_token is None:
106             raise ValueError(
107                 "This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer."
108             )
109 
110         labels = inputs.clone()
111         # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)
112         probability_matrix = torch.full(labels.shape, self.mlm_probability)
113         special_tokens_mask = [
114             self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
115         ]
116         probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)
117         if self.tokenizer._pad_token is not None:
118             padding_mask = labels.eq(self.tokenizer.pad_token_id)
119             probability_matrix.masked_fill_(padding_mask, value=0.0)
120         masked_indices = torch.bernoulli(probability_matrix).bool()
121         labels[~masked_indices] = -100  # We only compute loss on masked tokens
122 
123         # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
124         indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
125         inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
126 
127         # 10% of the time, we replace masked input tokens with random word
128         indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
129         random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)
130         inputs[indices_random] = random_words[indices_random]
131 
132         # The rest of the time (10% of the time) we keep the masked input tokens unchanged
133         return inputs, labels
134
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/transformers/data/data_collator.py b/src/transformers/data/data_collator.py
--- a/src/transformers/data/data_collator.py
+++ b/src/transformers/data/data_collator.py
@@ -42,10 +42,10 @@
     # Special handling for labels.
     # Ensure that tensor is created with the correct type
     # (it should be automatically the case, but let's make sure of it.)
-    if "label" in first:
+    if "label" in first and first["label"] is not None:
         dtype = torch.long if type(first["label"]) is int else torch.float
         batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
-    elif "label_ids" in first:
+    elif "label_ids" in first and first["label_ids"] is not None:
         if isinstance(first["label_ids"], torch.Tensor):
             batch["labels"] = torch.stack([f["label_ids"] for f in features])
         else:

{"golden_diff": "diff --git a/src/transformers/data/data_collator.py b/src/transformers/data/data_collator.py\n--- a/src/transformers/data/data_collator.py\n+++ b/src/transformers/data/data_collator.py\n@@ -42,10 +42,10 @@\n     # Special handling for labels.\n     # Ensure that tensor is created with the correct type\n     # (it should be automatically the case, but let's make sure of it.)\n-    if \"label\" in first:\n+    if \"label\" in first and first[\"label\"] is not None:\n         dtype = torch.long if type(first[\"label\"]) is int else torch.float\n         batch[\"labels\"] = torch.tensor([f[\"label\"] for f in features], dtype=dtype)\n-    elif \"label_ids\" in first:\n+    elif \"label_ids\" in first and first[\"label_ids\"] is not None:\n         if isinstance(first[\"label_ids\"], torch.Tensor):\n             batch[\"labels\"] = torch.stack([f[\"label_ids\"] for f in features])\n         else:\n", "issue": "data_collator.py does not allow NoneType labels for test set predictions on Glue\n# \ud83d\udc1b Bug\r\n\r\n## Information\r\n\r\nModel I am using (Bert, XLNet ...): Distilbert\r\n\r\nLanguage I am using the model on (English, Chinese ...): English\r\n\r\nThe problem arises when using:\r\n* [x] the official example scripts: (give details below) \r\n* [ ] my own modified scripts: (give details below)\r\n\r\nThe tasks I am working on is:\r\n* [x] an official GLUE/SQUaD task: (give the name)\r\n* [x] my own task or dataset: (give details below)\r\n\r\n\r\n## To reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\n1. Open the  example Colab for text-classification from here https://huggingface.co/transformers/examples.html\r\n2. Try to run the prediction function found in run_glue.py to predict on the official Glue test set. \r\n3. The error is as shown below. \r\n\r\nEarlier, this has worked with the exact same program, ever since an update recently, this error shows up. \r\n\r\n`TypeError                                 Traceback (most recent call last)\r\n<ipython-input-16-9eecdd4d48b1> in <module>()\r\n      2 output_mode = \"classification\"\r\n      3 \r\n----> 4 predictions = trainer.predict(test_dataset=test_dataset).predictions\r\n      5 if output_mode == \"classification\":\r\n      6     predictions = np.argmax(predictions, axis=1)\r\n\r\n7 frames\r\n/usr/local/lib/python3.6/dist-packages/transformers/data/data_collator.py in default_data_collator(features)\r\n     45     if \"label\" in first:\r\n     46         dtype = torch.long if type(first[\"label\"]) is int else torch.float\r\n---> 47         batch[\"labels\"] = torch.tensor([f[\"label\"] for f in features], dtype=dtype)\r\n     48     elif \"label_ids\" in first:\r\n     49         if isinstance(first[\"label_ids\"], torch.Tensor):\r\n\r\nTypeError: must be real number, not NoneType`\r\n\r\n<!-- If you have code snippets, error messages, stack traces please provide them here as well.\r\n     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting\r\n     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->\r\n\r\n## Expected behavior\r\nThe error can be seen in the Colab notebook here https://colab.research.google.com/drive/1H_92qdsOOql2hS210qNrfMEEMRcAoHD_?usp=sharing\r\n\r\n<!-- A clear and concise description of what you would expect to happen. -->\r\n\r\n## Environment info\r\n<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.\r\n     Don't forget to fill out the missing fields in that output! -->\r\n     \r\n- `transformers` version:\r\n- Platform: Colab\r\n- Python version: NA\r\n- PyTorch version (GPU?):\r\n- Tensorflow version (GPU?):\r\n- Using GPU in script?: Yes\r\n- Using distributed or parallel set-up in script?:\r\n\ndata_collator.py does not allow NoneType labels for test set predictions on Glue\n# \ud83d\udc1b Bug\r\n\r\n## Information\r\n\r\nModel I am using (Bert, XLNet ...): Distilbert\r\n\r\nLanguage I am using the model on (English, Chinese ...): English\r\n\r\nThe problem arises when using:\r\n* [x] the official example scripts: (give details below) \r\n* [ ] my own modified scripts: (give details below)\r\n\r\nThe tasks I am working on is:\r\n* [x] an official GLUE/SQUaD task: (give the name)\r\n* [x] my own task or dataset: (give details below)\r\n\r\n\r\n## To reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\n1. Open the  example Colab for text-classification from here https://huggingface.co/transformers/examples.html\r\n2. Try to run the prediction function found in run_glue.py to predict on the official Glue test set. \r\n3. The error is as shown below. \r\n\r\nEarlier, this has worked with the exact same program, ever since an update recently, this error shows up. \r\n\r\n`TypeError                                 Traceback (most recent call last)\r\n<ipython-input-16-9eecdd4d48b1> in <module>()\r\n      2 output_mode = \"classification\"\r\n      3 \r\n----> 4 predictions = trainer.predict(test_dataset=test_dataset).predictions\r\n      5 if output_mode == \"classification\":\r\n      6     predictions = np.argmax(predictions, axis=1)\r\n\r\n7 frames\r\n/usr/local/lib/python3.6/dist-packages/transformers/data/data_collator.py in default_data_collator(features)\r\n     45     if \"label\" in first:\r\n     46         dtype = torch.long if type(first[\"label\"]) is int else torch.float\r\n---> 47         batch[\"labels\"] = torch.tensor([f[\"label\"] for f in features], dtype=dtype)\r\n     48     elif \"label_ids\" in first:\r\n     49         if isinstance(first[\"label_ids\"], torch.Tensor):\r\n\r\nTypeError: must be real number, not NoneType`\r\n\r\n<!-- If you have code snippets, error messages, stack traces please provide them here as well.\r\n     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting\r\n     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->\r\n\r\n## Expected behavior\r\nThe error can be seen in the Colab notebook here https://colab.research.google.com/drive/1H_92qdsOOql2hS210qNrfMEEMRcAoHD_?usp=sharing\r\n\r\n<!-- A clear and concise description of what you would expect to happen. -->\r\n\r\n## Environment info\r\n<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.\r\n     Don't forget to fill out the missing fields in that output! -->\r\n     \r\n- `transformers` version:\r\n- Platform: Colab\r\n- Python version: NA\r\n- PyTorch version (GPU?):\r\n- Tensorflow version (GPU?):\r\n- Using GPU in script?: Yes\r\n- Using distributed or parallel set-up in script?:\r\n\n", "before_files": [{"content": "from dataclasses import dataclass\nfrom typing import Any, Callable, Dict, List, NewType, Tuple\n\nimport torch\nfrom torch.nn.utils.rnn import pad_sequence\n\nfrom ..tokenization_utils import PreTrainedTokenizer\n\n\nInputDataClass = NewType(\"InputDataClass\", Any)\n\n\"\"\"\nA DataCollator is a function that takes a list of samples from a Dataset\nand collate them into a batch, as a dictionary of Tensors.\n\"\"\"\nDataCollator = NewType(\"DataCollator\", Callable[[List[InputDataClass]], Dict[str, torch.Tensor]])\n\n\ndef default_data_collator(features: List[InputDataClass]) -> Dict[str, torch.Tensor]:\n    \"\"\"\n    Very simple data collator that:\n    - simply collates batches of dict-like objects\n    - Performs special handling for potential keys named:\n        - `label`: handles a single value (int or float) per object\n        - `label_ids`: handles a list of values per object\n    - does not do any additional preprocessing\n\n    i.e., Property names of the input object will be used as corresponding inputs to the model.\n    See glue and ner for example of how it's useful.\n    \"\"\"\n\n    # In this function we'll make the assumption that all `features` in the batch\n    # have the same attributes.\n    # So we will look at the first element as a proxy for what attributes exist\n    # on the whole batch.\n    if not isinstance(features[0], dict):\n        features = [vars(f) for f in features]\n\n    first = features[0]\n    batch = {}\n\n    # Special handling for labels.\n    # Ensure that tensor is created with the correct type\n    # (it should be automatically the case, but let's make sure of it.)\n    if \"label\" in first:\n        dtype = torch.long if type(first[\"label\"]) is int else torch.float\n        batch[\"labels\"] = torch.tensor([f[\"label\"] for f in features], dtype=dtype)\n    elif \"label_ids\" in first:\n        if isinstance(first[\"label_ids\"], torch.Tensor):\n            batch[\"labels\"] = torch.stack([f[\"label_ids\"] for f in features])\n        else:\n            dtype = torch.long if type(first[\"label_ids\"][0]) is int else torch.float\n            batch[\"labels\"] = torch.tensor([f[\"label_ids\"] for f in features], dtype=dtype)\n\n    # Handling of all other possible keys.\n    # Again, we will use the first element to figure out which key/values are not None for this model.\n    for k, v in first.items():\n        if k not in (\"label\", \"label_ids\") and v is not None and not isinstance(v, str):\n            if isinstance(v, torch.Tensor):\n                batch[k] = torch.stack([f[k] for f in features])\n            else:\n                batch[k] = torch.tensor([f[k] for f in features], dtype=torch.long)\n\n    return batch\n\n\n@dataclass\nclass DataCollatorForLanguageModeling:\n    \"\"\"\n    Data collator used for language modeling.\n    - collates batches of tensors, honoring their tokenizer's pad_token\n    - preprocesses batches for masked language modeling\n    \"\"\"\n\n    tokenizer: PreTrainedTokenizer\n    mlm: bool = True\n    mlm_probability: float = 0.15\n\n    def __call__(self, examples: List[torch.Tensor]) -> Dict[str, torch.Tensor]:\n        batch = self._tensorize_batch(examples)\n        if self.mlm:\n            inputs, labels = self.mask_tokens(batch)\n            return {\"input_ids\": inputs, \"labels\": labels}\n        else:\n            return {\"input_ids\": batch, \"labels\": batch}\n\n    def _tensorize_batch(self, examples: List[torch.Tensor]) -> torch.Tensor:\n        length_of_first = examples[0].size(0)\n        are_tensors_same_length = all(x.size(0) == length_of_first for x in examples)\n        if are_tensors_same_length:\n            return torch.stack(examples, dim=0)\n        else:\n            if self.tokenizer._pad_token is None:\n                raise ValueError(\n                    \"You are attempting to pad samples but the tokenizer you are using\"\n                    f\" ({self.tokenizer.__class__.__name__}) does not have one.\"\n                )\n            return pad_sequence(examples, batch_first=True, padding_value=self.tokenizer.pad_token_id)\n\n    def mask_tokens(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.\n        \"\"\"\n\n        if self.tokenizer.mask_token is None:\n            raise ValueError(\n                \"This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer.\"\n            )\n\n        labels = inputs.clone()\n        # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)\n        probability_matrix = torch.full(labels.shape, self.mlm_probability)\n        special_tokens_mask = [\n            self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()\n        ]\n        probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)\n        if self.tokenizer._pad_token is not None:\n            padding_mask = labels.eq(self.tokenizer.pad_token_id)\n            probability_matrix.masked_fill_(padding_mask, value=0.0)\n        masked_indices = torch.bernoulli(probability_matrix).bool()\n        labels[~masked_indices] = -100  # We only compute loss on masked tokens\n\n        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])\n        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices\n        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)\n\n        # 10% of the time, we replace masked input tokens with random word\n        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced\n        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)\n        inputs[indices_random] = random_words[indices_random]\n\n        # The rest of the time (10% of the time) we keep the masked input tokens unchanged\n        return inputs, labels\n", "path": "src/transformers/data/data_collator.py"}], "after_files": [{"content": "from dataclasses import dataclass\nfrom typing import Any, Callable, Dict, List, NewType, Tuple\n\nimport torch\nfrom torch.nn.utils.rnn import pad_sequence\n\nfrom ..tokenization_utils import PreTrainedTokenizer\n\n\nInputDataClass = NewType(\"InputDataClass\", Any)\n\n\"\"\"\nA DataCollator is a function that takes a list of samples from a Dataset\nand collate them into a batch, as a dictionary of Tensors.\n\"\"\"\nDataCollator = NewType(\"DataCollator\", Callable[[List[InputDataClass]], Dict[str, torch.Tensor]])\n\n\ndef default_data_collator(features: List[InputDataClass]) -> Dict[str, torch.Tensor]:\n    \"\"\"\n    Very simple data collator that:\n    - simply collates batches of dict-like objects\n    - Performs special handling for potential keys named:\n        - `label`: handles a single value (int or float) per object\n        - `label_ids`: handles a list of values per object\n    - does not do any additional preprocessing\n\n    i.e., Property names of the input object will be used as corresponding inputs to the model.\n    See glue and ner for example of how it's useful.\n    \"\"\"\n\n    # In this function we'll make the assumption that all `features` in the batch\n    # have the same attributes.\n    # So we will look at the first element as a proxy for what attributes exist\n    # on the whole batch.\n    if not isinstance(features[0], dict):\n        features = [vars(f) for f in features]\n\n    first = features[0]\n    batch = {}\n\n    # Special handling for labels.\n    # Ensure that tensor is created with the correct type\n    # (it should be automatically the case, but let's make sure of it.)\n    if \"label\" in first and first[\"label\"] is not None:\n        dtype = torch.long if type(first[\"label\"]) is int else torch.float\n        batch[\"labels\"] = torch.tensor([f[\"label\"] for f in features], dtype=dtype)\n    elif \"label_ids\" in first and first[\"label_ids\"] is not None:\n        if isinstance(first[\"label_ids\"], torch.Tensor):\n            batch[\"labels\"] = torch.stack([f[\"label_ids\"] for f in features])\n        else:\n            dtype = torch.long if type(first[\"label_ids\"][0]) is int else torch.float\n            batch[\"labels\"] = torch.tensor([f[\"label_ids\"] for f in features], dtype=dtype)\n\n    # Handling of all other possible keys.\n    # Again, we will use the first element to figure out which key/values are not None for this model.\n    for k, v in first.items():\n        if k not in (\"label\", \"label_ids\") and v is not None and not isinstance(v, str):\n            if isinstance(v, torch.Tensor):\n                batch[k] = torch.stack([f[k] for f in features])\n            else:\n                batch[k] = torch.tensor([f[k] for f in features], dtype=torch.long)\n\n    return batch\n\n\n@dataclass\nclass DataCollatorForLanguageModeling:\n    \"\"\"\n    Data collator used for language modeling.\n    - collates batches of tensors, honoring their tokenizer's pad_token\n    - preprocesses batches for masked language modeling\n    \"\"\"\n\n    tokenizer: PreTrainedTokenizer\n    mlm: bool = True\n    mlm_probability: float = 0.15\n\n    def __call__(self, examples: List[torch.Tensor]) -> Dict[str, torch.Tensor]:\n        batch = self._tensorize_batch(examples)\n        if self.mlm:\n            inputs, labels = self.mask_tokens(batch)\n            return {\"input_ids\": inputs, \"labels\": labels}\n        else:\n            return {\"input_ids\": batch, \"labels\": batch}\n\n    def _tensorize_batch(self, examples: List[torch.Tensor]) -> torch.Tensor:\n        length_of_first = examples[0].size(0)\n        are_tensors_same_length = all(x.size(0) == length_of_first for x in examples)\n        if are_tensors_same_length:\n            return torch.stack(examples, dim=0)\n        else:\n            if self.tokenizer._pad_token is None:\n                raise ValueError(\n                    \"You are attempting to pad samples but the tokenizer you are using\"\n                    f\" ({self.tokenizer.__class__.__name__}) does not have one.\"\n                )\n            return pad_sequence(examples, batch_first=True, padding_value=self.tokenizer.pad_token_id)\n\n    def mask_tokens(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.\n        \"\"\"\n\n        if self.tokenizer.mask_token is None:\n            raise ValueError(\n                \"This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer.\"\n            )\n\n        labels = inputs.clone()\n        # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)\n        probability_matrix = torch.full(labels.shape, self.mlm_probability)\n        special_tokens_mask = [\n            self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()\n        ]\n        probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)\n        if self.tokenizer._pad_token is not None:\n            padding_mask = labels.eq(self.tokenizer.pad_token_id)\n            probability_matrix.masked_fill_(padding_mask, value=0.0)\n        masked_indices = torch.bernoulli(probability_matrix).bool()\n        labels[~masked_indices] = -100  # We only compute loss on masked tokens\n\n        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])\n        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices\n        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)\n\n        # 10% of the time, we replace masked input tokens with random word\n        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced\n        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)\n        inputs[indices_random] = random_words[indices_random]\n\n        # The rest of the time (10% of the time) we keep the masked input tokens unchanged\n        return inputs, labels\n", "path": "src/transformers/data/data_collator.py"}]}