Repository: huggingface/transformers

Issue:
[examples] should all examples support the predict stage?
This is part of the ongoing effort to sync the example scripts.

In https://github.com/huggingface/transformers/issues/10437#issuecomment-789090858 it was flagged that some scripts have test/predict, whereas others don't.

Should we:
A. have all scripts have train/eval/predict
B. only have predict where it's desired

@sgugger, @patil-suraj, @LysandreJik 


Discussion:
I think we should have it on all scripts except the language-modeling ones -> it doesn't make much sense there.
@bhadreshpsavani, please feel free to make this part of your project or not. Please do not feel obliged as you can see a small need quickly expands into a much bigger one.

If not, then completing https://github.com/huggingface/transformers/issues/10437 would be a fantastic contribution. And I will then label it to be open to anybody else who would like to contribute.

Thank you!
Sure @stas00,
I will take this. Actually, I would love to contribute more. I really enjoy contributing to this community.
Excellent! If you run into any puzzles please don't hesitate to ask.

Thank you!
Hi @stas00,
In the QA examples for the test dataset, should we keep The preprocessing same as evaluate dataset preprocessing?
Do we need to apply any additional post-processing as well at the end?
Yes, that would be needed.
Hi, @stas00 and @sgugger, 

Adding predict function for the `run_qa` example is slightly complicated. 
In the eval section itself, we are generating two files `predictions.json` and `nbest_predictions.json` using  `postprocess_qa_predictions` 
 from `utils_qa`. In Predict function also the same file will be generated and override the same files which will not be very good behavior.

In the predict function currently `trainer.predict()` calculates metrics but metrics calculation might not require in case of `predict()` right? 

This issue might take a longer time for me to complete may be few weeks, is it fine?
> In Predict function also the same file will be generated and override the same files which will not be very good behavior.

This is simple: add `eval_` prefix for eval outputs and `test_` for predict.

related: eventually we will rename the latter prefix, but for now this is the convention used in all examples. Please see: https://github.com/huggingface/transformers/issues/10165

> In the predict function currently trainer.predict() calculates metrics but metrics calculation might not require in case of predict() right?

We always want the metrics for each stage, since now we report speed and memory usage, so for the quality metrics see what makes sense to report there - typically similar to eval.

> This issue might take a longer time for me to complete may be few weeks, is it fine?

Yes, of course, thank you for giving us heads up and the clarity of your needs, @bhadreshpsavani 
Hi @stas00,

When I run the evaluation on [example for question answering](https://github.com/bhadreshpsavani/UnderstandingNLP/blob/master/Check_Pretrained_Model_on_SQUAD2.ipynb),

I was getting below error,
```
Traceback (most recent call last):
  File "./transformers/examples/question-answering/run_qa.py", line 546, in <module>
    main()
  File "./transformers/examples/question-answering/run_qa.py", line 531, in main
    metrics = trainer.evaluate()
  File "/content/transformers/examples/question-answering/trainer_qa.py", line 63, in evaluate
    metrics = self.compute_metrics(eval_preds)
  File "./transformers/examples/question-answering/run_qa.py", line 492, in compute_metrics
    return metric.compute(predictions=p.predictions, references=p.label_ids)
  File "/usr/local/lib/python3.7/dist-packages/datasets/metric.py", line 403, in compute
    output = self._compute(predictions=predictions, references=references, **kwargs)
  File "/root/.cache/huggingface/modules/datasets_modules/metrics/squad/c0855591f1a2c2af8b7949e3146b9c86a6b7f536b4154019b03472639d310181/squad.py", line 109, in _compute
    score = evaluate(dataset=dataset, predictions=pred_dict)
  File "/root/.cache/huggingface/modules/datasets_modules/metrics/squad/c0855591f1a2c2af8b7949e3146b9c86a6b7f536b4154019b03472639d310181/evaluate.py", line 68, in evaluate
    exact_match += metric_max_over_ground_truths(exact_match_score, prediction, ground_truths)
  File "/root/.cache/huggingface/modules/datasets_modules/metrics/squad/c0855591f1a2c2af8b7949e3146b9c86a6b7f536b4154019b03472639d310181/evaluate.py", line 53, in metric_max_over_ground_truths
    return max(scores_for_ground_truths)
ValueError: max() arg is an empty sequence
```
While calculating metrics at the end,
I could not found anything to resolve this!



I am able to reproduce it - will have a look shortly. Thank you for the notebook - makes it super-easy to reproduce!

BTW, I recommend you use `--max_val_samples 10` or similar to make your testing much faster ;)
OK, so we have a case of a broken dataset and the metrics evaluation code that doesn't check if the input data is valid.

While ideally the eval code should be robust against an occasional corrupted input, the dataset should not have broken entries, observe the kind of data it evaluates against (`ground_truths`) 

```
[{'answers': {'answer_start': [159, 159, 159, 159], 'text': ['France', 'France', 'France', 'France']}, 'id': '56ddde6b9a695914005b9628'},
{'answers': {'answer_start': [94, 87, 94, 94], 'text': ['10th and 11th centuries', 'in the 10th and 11th centuries', '10th and 11th centuries', '10th and 11th centuries']}, 'id': '56ddde6b9a695914005b9629'},

 {'answers': {'answer_start': [], 'text': []}, 'id': '5ad39d53604f3c001a3fe8d1'},
 {'answers': {'answer_start': [], 'text': []}, 'id': '5ad39d53604f3c001a3fe8d2'},
 {'answers': {'answer_start': [], 'text': []}, 'id': '5ad39d53604f3c001a3fe8d3'},
 {'answers': {'answer_start': [], 'text': []}, 'id': '5ad39d53604f3c001a3fe8d4'}, 
```
- and it has a ton of those - so this is a big problem on the dataset-level.

Having a quick look at the viewer:
https://huggingface.co/datasets/viewer/?dataset=squad_v2
I don't immediately see missing answers (scroll horizontally to the right to see it), but it doesn't mean anything.

Which most likely means that either the dataset conversion from the original to `datasets` format is borked, or there is some odd bug in the dataloader. But most likely it's the former. that is I'd debug this first and validate that it generates the answers correctly:
https://github.com/huggingface/datasets/blob/b51cb81e736b86103089a584daa6e43db3c88bb5/datasets/squad_v2/squad_v2.py#L101


if so then proceed to where they are loaded in the script and debug there, and so on - until you find a place where some of the records disappear in what it appears more than half of samples. There you will find the bug.

Please let me know if you'd like to try to investigate this and whether my suggestion at how to potentially approach this is clear. If it sounds too complicated please don't hesitate to say so and we will find another way to resolve this. Either way works.

Meanwhile you can also try with a different model/dataset pair and see if it works there, which would also help isolate the problem. (if another dataset works then we know for sure the issue lies with `squad_v2`.
Hi @stas00,
For `squad_v1` = `squad` it is working fine.

But when i run the following script,
```
!python ./transformers/examples/question-answering/run_qa.py \
--model_name_or_path distilbert-base-uncased \
--train_file ./transformers/tests/fixtures/tests_samples/SQUAD/sample.json \
--validation_file ./transformers/tests/fixtures/tests_samples/SQUAD/sample.json \
--do_eval \
--max_val_sample 10 \
--per_device_train_batch_size 12 \
--learning_rate 3e-5 \
--num_train_epochs 2 \
--max_seq_length 384 \
--doc_stride 128 \
--output_dir /tmp/debug_squad/
```
which again data, I got the same error

In the earlier error logs it is giving like this for this issue,
```
File "/root/.cache/huggingface/modules/datasets_modules/metrics/squad/c0855591f1a2c2af8b7949e3146b9c86a6b7f536b4154019b03472639d310181/squad.py", line 109, in _compute
```
But for the squad_v2 dataset, it should be like this, 
```
"/root/.cache/huggingface/modules/datasets_modules/metrics/squad_v2/c0855591f1a2c2af8b7949e3146b9c86a6b7f536b4154019b03472639d310181/squad_v2.py", line 109, in _compute
```
i mean it is picking wrong metrics somehow
OK, I found this in `run_qa.py`:
```
    version_2_with_negative: bool = field(
        default=False, metadata={"help": "If true, some of the examples do not have an answer."}
    )
```

That is add: `--version_2_with_negative` to the cl args and it does:
```
 metric = load_metric("squad_v2" if data_args.version_2_with_negative else "squad")
```

Except it's broken too:
```
examples/question-answering/run_qa.py --model_name_or_path ktrapeznikov/albert-xlarge-v2-squad-v2 --dataset_name squad_v2 --do_eval --per_device_train_batch_size 12 --learning_rate 3e-5 --num_train_epochs 2 --max_seq_length 384 --doc_stride 128 --output_dir /tmp/debug_squad/ --max_val_samples 10 --version_2_with_negative
Traceback (most recent call last):                                                                                         | 0/11873 [00:00<?, ?it/s]
  File "examples/question-answering/run_qa.py", line 553, in <module>
    main()
  File "examples/question-answering/run_qa.py", line 538, in main
    metrics = trainer.evaluate()
  File "/mnt/nvme1/code/huggingface/transformers-master/examples/question-answering/trainer_qa.py", line 62, in evaluate
    eval_preds = self.post_process_function(eval_examples, eval_dataset, output.predictions)
  File "examples/question-answering/run_qa.py", line 475, in post_processing_function
    predictions = postprocess_qa_predictions(
  File "/mnt/nvme1/code/huggingface/transformers-master/examples/question-answering/utils_qa.py", line 159, in postprocess_qa_predictions
    null_score = min_null_prediction["score"]
TypeError: 'NoneType' object is not subscriptable
```
this fails with the same error:
```
python ./examples/question-answering/run_qa.py --model_name_or_path distilbert-base-uncased --train_file tests/fixtures/tests_samples/SQUAD/sample.json --validation_file tests/fixtures/tests_samples/SQUAD/sample.json --do_eval --max_val_sample 10 --per_device_train_batch_size 12 --learning_rate 3e-5 --num_train_epochs 2 --max_seq_length 384 --doc_stride 128 --output_dir /tmp/debug_squad/ --version_2_with_negative
```

So it looks like this feature is half-baked, and ideally we should have a test for it.

Let me know if you're ok to investigate this new issue.

Plus ideally `run_qa.py` should signal the user that if the dataset has missing answers it should bail and recommend using `--version_2_with_negative`.

(and the choice of this flag's name could probably be improved too to be more intuitive at what it does, but that's another story)

Sure @stas00,

I can investigate this issue.

One more thing currently `run_qa.py` only supports `squad_v1` and `squad_v2` datasets completely, right? Shouldn't it support all the different question-answering datasets of the reading comprehension task? I mean preprocessing and postprocessing might be different for all the tasks but the concept is the same for all the tasks. Please correct me if I am wrong.

I tried narrativeQA it was not supporting that one.
Awesome!

> One more thing currently run_qa.py only supports squad_v1 and squad_v2 datasets completely, right? Shouldn't it support all the different question-answering datasets of the reading comprehension task? I mean preprocessing and postprocessing might be different for all the tasks but the concept is the same for all the tasks. Please correct me if I am wrong.

Honestly, I have no idea as I didn't write it. 

@sgugger, @patrickvonplaten, do we want `run_qa.py` to support more than `squad_v1` and `squad_v2` datasets? Thank you!
Like all other examples, the script is given as just that, an example. As said in the main README under "Why shouldn't I use Transformers?":
```
While we strive to present as many use cases as possible, the scripts in our examples folder are just that: examples. It is expected that they won't work out-of-the box on your specific problem and that you will be required to change a few lines of code to adapt them to your needs.
```

So like all other example scripts, the `run_qa` script will support any dataset that is structured the same way as the original dataset that was used with it (squad) but if the user wants the script to work on another dataset structured differently they will need to tweak it to their needs.
Hi @stas00 and @sgugger,

I figured the cause for that squad2 issue while using `max_sample_*` arguments, I fixed it locally. The cause of the error was in the below line
https://github.com/huggingface/transformers/blob/5f19c07a704eca4db376b56f950b729dcaa73039/examples/question-answering/run_qa.py#L493

It uses entire data for reference while it should use only `max_sample_*` data only.

I tried to fix it by below code,
```python
if training_args.do_eval:
  if "validation" not in datasets:
      raise ValueError("--do_eval requires a validation dataset")
  eval_examples = datasets["validation"]
  if data_args.max_val_samples is not None:
      # We will select sample from whole data
      eval_examples = eval_examples.select(range(data_args.max_val_samples))
  # Validation Feature Creation
  eval_dataset = eval_examples.map(
      prepare_validation_features,
      batched=True,
      num_proc=data_args.preprocessing_num_workers,
      remove_columns=column_names,
      load_from_cache_file=not data_args.overwrite_cache,
  )
  if data_args.max_val_samples is not None:
      # During Feature creation dataset samples might increase, we will select required samples again
      eval_dataset = eval_dataset.select(range(data_args.max_val_samples))
```
and I used `eval_examples ` instead `datasets["validation"]` in the above-mentioned line, This is working fine for data I tested, but as we know after applying `prepare_validation_features`, `eval_dataset ` and `eval_examples` might not have the same data (length will be same but because of sliding window example and feature might not be representing the same item)

For `max_val_samples=10` these changes are working fine so I added predict/test method, It seems to be working fine.

https://github.com/huggingface/transformers/blob/5f19c07a704eca4db376b56f950b729dcaa73039/examples/question-answering/utils_qa.py#L214-L226

i modified in below lines,
```python
  prediction_file = os.path.join(
      output_dir, "predictions.json" if prefix is None else f"{prefix}_predictions.json"
  )
  nbest_file = os.path.join(
      output_dir, "nbest_predictions.json" if prefix is None else f"{prefix}_nbest_predictions.json"
  )
  if version_2_with_negative:
      null_odds_file = os.path.join(
          output_dir, "null_odds.json" if prefix is None else f"{prefix}_null_odds_{prefix}.json"
```
because while passing prefix and running it I was getting error like `string can't have json attribute`
it will save test and eval files like `test_predictions.json` and `eval_predictions.json`
 
Good job getting to the root of the issue, I hadn't thought of that when you added the `max_sample_xxx` but this task uses a subclass of the main `Trainer` that does require the original `eval_examples`.

The fix you propose appears good to me and you should definitely make a PR with it as soon as you can :-)

For the predicting stage, note that the subclass of the `Trainer` will require a test_dataset and test_examples to work (basically to interpret the predictions of the model as spans of the original texts, the `Trainer` needs the original texts). I do think adding a `--do_predict` to `run_qa` is going to be a bit complex so should be treated separately so my advise would be to:

1. make a PR with the fix for evaluation in run_qa/run_qa_beam_search when `max_val_samples` is passed
2. make a PR to add predict in all but run_qa/run_qa_beam_search scripts (when it makes sense of course)
3. make a PR to add predict in run_qa/run_qa_beam_search

Let me know if that makes sense to you and if you need any help along the way (or don't want to do one of those steps yourself).
Really awesome, @bhadreshpsavani! Glad you were able to find out the cause!
I gone through the traceback(stack trace) by running example with different stage combination and figured the root. Since I wrote earlier code about max_sample and the error was related I was able to find it!

Thanks @stas00 and @sgugger for your constant guidance, Now opensource contribution don't seems very difficult like earlier!

Now I need to figure out for two more examples `run_swag.py` and `run_xlni.py`, all the other examples has predict if we ignore language modeling examples
Hi @sgugger,

For `run_swag.py` and `run_xlni.py` changes are still remaining
Hi @stas00,

I was working to update `run_swag.py` and I got logs like for test metrics, 
```
INFO|trainer_pt_utils.py:656] 2021-03-21 20:56:40,939 >> ***** test metrics *****
[INFO|trainer_pt_utils.py:661] 2021-03-21 20:56:40,939 >>   eval_accuracy             =    1.0
[INFO|trainer_pt_utils.py:661] 2021-03-21 20:56:40,939 >>   eval_loss                 = 0.2582
[INFO|trainer_pt_utils.py:661] 2021-03-21 20:56:40,939 >>   eval_runtime              = 4.1585
[INFO|trainer_pt_utils.py:661] 2021-03-21 20:56:40,941 >>   eval_samples_per_second   =  2.405
[INFO|trainer_pt_utils.py:661] 2021-03-21 20:56:40,941 >>   test_mem_cpu_alloc_delta  =    0MB
[INFO|trainer_pt_utils.py:661] 2021-03-21 20:56:40,941 >>   test_mem_cpu_peaked_delta =    0MB
[INFO|trainer_pt_utils.py:661] 2021-03-21 20:56:40,941 >>   test_samples              =     10
```
I think we need to write `trainer_swag.py` because even predictions files need to be saved like `test_prediction.json` / `eval_prediction.json`.
If I understood your correctly you suggest to write `MultipleChoiceTrainer` subclass.

If so unlike the `question-answering` folder that has 2 scripts, `multiple-choice` has only one script so one way is to add it directly in `run_swag.py` or if there is a convention then as you suggest as `trainer_mc.py`. I think follow you instinct and then we can decide at PR point whether to have it in a separate file. 
Hi, @stas00 @sgugger,

The changes are ready for both examples, `run_xlni.py` works perfectly.

`run_swag.py` has an issue after adding changes of predict,
```
python ./examples/multiple-choice/run_swag.py --model_name_or_path distilbert-base-uncased --do_train --do_eval --do_predict --max_train_samples 5 --max_val_samples 5 --max_test_samples 5 --learning_rate 5e-5 --num_train_epochs 3 --output_dir D:/tmp/swag_base --per_gpu_eval_batch_size=16 --per_device_train_batch_size=16 --overwrite_output
```
it gives the following error while prediction,
```
Traceback (most recent call last):
  File "./examples/multiple-choice/run_swag.py", line 481, in <module>
    main()
  File "./examples/multiple-choice/run_swag.py", line 466, in main
    predictions, labels, metrics = trainer.predict(test_dataset=test_dataset)
  File "d:\transformers\src\transformers\trainer.py", line 1762, in predict
    output = self.prediction_loop(
  File "d:\transformers\src\transformers\trainer.py", line 1829, in prediction_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "d:\transformers\src\transformers\trainer.py", line 1943, in prediction_step
    loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
  File "d:\transformers\src\transformers\trainer.py", line 1504, in compute_loss
    outputs = model(**inputs)
  File "C:\Users\Bhadr\Miniconda3\envs\trans-env2\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "d:\transformers\src\transformers\models\distilbert\modeling_distilbert.py", line 929, in forward
    loss = loss_fct(reshaped_logits, labels)
  File "C:\Users\Bhadr\Miniconda3\envs\trans-env2\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "C:\Users\Bhadr\Miniconda3\envs\trans-env2\lib\site-packages\torch\nn\modules\loss.py", line 1047, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "C:\Users\Bhadr\Miniconda3\envs\trans-env2\lib\site-packages\torch\nn\functional.py", line 2690, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "C:\Users\Bhadr\Miniconda3\envs\trans-env2\lib\site-packages\torch\nn\functional.py", line 2385, in nll_loss
    ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
IndexError: Target -1 is out of bounds.
```
The problem is with test dataset,
```python
dataset = load_dataset("swag", "regular")
pd.Series(dataset['test']['label']).value_counts()
"""
Output:
-1    20005
dtype: int64
"""
```
It should have label In range(0, 3) like eval dataset not -1,
```python
pd.Series(dataset['validation']['label']).value_counts()
"""
Output:
2    5038
1    5029
3    5006
0    4933
dtype: int64
"""
```
I am not sure how can we fix this since it's with the dataset and this example is for only swag. when we pass the sample data present in `tests/fixtures/tests_samples/swag/sample.json` it works fine since it has label 0.

There is another small issue that needs to be fixed as well in the `trainer.py`
https://github.com/huggingface/transformers/blob/a8d4d6776dd8a759324d0f57c60e8a738e7977a4/src/transformers/trainer.py#L1724-L1726
Because of this current test files are being saved with eval prefix, the issue that I mentioned in my previous comment.
It should be `metric_key_prefix: str = "test"`.

Please correct me if I am wrong.

I'm not sure we need the predict stage for `run_swag`, especially if the test dataset is "broken".
Filing an Issue with https://github.com/huggingface/datasets/issues and meanwhile using the local test file for the 
README.md?

> There is another small issue that needs to be fixed as well in the trainer.py

Indeed. In the default value and also the docstring - I'd say please make a separate PR for that as it's a standalone bug.

And then we will eventually take care of https://github.com/huggingface/transformers/issues/10165 but for now let's use `test`.
Hi @stas00,
The original [dataset](https://github.com/rowanz/swagaf/tree/master/data) doesn't have a label in the test dataset so it's using `-1` as a label. 

`meanwhile using the local test file for the
README.md` - I don't understand this part, can you please tell me a bit more about this?

I will create a PR with `run_xlni.py`, Please let me know if we need to add changes for `run_swag` (On sample test data predict is working fine)

Resolving patch:
diff --git a/examples/text-classification/run_xnli.py b/examples/text-classification/run_xnli.py
--- a/examples/text-classification/run_xnli.py
+++ b/examples/text-classification/run_xnli.py
@@ -207,14 +207,22 @@ def main():
     # In distributed training, the load_dataset function guarantees that only one local process can concurrently
     # download the dataset.
     # Downloading and loading xnli dataset from the hub.
-    if model_args.train_language is None:
-        train_dataset = load_dataset("xnli", model_args.language, split="train")
-    else:
-        train_dataset = load_dataset("xnli", model_args.train_language, split="train")
+    if training_args.do_train:
+        if model_args.train_language is None:
+            train_dataset = load_dataset("xnli", model_args.language, split="train")
+        else:
+            train_dataset = load_dataset("xnli", model_args.train_language, split="train")
+        label_list = train_dataset.features["label"].names
+
+    if training_args.do_eval:
+        eval_dataset = load_dataset("xnli", model_args.language, split="validation")
+        label_list = eval_dataset.features["label"].names
+
+    if training_args.do_predict:
+        test_dataset = load_dataset("xnli", model_args.language, split="test")
+        label_list = test_dataset.features["label"].names
 
-    eval_dataset = load_dataset("xnli", model_args.language, split="validation")
     # Labels
-    label_list = train_dataset.features["label"].names
     num_labels = len(label_list)
 
     # Load pretrained model and tokenizer
@@ -271,6 +279,9 @@ def preprocess_function(examples):
             batched=True,
             load_from_cache_file=not data_args.overwrite_cache,
         )
+        # Log a few random samples from the training set:
+        for index in random.sample(range(len(train_dataset)), 3):
+            logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
 
     if training_args.do_eval:
         if data_args.max_val_samples is not None:
@@ -281,9 +292,14 @@ def preprocess_function(examples):
             load_from_cache_file=not data_args.overwrite_cache,
         )
 
-    # Log a few random samples from the training set:
-    for index in random.sample(range(len(train_dataset)), 3):
-        logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+    if training_args.do_predict:
+        if data_args.max_test_samples is not None:
+            test_dataset = test_dataset.select(range(data_args.max_test_samples))
+        test_dataset = test_dataset.map(
+            preprocess_function,
+            batched=True,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
 
     # Get the metric function
     metric = load_metric("xnli")
@@ -307,7 +323,7 @@ def compute_metrics(p: EvalPrediction):
     trainer = Trainer(
         model=model,
         args=training_args,
-        train_dataset=train_dataset,
+        train_dataset=train_dataset if training_args.do_train else None,
         eval_dataset=eval_dataset if training_args.do_eval else None,
         compute_metrics=compute_metrics,
         tokenizer=tokenizer,
@@ -346,6 +362,26 @@ def compute_metrics(p: EvalPrediction):
         trainer.log_metrics("eval", metrics)
         trainer.save_metrics("eval", metrics)
 
+    # Prediction
+    if training_args.do_predict:
+        logger.info("*** Predict ***")
+        predictions, labels, metrics = trainer.predict(test_dataset)
+
+        max_test_samples = data_args.max_test_samples if data_args.max_test_samples is not None else len(test_dataset)
+        metrics["test_samples"] = min(max_test_samples, len(test_dataset))
+
+        trainer.log_metrics("test", metrics)
+        trainer.save_metrics("test", metrics)
+
+        predictions = np.argmax(predictions, axis=1)
+        output_test_file = os.path.join(training_args.output_dir, "test_predictions.txt")
+        if trainer.is_world_process_zero():
+            with open(output_test_file, "w") as writer:
+                writer.write("index\tprediction\n")
+                for index, item in enumerate(predictions):
+                    item = label_list[item]
+                    writer.write(f"{index}\t{item}\n")
+
 
 if __name__ == "__main__":
     main()
diff --git a/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py b/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
--- a/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
+++ b/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
@@ -139,6 +139,10 @@ class DataTrainingArguments:
         default=None,
         metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
     )
+    test_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input test data file to predict the label on (a text file)."},
+    )
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
@@ -160,10 +164,22 @@ class DataTrainingArguments:
             "value if set."
         },
     )
+    max_test_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of test examples to this "
+            "value if set."
+        },
+    )
 
     def __post_init__(self):
-        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
-            raise ValueError("Need either a dataset name or a training/validation file.")
+        if (
+            self.dataset_name is None
+            and self.train_file is None
+            and self.validation_file is None
+            and self.test_file is None
+        ):
+            raise ValueError("Need either a dataset name or a training/validation/test file.")
         else:
             if self.train_file is not None:
                 extension = self.train_file.split(".")[-1]
@@ -171,6 +187,9 @@ def __post_init__(self):
             if self.validation_file is not None:
                 extension = self.validation_file.split(".")[-1]
                 assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
+            if self.test_file is not None:
+                extension = self.test_file.split(".")[-1]
+                assert extension in ["csv", "json", "txt"], "`test_file` should be a csv, a json or a txt file."
 
 
 def main():
@@ -238,9 +257,13 @@ def main():
         data_files = {}
         if data_args.train_file is not None:
             data_files["train"] = data_args.train_file
+            extension = data_args.train_file.split(".")[-1]
         if data_args.validation_file is not None:
             data_files["validation"] = data_args.validation_file
-        extension = data_args.train_file.split(".")[-1]
+            extension = data_args.validation_file.split(".")[-1]
+        if data_args.test_file is not None:
+            data_files["test"] = data_args.test_file
+            extension = data_args.test_file.split(".")[-1]
         if extension == "txt":
             extension = "text"
         datasets = load_dataset(extension, data_files=data_files)
@@ -326,8 +349,10 @@ def main():
     # First we tokenize all the texts.
     if training_args.do_train:
         column_names = datasets["train"].column_names
-    else:
+    elif training_args.do_eval:
         column_names = datasets["validation"].column_names
+    elif training_args.do_predict:
+        column_names = datasets["test"].column_names
     text_column_name = "text" if "text" in column_names else column_names[0]
 
     def tokenize_function(examples):
@@ -365,6 +390,22 @@ def tokenize_function(examples):
             load_from_cache_file=not data_args.overwrite_cache,
         )
 
+    if training_args.do_predict:
+        if "test" not in datasets:
+            raise ValueError("--do_predict requires a test dataset")
+        test_dataset = datasets["test"]
+        # Selecting samples from dataset
+        if data_args.max_test_samples is not None:
+            test_dataset = test_dataset.select(range(data_args.max_test_samples))
+        # tokenize test dataset
+        test_dataset = test_dataset.map(
+            tokenize_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=[text_column_name],
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
     # Data collator
     data_collator=default_data_collator if not training_args.fp16 else DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
 
@@ -420,6 +461,18 @@ def tokenize_function(examples):
         trainer.log_metrics("eval", metrics)
         trainer.save_metrics("eval", metrics)
 
+    # Prediction
+    if training_args.do_predict:
+        logger.info("*** Predict ***")
+        predictions, labels, metrics = trainer.predict(test_dataset)
+
+        max_test_samples = data_args.max_test_samples if data_args.max_test_samples is not None else len(test_dataset)
+        metrics["test_samples"] = min(max_test_samples, len(test_dataset))
+
+        trainer.log_metrics("test", metrics)
+        trainer.save_metrics("test", metrics)
+        
+        # write custom code for saving predictions according to task
 
 def _mp_fn(index):
     # For xla_spawn (TPUs)


Test patch: