Gradient checkpointing

* Try -j option

* Try other thing

* Bigger machine

* Test lower sphinx version

* Remove trailing space

.circleci/config.yml

@@ -77,16 +77,17 @@ jobs:
                       - v0.4-torch_and_tf-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,torch,testing]
             - save_cache:
                 key: v0.4-{{ checksum "setup.py" }}
                 paths:
                     - '~/.cache/pip'
-            - run: RUN_PT_TF_CROSS_TESTS=1 python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ -m is_pt_tf_cross_test --durations=0 | tee output.txt
+            - run: RUN_PT_TF_CROSS_TESTS=1 python -m pytest -n 8 --dist=loadfile -rA -s --make-reports=tests_torch_and_tf ./tests/ -m is_pt_tf_cross_test --durations=0 | tee tests_output.txt
             - store_artifacts:
-                  path: ~/transformers/output.txt
-                  destination: test_output.txt
+                  path: ~/transformers/tests_output.txt
+            - store_artifacts:
+                  path: ~/transformers/reports
+
     run_tests_torch:
         working_directory: ~/transformers
         docker:
@@ -102,13 +103,12 @@ jobs:
                       - v0.4-torch-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,torch,testing]
             - save_cache:
                   key: v0.4-torch-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: python -m pytest -n 8 --dist=loadfile -s --make_reports=tests ./tests/ | tee tests_output.txt
+            - run: python -m pytest -n 8 --dist=loadfile -s --make-reports=tests_torch ./tests/ | tee tests_output.txt
             - store_artifacts:
                   path: ~/transformers/tests_output.txt
             - store_artifacts:
@@ -129,16 +129,17 @@ jobs:
                       - v0.4-tf-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,testing]
             - save_cache:
                   key: v0.4-tf-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ | tee output.txt
+            - run: python -m pytest -n 8 --dist=loadfile -rA -s --make-reports=tests_tf ./tests/ | tee tests_output.txt
             - store_artifacts:
-               path: ~/transformers/output.txt
-               destination: test_output.txt
+                  path: ~/transformers/tests_output.txt
+            - store_artifacts:
+                  path: ~/transformers/reports
+
     run_tests_flax:
         working_directory: ~/transformers
         docker:
@@ -154,16 +155,17 @@ jobs:
                     - v0.4-flax-{{ checksum "setup.py" }}
                     - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: sudo pip install .[flax,sklearn,torch,testing]
             - save_cache:
                   key: v0.4-flax-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ | tee output.txt
+            - run: python -m pytest -n 8 --dist=loadfile -rA -s --make-reports=tests_flax ./tests/ | tee tests_output.txt
             - store_artifacts:
-                  path: ~/transformers/output.txt
-                  destination: test_output.txt
+                  path: ~/transformers/tests_output.txt
+            - store_artifacts:
+                  path: ~/transformers/reports
+
     run_tests_pipelines_torch:
         working_directory: ~/transformers
         docker:
@@ -179,16 +181,17 @@ jobs:
                       - v0.4-torch-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,torch,testing]
             - save_cache:
                   key: v0.4-torch-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: RUN_PIPELINE_TESTS=1 python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ -m is_pipeline_test | tee output.txt
+            - run: RUN_PIPELINE_TESTS=1 python -m pytest -n 8 --dist=loadfile -rA -s --make-reports=tests_pipelines_torch -m is_pipeline_test ./tests/ | tee tests_output.txt
             - store_artifacts:
-                  path: ~/transformers/output.txt
-                  destination: test_output.txt
+                  path: ~/transformers/tests_output.txt
+            - store_artifacts:
+                  path: ~/transformers/reports
+
     run_tests_pipelines_tf:
         working_directory: ~/transformers
         docker:
@@ -204,16 +207,17 @@ jobs:
                       - v0.4-tf-{{ checksum "setup.py" }}
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,testing]
             - save_cache:
                   key: v0.4-tf-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: RUN_PIPELINE_TESTS=1 python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ -m is_pipeline_test | tee output.txt
+            - run: RUN_PIPELINE_TESTS=1 python -m pytest -n 8 --dist=loadfile -rA -s --make-reports=tests_pipelines_tf ./tests/ -m is_pipeline_test | tee tests_output.txt
             - store_artifacts:
-               path: ~/transformers/output.txt
-               destination: test_output.txt
+                  path: ~/transformers/tests_output.txt
+            - store_artifacts:
+                  path: ~/transformers/reports
+
     run_tests_custom_tokenizers:
         working_directory: ~/transformers
         docker:
@@ -233,10 +237,12 @@ jobs:
                   key: v0.4-custom_tokenizers-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: python -m pytest -s ./tests/test_tokenization_bert_japanese.py | tee output.txt
+            - run: python -m pytest -s --make-reports=tests_custom_tokenizers ./tests/test_tokenization_bert_japanese.py | tee tests_output.txt
             - store_artifacts:
-                path: ~/transformers/output.txt
-                destination: test_output.txt
+                  path: ~/transformers/tests_output.txt
+            - store_artifacts:
+                  path: ~/transformers/reports
+
     run_examples_torch:
         working_directory: ~/transformers
         docker:
@@ -258,11 +264,12 @@ jobs:
                   key: v0.4-torch_examples-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: python -m pytest -n 8 --dist=loadfile -s --make_reports=examples ./examples/ | tee examples_output.txt
+            - run: python -m pytest -n 8 --dist=loadfile -s --make-reports=examples_torch ./examples/ | tee examples_output.txt
             - store_artifacts:
                   path: ~/transformers/examples_output.txt
             - store_artifacts:
                   path: ~/transformers/reports
+
     build_doc:
         working_directory: ~/transformers
         docker:
@@ -282,6 +289,7 @@ jobs:
             - run: cd docs && make html SPHINXOPTS="-W"
             - store_artifacts:
                 path: ./docs/_build
+
     deploy_doc:
         working_directory: ~/transformers
         docker:
@@ -301,6 +309,7 @@ jobs:
                   paths:
                       - '~/.cache/pip'
             - run: ./.circleci/deploy.sh
+
     check_code_quality:
         working_directory: ~/transformers
         docker:
@@ -327,6 +336,7 @@ jobs:
             - run: python utils/check_copies.py
             - run: python utils/check_dummies.py
             - run: python utils/check_repo.py
+
     check_repository_consistency:
         working_directory: ~/transformers
         docker:
@@ -357,6 +367,7 @@ jobs:
             - setup_remote_docker
             - *build_push_docker
             - *deploy_cluster
+
     cleanup-gke-jobs:
         docker:
             - image: circleci/python:3.6
@@ -366,6 +377,7 @@ jobs:
                   cluster: $GKE_CLUSTER
                   perform-login: true
             - *delete_gke_jobs
+
 workflow_filters: &workflow_filters
     filters:
         branches:

.github/workflows/self-push.yml

@@ -61,7 +61,19 @@ jobs:
           CUDA_VISIBLE_DEVICES: 0
         run: |
           source .env/bin/activate
-          python -m pytest -n 2 --dist=loadfile -s ./tests/
+          python -m pytest -n 2 --dist=loadfile -s --make-reports=tests_torch_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_torch_gpu_failures_short.txt
+        
+      - name: Test suite reports artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v2
+        with:
+          name: run_all_tests_torch_gpu_test_reports
+          path: reports
+                  
 
   run_tests_tf_gpu:
     runs-on: [self-hosted, single-gpu]
@@ -110,7 +122,18 @@ jobs:
           CUDA_VISIBLE_DEVICES: 0
         run: |
           source .env/bin/activate
-          python -m pytest -n 2 --dist=loadfile -s ./tests/
+          python -m pytest -n 2 --dist=loadfile -s --make-reports=tests_tf_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_tf_gpu_failures_short.txt
+        
+      - name: Test suite reports artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v2
+        with:
+          name: run_all_tests_tf_gpu_test_reports
+          path: reports
 
   run_tests_torch_multiple_gpu:
     runs-on: [self-hosted, multi-gpu]
@@ -158,7 +181,18 @@ jobs:
           OMP_NUM_THREADS: 1
         run: |
           source .env/bin/activate
-          python -m pytest -n 2 --dist=loadfile -s ./tests/
+          python -m pytest -n 2 --dist=loadfile -s --make-reports=tests_torch_multiple_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_torch_multiple_gpu_failures_short.txt          
+
+      - name: Test suite reports artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v2
+        with:
+          name: run_all_tests_torch_multi_gpu_test_reports
+          path: reports
 
   run_tests_tf_multiple_gpu:
     runs-on: [self-hosted, multi-gpu]
@@ -206,4 +240,16 @@ jobs:
           OMP_NUM_THREADS: 1
         run: |
           source .env/bin/activate
-          python -m pytest -n 2 --dist=loadfile -s ./tests/
+          python -m pytest -n 2 --dist=loadfile -s --make-reports=tests_tf_multiple_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_tf_multiple_gpu_failures_short.txt
+
+      - name: Test suite reports artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v2
+        with:
+          name: run_all_tests_tf_multi_gpu_test_reports
+          path: reports
+          

.github/workflows/self-scheduled.yml

@@ -4,14 +4,186 @@ on:
   push:
     branches:
       - ci_*
-      - trigger-scheduled
   repository_dispatch:
   schedule:
     - cron: "0 0 * * *"
 
 jobs:
+  run_all_tests_torch_gpu:
+    runs-on: [self-hosted, single-gpu]
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Loading cache.
+        uses: actions/cache@v2
+        id: cache
+        with:
+          path: .env
+          key: v  1.1-slow_tests_torch_gpu-${{ hashFiles('setup.py') }}
+
+      - name: Python version
+        run: |
+          which python
+          python --version
+          pip --version
+
+      - name: Current dir
+        run: pwd
+      - run: nvidia-smi
+
+      - name: Create new python env (on self-hosted runners we have to handle isolation ourselves)
+        if: steps.cache.outputs.cache-hit != 'true'
+        run: |
+          python -m venv .env
+          source .env/bin/activate
+          which python
+          python --version
+          pip --version
+
+      - name: Install dependencies
+        run: |
+          source .env/bin/activate
+          pip install --upgrade pip
+          pip install .[torch,sklearn,testing,onnxruntime]
+          pip install git+https://github.com/huggingface/datasets
+          pip list
+
+      - name: Are GPUs recognized by our DL frameworks
+        run: |
+          source .env/bin/activate
+          python -c "import torch; print('Cuda available:', torch.cuda.is_available())"
+          python -c "import torch; print('Number of GPUs available:', torch.cuda.device_count())"
+
+      - name: Run all tests on GPU
+        env:
+          OMP_NUM_THREADS: 1
+          RUN_SLOW: yes
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 1 --dist=loadfile -s --make-reports=tests_torch_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_torch_gpu_failures_short.txt
+        
+      - name: Run examples tests on GPU
+        if: ${{ always() }}
+        env:
+          OMP_NUM_THREADS: 1
+          RUN_SLOW: yes
+        run: |
+          source .env/bin/activate
+          pip install -r examples/requirements.txt
+          python -m pytest -n 1 --dist=loadfile -s --make-reports=examples_torch_gpu examples
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/examples_torch_gpu_failures_short.txt
+
+      - name: Run all pipeline tests on GPU
+        if: ${{ always() }}
+        env:
+          TF_FORCE_GPU_ALLOW_GROWTH: "true"
+          OMP_NUM_THREADS: 1
+          RUN_SLOW: yes
+          RUN_PIPELINE_TESTS: yes
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 1 --dist=loadfile -s -m is_pipeline_test --make-reports=tests_torch_pipeline_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_torch_pipeline_gpu_failures_short.txt
+
+      - name: Test suite reports artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v2
+        with:
+          name: run_all_tests_torch_gpu_test_reports
+          path: reports
+
+
+  run_all_tests_tf_gpu:
+    runs-on: [self-hosted, single-gpu]
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Loading cache.
+        uses: actions/cache@v2
+        id: cache
+        with:
+          path: .env
+          key: v1.1-slow_tests_tf_gpu-${{ hashFiles('setup.py') }}
+
+      - name: Python version
+        run: |
+          which python
+          python --version
+          pip --version
+
+      - name: Current dir
+        run: pwd
+      - run: nvidia-smi
+
+      - name: Create new python env (on self-hosted runners we have to handle isolation ourselves)
+        if: steps.cache.outputs.cache-hit != 'true'
+        run: |
+          python -m venv .env
+          source .env/bin/activate
+          which python
+          python --version
+          pip --version
+
+      - name: Install dependencies
+        run: |
+          source .env/bin/activate
+          pip install --upgrade pip
+          pip install .[tf,sklearn,testing,onnxruntime]
+          pip install git+https://github.com/huggingface/datasets
+          pip list
+
+      - name: Are GPUs recognized by our DL frameworks
+        run: |
+          source .env/bin/activate
+          TF_CPP_MIN_LOG_LEVEL=3 python -c "import tensorflow as tf; print('TF GPUs available:', bool(tf.config.list_physical_devices('GPU')))"
+          TF_CPP_MIN_LOG_LEVEL=3 python -c "import tensorflow as tf; print('Number of TF GPUs available:', len(tf.config.list_physical_devices('GPU')))"
+
+      - name: Run all tests on GPU
+        env:
+          OMP_NUM_THREADS: 1
+          RUN_SLOW: yes
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 1 --dist=loadfile -s --make-reports=tests_tf_gpu tests
+          
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_tf_gpu_failures_short.txt
+
+      - name: Run all pipeline tests on GPU
+        if: ${{ always() }}
+        env:
+          TF_FORCE_GPU_ALLOW_GROWTH: "true"
+          OMP_NUM_THREADS: 1
+          RUN_SLOW: yes
+          RUN_PIPELINE_TESTS: yes
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 1 --dist=loadfile -s -m is_pipeline_test --make-reports=tests_tf_pipelines_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_tf_pipelines_gpu_failures_short.txt
+
+      - name: Test suite reports artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v2
+        with:
+          name: run_all_tests_tf_gpu_test_reports
+          path: reports
+          
   run_all_tests_torch_multiple_gpu:
-    runs-on: [self-hosted, multi-gpu-tests]
+    runs-on: [self-hosted, multi-gpu]
     steps:
       - uses: actions/checkout@v2
 
@@ -61,11 +233,90 @@ jobs:
           RUN_SLOW: yes
         run: |
           source .env/bin/activate
-          python -m pytest -n 1 --dist=loadfile -s --make_reports=tests_torch tests
+          python -m pytest -n 1 --dist=loadfile -s --make-reports=tests_torch_multiple_gpu tests
 
       - name: Failure short reports
         if: ${{ always() }}
-        run: cat reports/report_tests_torch_failures_short.txt
+        run: cat reports/tests_torch_multiple_gpu_failures_short.txt
+        
+      - name: Run all pipeline tests on GPU
+        if: ${{ always() }}
+        env:
+          TF_FORCE_GPU_ALLOW_GROWTH: "true"
+          OMP_NUM_THREADS: 1
+          RUN_SLOW: yes
+          RUN_PIPELINE_TESTS: yes
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 1 --dist=loadfile -s -m is_pipeline_test --make-reports=tests_torch_pipeline_multiple_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_torch_pipeline_multiple_gpu_failures_short.txt
+
+      - name: Test suite reports artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v2
+        with:
+          name: run_all_tests_torch_multi_gpu_test_reports
+          path: reports
+
+  run_all_tests_tf_multiple_gpu:
+    runs-on: [self-hosted, multi-gpu]
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Loading cache.
+        uses: actions/cache@v2
+        id: cache
+        with:
+          path: .env
+          key: v1.1-slow_tests_tf_multi_gpu-${{ hashFiles('setup.py') }}
+
+      - name: Python version
+        run: |
+          which python
+          python --version
+          pip --version
+
+      - name: Current dir
+        run: pwd
+      - run: nvidia-smi
+
+      - name: Create new python env (on self-hosted runners we have to handle isolation ourselves)
+        if: steps.cache.outputs.cache-hit != 'true'
+        run: |
+          python -m venv .env
+          source .env/bin/activate
+          which python
+          python --version
+          pip --version
+
+      - name: Install dependencies
+        run: |
+          source .env/bin/activate
+          pip install --upgrade pip
+          pip install .[tf,sklearn,testing,onnxruntime]
+          pip install git+https://github.com/huggingface/datasets
+          pip list
+
+      - name: Are GPUs recognized by our DL frameworks
+        run: |
+          source .env/bin/activate
+          TF_CPP_MIN_LOG_LEVEL=3 python -c "import tensorflow as tf; print('TF GPUs available:', bool(tf.config.list_physical_devices('GPU')))"
+          TF_CPP_MIN_LOG_LEVEL=3 python -c "import tensorflow as tf; print('Number of TF GPUs available:', len(tf.config.list_physical_devices('GPU')))"
+
+      - name: Run all tests on GPU
+        env:
+          OMP_NUM_THREADS: 1
+          RUN_SLOW: yes
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 1 --dist=loadfile -s --make-reports=tests_tf_multiple_gpu tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/tests_tf_multiple_gpu_failures_short.txt
 
       - name: Run all pipeline tests on GPU
         if: ${{ always() }}
@@ -76,16 +327,16 @@ jobs:
           RUN_PIPELINE_TESTS: yes
         run: |
           source .env/bin/activate
-          python -m pytest -n 1 --dist=loadfile -s -m is_pipeline_test --make_reports=tests_torch_pipeline tests
-
+          python -m pytest -n 1 --dist=loadfile -s -m is_pipeline_test --make-reports=tests_tf_pipelines_multiple_gpu tests
+          
       - name: Failure short reports
         if: ${{ always() }}
-        run: cat reports/report_tests_torch_pipeline_failures_short.txt
+        run: cat reports/tests_tf_multiple_gpu_pipelines_failures_short.txt
 
       - name: Test suite reports artifacts
         if: ${{ always() }}
         uses: actions/upload-artifact@v2
         with:
-          name: run_all_tests_torch_multi_gpu_test_reports
+          name: run_all_tests_tf_multi_gpu_test_reports
           path: reports
-
+          

examples/bert-loses-patience/pabee/modeling_pabee_albert.py

@@ -65,7 +65,7 @@ class AlbertModelWithPabee(AlbertModel):
 
         self.encoder = AlbertTransformerWithPabee(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
         self.patience = 0
         self.inference_instances_num = 0
         self.inference_layers_num = 0
@@ -228,7 +228,7 @@ class AlbertForSequenceClassificationWithPabee(AlbertPreTrainedModel):
             [nn.Linear(config.hidden_size, self.config.num_labels) for _ in range(config.num_hidden_layers)]
         )
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ALBERT_INPUTS_DOCSTRING)
     def forward(

examples/bert-loses-patience/pabee/modeling_pabee_bert.py

@@ -70,7 +70,7 @@ class BertModelWithPabee(BertModel):
 
         self.encoder = BertEncoderWithPabee(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
         self.patience = 0
         self.inference_instances_num = 0
         self.inference_layers_num = 0
@@ -252,7 +252,7 @@ class BertForSequenceClassificationWithPabee(BertPreTrainedModel):
             [nn.Linear(config.hidden_size, self.config.num_labels) for _ in range(config.num_hidden_layers)]
         )
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING)
     def forward(

examples/conftest.py

@@ -17,17 +17,14 @@ warnings.simplefilter(action="ignore", category=FutureWarning)
 
 
 def pytest_addoption(parser):
-    parser.addoption(
-        "--make_reports",
-        action="store",
-        default=False,
-        help="generate report files - the value will be used as a `report_`+val+`reportname.txt`",
-    )
+    from transformers.testing_utils import pytest_addoption_shared
+
+    pytest_addoption_shared(parser)
 
 
 def pytest_terminal_summary(terminalreporter):
     from transformers.testing_utils import pytest_terminal_summary_main
 
-    make_reports = terminalreporter.config.getoption("--make_reports")
+    make_reports = terminalreporter.config.getoption("--make-reports")
     if make_reports:
         pytest_terminal_summary_main(terminalreporter, id=make_reports)

examples/deebert/src/modeling_highway_bert.py

@@ -115,7 +115,7 @@ class DeeBertModel(BertPreTrainedModel):
         self.encoder = DeeBertEncoder(config)
         self.pooler = BertPooler(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def init_highway_pooler(self):
         self.encoder.init_highway_pooler(self.pooler)
@@ -286,7 +286,7 @@ class DeeBertForSequenceClassification(BertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING)
     def forward(

examples/deebert/src/modeling_highway_roberta.py

@@ -23,7 +23,7 @@ class DeeRobertaModel(DeeBertModel):
         super().__init__(config)
 
         self.embeddings = RobertaEmbeddings(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
 
 @add_start_docstrings(

examples/language-modeling/README.md

@@ -10,7 +10,7 @@ These scripts leverage the 🤗 Datasets library and the Trainer API. You can ea
 need extra processing on your datasets.
 
 **Note:** The old script `run_language_modeling.py` is still available
-[here](https://github.com/huggingface/transformers/blob/master/examples/contrib/legacy/language-modeling/run_language_modeling.py).
+[here](https://github.com/huggingface/transformers/blob/master/examples/contrib/legacy/run_language_modeling.py).
 
 The following examples, will run on a datasets hosted on our [hub](https://huggingface.co/datasets) or with your own
 text files for training and validation. We give examples of both below.
@@ -68,13 +68,13 @@ python run_mlm.py \
 To run on your own training and validation files, use the following command:
 
 ```bash
-python run_clm.py \
+python run_mlm.py \
     --model_name_or_path roberta-base \
     --train_file path_to_train_file \
     --validation_file path_to_validation_file \
     --do_train \
     --do_eval \
-    --output_dir /tmp/test-clm
+    --output_dir /tmp/test-mlm
 ```
 
 If your dataset is organized with one sample per line, you can use the `--line_by_line` flag (otherwise the script

examples/language-modeling/run_clm.py

@@ -190,7 +190,7 @@ def main():
         if data_args.train_file is not None:
             data_files["train"] = data_args.train_file
         if data_args.validation_file is not None:
-            data_files["validation"] = data_args.train_file
+            data_files["validation"] = data_args.validation_file
         extension = data_args.train_file.split(".")[-1]
         if extension == "txt":
             extension = "text"

examples/language-modeling/run_mlm.py

@@ -201,7 +201,7 @@ def main():
         if data_args.train_file is not None:
             data_files["train"] = data_args.train_file
         if data_args.validation_file is not None:
-            data_files["validation"] = data_args.train_file
+            data_files["validation"] = data_args.validation_file
         extension = data_args.train_file.split(".")[-1]
         if extension == "txt":
             extension = "text"

examples/language-modeling/run_mlm_wwm.py

@@ -204,7 +204,7 @@ def main():
     if data_args.train_file is not None:
         data_files["train"] = data_args.train_file
     if data_args.validation_file is not None:
-        data_files["validation"] = data_args.train_file
+        data_files["validation"] = data_args.validation_file
     extension = data_args.train_file.split(".")[-1]
     if extension == "txt":
         extension = "text"

examples/language-modeling/run_plm.py

@@ -198,7 +198,7 @@ def main():
         if data_args.train_file is not None:
             data_files["train"] = data_args.train_file
         if data_args.validation_file is not None:
-            data_files["validation"] = data_args.train_file
+            data_files["validation"] = data_args.validation_file
         extension = data_args.train_file.split(".")[-1]
         if extension == "txt":
             extension = "text"

examples/movement-pruning/emmental/modeling_bert_masked.py

@@ -482,7 +482,7 @@ class MaskedBertModel(MaskedBertPreTrainedModel):
         self.encoder = BertEncoder(config)
         self.pooler = BertPooler(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -669,7 +669,7 @@ class MaskedBertForSequenceClassification(MaskedBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MASKED_BERT_INPUTS_DOCSTRING)
     def forward(
@@ -754,7 +754,7 @@ class MaskedBertForMultipleChoice(MaskedBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MASKED_BERT_INPUTS_DOCSTRING)
     def forward(
@@ -844,7 +844,7 @@ class MaskedBertForTokenClassification(MaskedBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MASKED_BERT_INPUTS_DOCSTRING)
     def forward(
@@ -930,7 +930,7 @@ class MaskedBertForQuestionAnswering(MaskedBertPreTrainedModel):
         self.bert = MaskedBertModel(config)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MASKED_BERT_INPUTS_DOCSTRING)
     def forward(

examples/seq2seq/finetune_trainer.py

@@ -4,16 +4,9 @@ import sys
 from dataclasses import dataclass, field
 from typing import Optional
 
-from seq2seq_trainer import Seq2SeqTrainer, arg_to_scheduler_choices
-from transformers import (
-    AutoConfig,
-    AutoModelForSeq2SeqLM,
-    AutoTokenizer,
-    HfArgumentParser,
-    MBartTokenizer,
-    TrainingArguments,
-    set_seed,
-)
+from seq2seq_trainer import Seq2SeqTrainer
+from seq2seq_training_args import Seq2SeqTrainingArguments
+from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, HfArgumentParser, MBartTokenizer, set_seed
 from transformers.trainer_utils import EvaluationStrategy
 from utils import (
     Seq2SeqDataCollator,
@@ -33,41 +26,6 @@ from utils import (
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class Seq2SeqTrainingArguments(TrainingArguments):
-    """
-    Parameters:
-        label_smoothing (:obj:`float`, `optional`, defaults to 0):
-            The label smoothing epsilon to apply (if not zero).
-        sortish_sampler (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to SortishSamler or not. It sorts the inputs according to lenghts in-order to minimizing the padding size.
-        predict_with_generate (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to use generate to calculate generative metrics (ROUGE, BLEU).
-    """
-
-    label_smoothing: Optional[float] = field(
-        default=0.0, metadata={"help": "The label smoothing epsilon to apply (if not zero)."}
-    )
-    sortish_sampler: bool = field(default=False, metadata={"help": "Whether to SortishSamler or not."})
-    predict_with_generate: bool = field(
-        default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}
-    )
-    adafactor: bool = field(default=False, metadata={"help": "whether to use adafactor"})
-    encoder_layerdrop: Optional[float] = field(
-        default=None, metadata={"help": "Encoder layer dropout probability. Goes into model.config."}
-    )
-    decoder_layerdrop: Optional[float] = field(
-        default=None, metadata={"help": "Decoder layer dropout probability. Goes into model.config."}
-    )
-    dropout: Optional[float] = field(default=None, metadata={"help": "Dropout probability. Goes into model.config."})
-    attention_dropout: Optional[float] = field(
-        default=None, metadata={"help": "Attention dropout probability. Goes into model.config."}
-    )
-    lr_scheduler: Optional[str] = field(
-        default="linear", metadata={"help": f"Which lr scheduler to use. Selected in {arg_to_scheduler_choices}"}
-    )
-
-
 @dataclass
 class ModelArguments:
     """

examples/seq2seq/seq2seq_trainer.py

@@ -30,7 +30,6 @@ arg_to_scheduler = {
     "constant": get_constant_schedule,
     "constant_w_warmup": get_constant_schedule_with_warmup,
 }
-arg_to_scheduler_choices = sorted(arg_to_scheduler.keys())
 
 
 class Seq2SeqTrainer(Trainer):

examples/seq2seq/seq2seq_training_args.py

@@ -0,0 +1,45 @@
+import logging
+from dataclasses import dataclass, field
+from typing import Optional
+
+from seq2seq_trainer import arg_to_scheduler
+from transformers import TrainingArguments
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Seq2SeqTrainingArguments(TrainingArguments):
+    """
+    Parameters:
+        label_smoothing (:obj:`float`, `optional`, defaults to 0):
+            The label smoothing epsilon to apply (if not zero).
+        sortish_sampler (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to SortishSamler or not. It sorts the inputs according to lenghts in-order to minimizing the padding size.
+        predict_with_generate (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to use generate to calculate generative metrics (ROUGE, BLEU).
+    """
+
+    label_smoothing: Optional[float] = field(
+        default=0.0, metadata={"help": "The label smoothing epsilon to apply (if not zero)."}
+    )
+    sortish_sampler: bool = field(default=False, metadata={"help": "Whether to SortishSamler or not."})
+    predict_with_generate: bool = field(
+        default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}
+    )
+    adafactor: bool = field(default=False, metadata={"help": "whether to use adafactor"})
+    encoder_layerdrop: Optional[float] = field(
+        default=None, metadata={"help": "Encoder layer dropout probability. Goes into model.config."}
+    )
+    decoder_layerdrop: Optional[float] = field(
+        default=None, metadata={"help": "Decoder layer dropout probability. Goes into model.config."}
+    )
+    dropout: Optional[float] = field(default=None, metadata={"help": "Dropout probability. Goes into model.config."})
+    attention_dropout: Optional[float] = field(
+        default=None, metadata={"help": "Attention dropout probability. Goes into model.config."}
+    )
+    lr_scheduler: Optional[str] = field(
+        default="linear",
+        metadata={"help": f"Which lr scheduler to use. Selected in {sorted(arg_to_scheduler.keys())}"},
+    )

model_cards/deepset/roberta-base-squad2-v2/README.md

@@ -0,0 +1,117 @@
+---
+datasets:
+- squad_v2
+---
+
+# roberta-base for QA 
+
+## Overview
+**Language model:** roberta-base  
+**Language:** English  
+**Downstream-task:** Extractive QA  
+**Training data:** SQuAD 2.0  
+**Eval data:** SQuAD 2.0  
+**Code:**  See [example](https://github.com/deepset-ai/FARM/blob/master/examples/question_answering.py) in [FARM](https://github.com/deepset-ai/FARM/blob/master/examples/question_answering.py)  
+**Infrastructure**: 4x Tesla v100
+
+## Hyperparameters
+
+```
+batch_size = 96
+n_epochs = 2
+base_LM_model = "roberta-base"
+max_seq_len = 386
+learning_rate = 3e-5
+lr_schedule = LinearWarmup
+warmup_proportion = 0.2
+doc_stride=128
+max_query_length=64
+``` 
+
+## Performance
+Evaluated on the SQuAD 2.0 dev set with the [official eval script](https://worksheets.codalab.org/rest/bundles/0x6b567e1cf2e041ec80d7098f031c5c9e/contents/blob/).
+
+```
+"exact": 79.97136359807968
+"f1": 83.00449234495325
+
+"total": 11873
+"HasAns_exact": 78.03643724696356
+"HasAns_f1": 84.11139298441825
+"HasAns_total": 5928
+"NoAns_exact": 81.90075693860386
+"NoAns_f1": 81.90075693860386
+"NoAns_total": 5945
+```
+
+## Usage
+
+### In Transformers
+```python
+from transformers.pipelines import pipeline
+from transformers.modeling_auto import AutoModelForQuestionAnswering
+from transformers.tokenization_auto import AutoTokenizer
+
+model_name = "deepset/roberta-base-squad2-v2"
+
+# a) Get predictions
+nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)
+QA_input = {
+    'question': 'Why is model conversion important?',
+    'context': 'The option to convert models between FARM and transformers gives freedom to the user and let people easily switch between frameworks.'
+}
+res = nlp(QA_input)
+
+# b) Load model & tokenizer
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+```
+
+### In FARM
+
+```python
+from farm.modeling.adaptive_model import AdaptiveModel
+from farm.modeling.tokenization import Tokenizer
+from farm.infer import Inferencer
+
+model_name = "deepset/roberta-base-squad2-v2"
+
+# a) Get predictions
+nlp = Inferencer.load(model_name, task_type="question_answering")
+QA_input = [{"questions": ["Why is model conversion important?"],
+             "text": "The option to convert models between FARM and transformers gives freedom to the user and let people easily switch between frameworks."}]
+res = nlp.inference_from_dicts(dicts=QA_input, rest_api_schema=True)
+
+# b) Load model & tokenizer
+model = AdaptiveModel.convert_from_transformers(model_name, device="cpu", task_type="question_answering")
+tokenizer = Tokenizer.load(model_name)
+```
+
+### In haystack
+For doing QA at scale (i.e. many docs instead of single paragraph), you can load the model also in [haystack](https://github.com/deepset-ai/haystack/):
+```python
+reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2")
+# or 
+reader = TransformersReader(model="deepset/roberta-base-squad2",tokenizer="deepset/roberta-base-squad2")
+```
+
+
+## Authors
+Branden Chan: `branden.chan [at] deepset.ai`
+Timo Möller: `timo.moeller [at] deepset.ai`
+Malte Pietsch: `malte.pietsch [at] deepset.ai`
+Tanay Soni: `tanay.soni [at] deepset.ai`
+
+## About us
+![deepset logo](https://raw.githubusercontent.com/deepset-ai/FARM/master/docs/img/deepset_logo.png)
+
+We bring NLP to the industry via open source!  
+Our focus: Industry specific language models & large scale QA systems.  
+  
+Some of our work: 
+- [German BERT (aka "bert-base-german-cased")](https://deepset.ai/german-bert)
+- [FARM](https://github.com/deepset-ai/FARM)
+- [Haystack](https://github.com/deepset-ai/haystack/)
+
+Get in touch:
+[Twitter](https://twitter.com/deepset_ai) | [LinkedIn](https://www.linkedin.com/company/deepset-ai/) | [Website](https://deepset.ai)

model_cards/deepset/roberta-base-squad2/README.md

@@ -5,6 +5,8 @@ datasets:
 
 # roberta-base for QA 
 
+NOTE: This model has been superseded by deepset/roberta-base-squad2-v2. For an explanation of why, see [this github issue](https://github.com/deepset-ai/FARM/issues/552) from the FARM repository.
+
 ## Overview
 **Language model:** roberta-base  
 **Language:** English  
@@ -106,4 +108,3 @@ Some of our work:
 
 Get in touch:
 [Twitter](https://twitter.com/deepset_ai) | [LinkedIn](https://www.linkedin.com/company/deepset-ai/) | [Website](https://deepset.ai)
-

model_cards/distilbert-base-cased-README.md

@@ -1,3 +1,40 @@
 ---
+language: en
 license: apache-2.0
+datasets:
+- bookcorpus
+- wikipedia
 ---
+
+# DistilBERT base model (cased)
+
+This model is a distilled version of the [BERT base model](https://huggingface.co/bert-base-cased).
+It was introduced in [this paper](https://arxiv.org/abs/1910.01108).
+The code for the distillation process can be found
+[here](https://github.com/huggingface/transformers/tree/master/examples/distillation).
+This model is cased: it does make a difference between english and English.
+
+All the training details on the pre-training, the uses, limitations and potential biases are the same as for [DistilBERT-base-uncased](https://huggingface.co/distilbert-base-uncased).
+We highly encourage to check it if you want to know more.
+
+## Evaluation results
+
+When fine-tuned on downstream tasks, this model achieves the following results:
+
+Glue test results:
+
+| Task | MNLI | QQP  | QNLI | SST-2 | CoLA | STS-B | MRPC | RTE  |
+|:----:|:----:|:----:|:----:|:-----:|:----:|:-----:|:----:|:----:|
+|      | 81.5 | 87.8 | 88.2 | 90.4  | 47.2 | 85.5  | 85.6 | 60.6 |
+
+### BibTeX entry and citation info
+
+```bibtex
+@article{Sanh2019DistilBERTAD,
+  title={DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter},
+  author={Victor Sanh and Lysandre Debut and Julien Chaumond and Thomas Wolf},
+  journal={ArXiv},
+  year={2019},
+  volume={abs/1910.01108}
+}
+```

model_cards/distilbert-base-cased-distilled-squad-README.md

@@ -6,3 +6,8 @@ metrics:
 - squad
 license: apache-2.0
 ---
+
+# DistilBERT base cased distilled SQuAD
+
+This model is a fine-tune checkpoint of [DistilBERT-base-cased](https://huggingface.co/distilbert-base-cased), fine-tuned using (a second step of) knowledge distillation on SQuAD v1.1.
+This model reaches a F1 score of 87.1 on the dev set (for comparison, BERT bert-base-cased version reaches a F1 score of 88.7).

model_cards/distilbert-base-multilingual-cased-README.md

@@ -1,4 +1,35 @@
 ---
 language: multilingual
 license: apache-2.0
+datasets:
+- wikipedia
 ---
+
+# DistilBERT base multilingual model (cased)
+
+This model is a distilled version of the [BERT base multilingual model](bert-base-multilingual-cased). The code for the distillation process can be found
+[here](https://github.com/huggingface/transformers/tree/master/examples/distillation). This model is cased: it does make a difference between english and English.
+
+The model is trained on the concatenation of Wikipedia in 104 different languages listed [here](https://github.com/google-research/bert/blob/master/multilingual.md#list-of-languages).
+The model has 6 layers, 768 dimension and 12 heads, totalizing 134M parameters (compared to 177M parameters for mBERT-base).
+On average DistilmBERT is twice as fast as mBERT-base.
+
+We encourage to check [BERT base multilingual model](bert-base-multilingual-cased) to know more about usage, limitations and potential biases.
+
+| Model                        | English | Spanish | Chinese | German | Arabic  | Urdu |
+| :---:                        | :---:   | :---:   | :---:   | :---:  | :---:   | :---:|
+| mBERT base cased (computed)  | 82.1    | 74.6    | 69.1    | 72.3   | 66.4    | 58.5 |
+| mBERT base uncased (reported)| 81.4    | 74.3    | 63.8    | 70.5   | 62.1    | 58.3 |
+| DistilmBERT                  | 78.2    | 69.1    | 64.0    | 66.3   | 59.1    | 54.7 |
+
+### BibTeX entry and citation info
+
+```bibtex
+@article{Sanh2019DistilBERTAD,
+  title={DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter},
+  author={Victor Sanh and Lysandre Debut and Julien Chaumond and Thomas Wolf},
+  journal={ArXiv},
+  year={2019},
+  volume={abs/1910.01108}
+}
+```

model_cards/distilbert-base-uncased-README.md

@@ -10,7 +10,7 @@ datasets:
 
 # DistilBERT base model (uncased)
 
-This model is a distilled version of the [BERT base mode](https://huggingface.co/distilbert-base-uncased). It was
+This model is a distilled version of the [BERT base model](https://huggingface.co/bert-base-uncased). It was
 introduced in [this paper](https://arxiv.org/abs/1910.01108). The code for the distillation process can be found
 [here](https://github.com/huggingface/transformers/tree/master/examples/distillation). This model is uncased: it does
 not make a difference between english and English.
@@ -102,7 +102,7 @@ output = model(encoded_input)
 
 Even if the training data used for this model could be characterized as fairly neutral, this model can have biased
 predictions. It also inherits some of
-[the bias of its teacher model](https://huggingface.co/bert-base-uncased#limitations-and-bias). 
+[the bias of its teacher model](https://huggingface.co/bert-base-uncased#limitations-and-bias).
 
 ```python
 >>> from transformers import pipeline
@@ -196,9 +196,9 @@ When fine-tuned on downstream tasks, this model achieves the following results:
 
 Glue test results:
 
-| Task | MNLI | QQP  | QNLI | SST-2 | CoLA | STS-B | MRPC | RTE  | Average |
-|:----:|:----:|:----:|:----:|:-----:|:----:|:-----:|:----:|:----:|:-------:|
-|      | 82.2 | 88.5 | 89.2 | 91.3  | 51.3 | 85.8  | 87.5 | 59.9 | 77.0    |
+| Task | MNLI | QQP  | QNLI | SST-2 | CoLA | STS-B | MRPC | RTE  |
+|:----:|:----:|:----:|:----:|:-----:|:----:|:-----:|:----:|:----:|
+|      | 82.2 | 88.5 | 89.2 | 91.3  | 51.3 | 85.8  | 87.5 | 59.9 |
 
 
 ### BibTeX entry and citation info

model_cards/distilbert-base-uncased-distilled-squad-README.md

@@ -1,4 +1,5 @@
 ---
+language: en
 datasets:
 - squad
 widget:
@@ -8,3 +9,8 @@ widget:
   context: "The Amazon rainforest (Portuguese: Floresta Amazônica or Amazônia; Spanish: Selva Amazónica, Amazonía or usually Amazonia; French: Forêt amazonienne; Dutch: Amazoneregenwoud), also known in English as Amazonia or the Amazon Jungle, is a moist broadleaf forest that covers most of the Amazon basin of South America. This basin encompasses 7,000,000 square kilometres (2,700,000 sq mi), of which 5,500,000 square kilometres (2,100,000 sq mi) are covered by the rainforest. This region includes territory belonging to nine nations. The majority of the forest is contained within Brazil, with 60% of the rainforest, followed by Peru with 13%, Colombia with 10%, and with minor amounts in Venezuela, Ecuador, Bolivia, Guyana, Suriname and French Guiana. States or departments in four nations contain \"Amazonas\" in their names. The Amazon represents over half of the planet's remaining rainforests, and comprises the largest and most biodiverse tract of tropical rainforest in the world, with an estimated 390 billion individual trees divided into 16,000 species."
 license: apache-2.0
 ---
+
+# DistilBERT base uncased distilled SQuAD
+
+This model is a fine-tune checkpoint of [DistilBERT-base-uncased](https://huggingface.co/distilbert-base-uncased), fine-tuned using (a second step of) knowledge distillation on SQuAD v1.1.
+This model reaches a F1 score of 86.9 on the dev set (for comparison, Bert bert-base-uncased version reaches a F1 score of 88.5).

model_cards/distilbert-base-uncased-finetuned-sst-2-english-README.md

@@ -1,3 +1,19 @@
 ---
+language: en
 license: apache-2.0
+datasets:
+- sst-2
 ---
+
+# DistilBERT base uncased finetuned SST-2
+
+This model is a fine-tune checkpoint of [DistilBERT-base-uncased](https://huggingface.co/distilbert-base-uncased), fine-tuned on SST-2.
+This model reaches an accuracy of 91.3 on the dev set (for comparison, Bert bert-base-uncased version reaches an accuracy of 92.7).
+
+# Fine-tuning hyper-parameters
+
+- learning_rate = 1e-5
+- batch_size = 32
+- warmup = 600
+- max_seq_length = 128
+- num_train_epochs = 3.0

model_cards/distilgpt2-README.md

@@ -1,10 +1,21 @@
 ---
+language: en
 tags:
 - exbert
 
 license: apache-2.0
+datasets:
+- openwebtext
 ---
 
+# DistilGPT2
+
+DistilGPT2 English language model pretrained with the supervision of [GPT2](https://huggingface.co/gpt2) (the smallest version of GPT2) on [OpenWebTextCorpus](https://skylion007.github.io/OpenWebTextCorpus/), a reproduction of OpenAI's WebText dataset. The model has 6 layers, 768 dimension and 12 heads, totalizing 82M parameters (compared to 124M parameters for GPT2). On average, DistilGPT2 is two times faster than GPT2.
+
+On the [WikiText-103](https://blog.einstein.ai/the-wikitext-long-term-dependency-language-modeling-dataset/) benchmark, GPT2 reaches a perplexity on the test set of 16.3 compared to 21.1 for DistilGPT2 (after fine-tuning on the train set).
+
+We encourage to check [GPT2](https://huggingface.co/gpt2) to know more about usage, limitations and potential biases.
+
 <a href="https://huggingface.co/exbert/?model=distilgpt2">
 	<img width="300px" src="https://cdn-media.huggingface.co/exbert/button.png">
 </a>

model_cards/distilroberta-base-README.md

@@ -1,10 +1,50 @@
 ---
+language: en
 tags:
 - exbert
 
 license: apache-2.0
+datasets:
+- openwebtext
 ---
 
+# DistilRoBERTa base model
+
+This model is a distilled version of the [RoBERTa-base model](https://huggingface.co/roberta-base). It follows the same training procedure as [DistilBERT](https://huggingface.co/distilbert-base-uncased).
+The code for the distillation process can be found [here](https://github.com/huggingface/transformers/tree/master/examples/distillation).
+This model is case-sensitive: it makes a difference between english and English.
+
+The model has 6 layers, 768 dimension and 12 heads, totalizing 82M parameters (compared to 125M parameters for RoBERTa-base).
+On average DistilRoBERTa is twice as fast as Roberta-base.
+
+We encourage to check [RoBERTa-base model](https://huggingface.co/roberta-base) to know more about usage, limitations and potential biases.
+
+## Training data
+
+DistilRoBERTa was pre-trained on [OpenWebTextCorpus](https://skylion007.github.io/OpenWebTextCorpus/), a reproduction of OpenAI's WebText dataset (it is ~4 times less training data than the teacher RoBERTa).
+
+## Evaluation results
+
+When fine-tuned on downstream tasks, this model achieves the following results:
+
+Glue test results:
+
+| Task | MNLI | QQP  | QNLI | SST-2 | CoLA | STS-B | MRPC | RTE  |
+|:----:|:----:|:----:|:----:|:-----:|:----:|:-----:|:----:|:----:|
+|      | 84.0 | 89.4 | 90.8 | 92.5  | 59.3 | 88.3  | 86.6 | 67.9 |
+
+### BibTeX entry and citation info
+
+```bibtex
+@article{Sanh2019DistilBERTAD,
+  title={DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter},
+  author={Victor Sanh and Lysandre Debut and Julien Chaumond and Thomas Wolf},
+  journal={ArXiv},
+  year={2019},
+  volume={abs/1910.01108}
+}
+```
+
 <a href="https://huggingface.co/exbert/?model=distilroberta-base">
 	<img width="300px" src="https://cdn-media.huggingface.co/exbert/button.png">
 </a>

model_cards/mrm8488/t5-base-finetuned-qasc/README.md

@@ -0,0 +1,66 @@
+---
+language: en
+datasets:
+- qasc
+---
+
+# T5-base fine-tuned on QASC 
+
+[Google's T5](https://ai.googleblog.com/2020/02/exploring-transfer-learning-with-t5.html) fine-tuned on [QASC](https://allenai.org/data/qasc) for **QA** (via *sentence composition*) downstream task.
+
+## Details of T5
+
+The **T5** model was presented in [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/pdf/1910.10683.pdf) by *Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu* in Here the abstract:
+
+Transfer learning, where a model is first pre-trained on a data-rich task before being fine-tuned on a downstream task, has emerged as a powerful technique in natural language processing (NLP). The effectiveness of transfer learning has given rise to a diversity of approaches, methodology, and practice. In this paper, we explore the landscape of transfer learning techniques for NLP by introducing a unified framework that converts every language problem into a text-to-text format. Our systematic study compares pre-training objectives, architectures, unlabeled datasets, transfer approaches, and other factors on dozens of language understanding tasks. By combining the insights from our exploration with scale and our new “Colossal Clean Crawled Corpus”, we achieve state-of-the-art results on many benchmarks covering summarization, question answering, text classification, and more. To facilitate future work on transfer learning for NLP, we release our dataset, pre-trained models, and code.
+
+![model image](https://i.imgur.com/jVFMMWR.png)
+
+
+## Details of the dataset 📚 
+
+**Question Answering via Sentence Composition** (QASC) is a question-answering dataset with a focus on sentence composition. It consists of 9,980 8-way multiple-choice questions about grade school science (8,134 train, 926 dev, 920 test), and comes with a corpus of 17M sentences.
+
+
+## Model fine-tuning 🏋️‍
+
+The training script is a slightly modified version of [this  awesome one](https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb) by [Suraj Patil](https://twitter.com/psuraj28). The **context** passed to the *encoder* is the combination of the 2 *facts* (`fact1` and `fact2`). The **question** is just the `formatted_question` field. The **answer** passed to the *decoder* is the`text` right answer instead of the `label` (A, B, C... See `choices` field). More details about the dataset format/fields [here](https://huggingface.co/nlp/viewer/?dataset=qasc)
+
+## Metrics on validation set 📋
+
+| Metric | Score |
+|--------|-------|
+|Accuracy (EM) | **97.73**|
+
+
+## Model in Action 🚀
+
+```python
+from transformers import AutoModelWithLMHead, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-qasc")
+model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-qasc")
+
+def get_response(question, context, max_length=64):
+  input_text = 'question: %s  context: %s' % (question, context)
+  features = tokenizer([input_text], return_tensors='pt')
+
+  output = model.generate(input_ids=features['input_ids'], 
+               attention_mask=features['attention_mask'],
+               max_length=max_length)
+
+  return tokenizer.decode(output[0])
+  
+fact_1 = 'a watch is used for measuring time'
+fact_2 = 'Times are measured in seconds.'
+context = fact_1 + ' ' + fact_2
+question = 'What can be used to measure seconds? (A) Watch (B) seconds (C) fluid (D) Ruler (E) goggles (F) glasses (G) Drill (H) Scale'
+
+get_response(question, context)
+
+# output: 'Watch'
+```
+
+> Created by [Manuel Romero/@mrm8488](https://twitter.com/mrm8488) | [LinkedIn](https://www.linkedin.com/in/manuel-romero-cs/)
+
+> Made with <span style="color: #e25555;">&hearts;</span> in Spain

setup.py

@@ -105,7 +105,7 @@ extras["sentencepiece"] = ["sentencepiece==0.1.91"]
 extras["retrieval"] = ["faiss-cpu", "datasets"]
 extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator", "parameterized", "psutil"] + extras["retrieval"]
 # sphinx-rtd-theme==0.5.0 introduced big changes in the style.
-extras["docs"] = ["recommonmark", "sphinx", "sphinx-markdown-tables", "sphinx-rtd-theme==0.4.3", "sphinx-copybutton"]
+extras["docs"] = ["recommonmark", "sphinx==3.2.1", "sphinx-markdown-tables", "sphinx-rtd-theme==0.4.3", "sphinx-copybutton"]
 extras["quality"] = ["black >= 20.8b1", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 
 

src/transformers/__init__.py

@@ -284,6 +284,7 @@ if is_torch_available():
         DataCollatorForNextSentencePrediction,
         DataCollatorForPermutationLanguageModeling,
         DataCollatorForSOP,
+        DataCollatorForTokenClassification,
         DataCollatorForWholeWordMask,
         DataCollatorWithPadding,
         default_data_collator,

src/transformers/benchmark/benchmark.py

@@ -254,6 +254,7 @@ class PyTorchBenchmark(Benchmark):
             else:
                 # cpu
                 memory_bytes = measure_peak_memory_cpu(func)
+                print("PEAK", memory_bytes)
                 memory = Memory(memory_bytes) if isinstance(memory_bytes, int) else memory_bytes
 
             if self.args.trace_memory_line_by_line:
@@ -261,6 +262,7 @@ class PyTorchBenchmark(Benchmark):
             else:
                 summary = None
 
+            print(memory, summary)
             return memory, summary
         except RuntimeError as e:
             self.print_fn("Doesn't fit on GPU. {}".format(e))

src/transformers/benchmark/benchmark_utils.py

@@ -76,12 +76,7 @@ def separate_process_wrapper_fn(func: Callable[[], None], do_multi_processing: b
         # run function in an individual
         # process to get correct memory
         def wrapper_func(queue: Queue, *args):
-            try:
-                result = func(*args)
-            except Exception as e:
-                logger.error(e)
-                print(e)
-                result = "N/A"
+            result = func(*args)
             queue.put(result)
 
         queue = Queue()
@@ -291,13 +286,13 @@ def measure_peak_memory_cpu(function: Callable[[], None], interval=0.5, device_i
                 # receive memory and num measurements
                 max_memory = parent_connection.recv()
                 num_measurements = parent_connection.recv()
-            except Exception:
+            except Exception as e:
                 # kill process in a clean way
                 parent = psutil.Process(os.getpid())
                 for child in parent.children(recursive=True):
                     os.kill(child.pid, SIGKILL)
                 mem_process.join(0)
-                raise RuntimeError("Process killed. Error in Process")
+                raise RuntimeError(f"Process killed. Error in Process: {e}")
 
             # run process at least 20 * interval or until it finishes
             mem_process.join(20 * interval)
@@ -692,7 +687,8 @@ class Benchmark(ABC):
                 for sequence_length in self.args.sequence_lengths:
                     if self.args.inference:
                         if self.args.memory:
-                            memory, inference_summary = self.inference_memory(model_name, batch_size, sequence_length)
+                            outputs = self.inference_memory(model_name, batch_size, sequence_length)
+                            memory, inference_summary = outputs if type(outputs) == tuple else (outputs, None)
                             inference_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                         if self.args.speed:
                             time = self.inference_speed(model_name, batch_size, sequence_length)
@@ -700,7 +696,8 @@ class Benchmark(ABC):
 
                     if self.args.training:
                         if self.args.memory:
-                            memory, train_summary = self.train_memory(model_name, batch_size, sequence_length)
+                            outputs = self.train_memory(model_name, batch_size, sequence_length)
+                            memory, train_summary = outputs if type(outputs) == tuple else (outputs, None)
                             train_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                         if self.args.speed:
                             time = self.train_speed(model_name, batch_size, sequence_length)

src/transformers/configuration_utils.py

@@ -170,6 +170,7 @@ class PretrainedConfig(object):
         self.torchscript = kwargs.pop("torchscript", False)  # Only used by PyTorch models
         self.use_bfloat16 = kwargs.pop("use_bfloat16", False)
         self.pruned_heads = kwargs.pop("pruned_heads", {})
+        self.gradient_checkpointing = kwargs.pop("gradient_checkpointing", False)
         self.tie_word_embeddings = kwargs.pop(
             "tie_word_embeddings", True
         )  # Whether input and output word embeddings should be tied for all MLM, LM and Seq2Seq models.

src/transformers/data/data_collator.py

@@ -114,6 +114,67 @@ class DataCollatorWithPadding:
         return batch
 
 
+@dataclass
+class DataCollatorForTokenClassification:
+    """
+    Data collator that will dynamically pad the inputs received, as well as the labels.
+
+    Args:
+        tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
+            The tokenizer used for encoding the data.
+        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        max_length (:obj:`int`, `optional`):
+            Maximum length of the returned list and optionally padding length (see above).
+        pad_to_multiple_of (:obj:`int`, `optional`):
+            If set will pad the sequence to a multiple of the provided value.
+
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+        label_pad_token_id (:obj:`int`, `optional`, defaults to -100):
+            The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    label_pad_token_id: int = -100
+
+    def __call__(self, features):
+        label_name = "label" if "label" in features[0].keys() else "labels"
+        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
+        batch = self.tokenizer.pad(
+            features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            # Conversion to tensors will fail if we have labels as they are not of the same length yet.
+            return_tensors="pt" if labels is None else None,
+        )
+
+        if labels is None:
+            return batch
+
+        sequence_length = torch.tensor(batch["input_ids"]).shape[1]
+        padding_side = self.tokenizer.padding_side
+        if padding_side == "right":
+            batch["labels"] = [label + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels]
+        else:
+            batch["labels"] = [[self.label_pad_token_id] * (sequence_length - len(label)) + label for label in labels]
+
+        batch = {k: torch.tensor(v, dtype=torch.int64) for k, v in batch.items()}
+        return batch
+
+
 @dataclass
 class DataCollatorForLanguageModeling:
     """

src/transformers/modeling_albert.py

@@ -600,7 +600,7 @@ class AlbertModel(AlbertPreTrainedModel):
             self.pooler = None
             self.pooler_activation = None
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -608,6 +608,9 @@ class AlbertModel(AlbertPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def get_layers(self):
+        return [layer for layer_group in self.encoder.albert_layer_groups for layer in layer_group.albert_layers]
+
     def _resize_token_embeddings(self, new_num_tokens):
         old_embeddings = self.embeddings.word_embeddings
         new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens)
@@ -650,11 +653,13 @@ class AlbertModel(AlbertPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -719,7 +724,7 @@ class AlbertForPreTraining(AlbertPreTrainedModel):
         self.predictions = AlbertMLMHead(config)
         self.sop_classifier = AlbertSOPHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.predictions.decoder
@@ -871,7 +876,7 @@ class AlbertForMaskedLM(AlbertPreTrainedModel):
         self.albert = AlbertModel(config, add_pooling_layer=False)
         self.predictions = AlbertMLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.predictions.decoder
@@ -965,7 +970,7 @@ class AlbertForSequenceClassification(AlbertPreTrainedModel):
         self.dropout = nn.Dropout(config.classifier_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ALBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1053,7 +1058,7 @@ class AlbertForTokenClassification(AlbertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ALBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1141,7 +1146,7 @@ class AlbertForQuestionAnswering(AlbertPreTrainedModel):
         self.albert = AlbertModel(config, add_pooling_layer=False)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ALBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1240,7 +1245,7 @@ class AlbertForMultipleChoice(AlbertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ALBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
     @add_code_sample_docstrings(

src/transformers/modeling_bart.py

@@ -842,7 +842,7 @@ class BartModel(PretrainedBartModel):
         self.encoder = BartEncoder(config, self.shared)
         self.decoder = BartDecoder(config, self.shared)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BART_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(
@@ -875,8 +875,10 @@ class BartModel(PretrainedBartModel):
         if decoder_input_ids is None:
             use_cache = False
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         use_cache = use_cache if use_cache is not None else self.config.use_cache

src/transformers/modeling_bert.py

@@ -455,32 +455,14 @@ class BertEncoder(nn.Module):
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
             layer_head_mask = head_mask[i] if head_mask is not None else None
-
-            if getattr(self.config, "gradient_checkpointing", False):
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs, output_attentions)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(layer_module),
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                )
-            else:
-                layer_outputs = layer_module(
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                    output_attentions,
-                )
+            layer_outputs = layer_module(
+                hidden_states,
+                attention_mask,
+                layer_head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
+            )
             hidden_states = layer_outputs[0]
             if output_attentions:
                 all_attentions = all_attentions + (layer_outputs[1],)
@@ -732,7 +714,7 @@ class BertModel(BertPreTrainedModel):
 
         self.pooler = BertPooler(config) if add_pooling_layer else None
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -740,6 +722,9 @@ class BertModel(BertPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def get_layers(self):
+        return self.encoder.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -780,11 +765,13 @@ class BertModel(BertPreTrainedModel):
             - 1 for tokens that are **not masked**,
             - 0 for tokens that are **masked**.
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -865,7 +852,7 @@ class BertForPreTraining(BertPreTrainedModel):
         self.bert = BertModel(config)
         self.cls = BertPreTrainingHeads(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.cls.predictions.decoder
@@ -978,7 +965,7 @@ class BertLMHeadModel(BertPreTrainedModel):
         self.bert = BertModel(config, add_pooling_layer=False)
         self.cls = BertOnlyMLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.cls.predictions.decoder
@@ -1098,7 +1085,7 @@ class BertForMaskedLM(BertPreTrainedModel):
         self.bert = BertModel(config, add_pooling_layer=False)
         self.cls = BertOnlyMLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.cls.predictions.decoder
@@ -1204,7 +1191,7 @@ class BertForNextSentencePrediction(BertPreTrainedModel):
         self.bert = BertModel(config)
         self.cls = BertOnlyNSPHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @replace_return_docstrings(output_type=NextSentencePredictorOutput, config_class=_CONFIG_FOR_DOC)
@@ -1298,7 +1285,7 @@ class BertForSequenceClassification(BertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1382,7 +1369,7 @@ class BertForMultipleChoice(BertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
     @add_code_sample_docstrings(
@@ -1477,7 +1464,7 @@ class BertForTokenClassification(BertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1567,7 +1554,7 @@ class BertForQuestionAnswering(BertPreTrainedModel):
         self.bert = BertModel(config, add_pooling_layer=False)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(

src/transformers/modeling_bert_generation.py

@@ -277,7 +277,7 @@ class BertGenerationEncoder(BertGenerationPreTrainedModel):
         self.embeddings = BertGenerationEmbeddings(config)
         self.encoder = BertEncoder(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -285,6 +285,9 @@ class BertGenerationEncoder(BertGenerationPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def get_layers(self):
+        return self.encoder.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -322,11 +325,13 @@ class BertGenerationEncoder(BertGenerationPreTrainedModel):
             the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``: ``1`` for
             tokens that are NOT MASKED, ``0`` for MASKED tokens.
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -416,7 +421,7 @@ class BertGenerationDecoder(BertGenerationPreTrainedModel):
         self.bert = BertGenerationEncoder(config)
         self.lm_head = BertGenerationOnlyLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head.decoder

src/transformers/modeling_ctrl.py

@@ -149,7 +149,7 @@ class MultiHeadAttention(torch.nn.Module):
             k = torch.cat((past_key, k), dim=-2)
             v = torch.cat((past_value, v), dim=-2)
 
-        if use_cache is True:
+        if use_cache:
             present = torch.stack((k, v))
         else:
             present = (None,)
@@ -334,7 +334,7 @@ class CTRLModel(CTRLPreTrainedModel):
         )
         self.layernorm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.w
@@ -342,6 +342,9 @@ class CTRLModel(CTRLPreTrainedModel):
     def set_input_embeddings(self, new_embeddings):
         self.w = new_embeddings
 
+    def get_layers(self):
+        return self.h
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
@@ -379,9 +382,11 @@ class CTRLModel(CTRLPreTrainedModel):
             past_key_values = kwargs.pop("past")
         assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        use_cache = torch.tensor(use_cache if use_cache is not None else self.config.use_cache)
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -459,17 +464,18 @@ class CTRLModel(CTRLPreTrainedModel):
         for i, (h, layer_past) in enumerate(zip(self.h, past_key_values)):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
+            layer_head_mask = head_mask[i] if head_mask is not None else None
             outputs = h(
                 hidden_states,
                 mask,
-                layer_past=layer_past,
-                attention_mask=attention_mask,
-                head_mask=head_mask[i],
-                use_cache=use_cache,
-                output_attentions=output_attentions,
+                layer_past,
+                attention_mask,
+                layer_head_mask,
+                use_cache,
+                output_attentions,
             )
             hidden_states, present = outputs[:2]
-            if use_cache is True:
+            if use_cache:
                 presents = presents + (present,)
 
             if output_attentions:
@@ -509,7 +515,7 @@ class CTRLLMHeadModel(CTRLPreTrainedModel):
         self.transformer = CTRLModel(config)
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head

src/transformers/modeling_deberta.py

@@ -406,9 +406,9 @@ class DebertaEncoder(nn.Module):
                 next_kv,
                 attention_mask,
                 output_attentions,
-                query_states=query_states,
-                relative_pos=relative_pos,
-                rel_embeddings=rel_embeddings,
+                query_states,
+                relative_pos,
+                rel_embeddings,
             )
             if output_attentions:
                 hidden_states, att_m = hidden_states
@@ -843,7 +843,7 @@ class DebertaModel(DebertaPreTrainedModel):
         self.encoder = DebertaEncoder(config)
         self.z_steps = 0
         self.config = config
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -851,6 +851,9 @@ class DebertaModel(DebertaPreTrainedModel):
     def set_input_embeddings(self, new_embeddings):
         self.embeddings.word_embeddings = new_embeddings
 
+    def get_layers(self):
+        return self.encoder.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -876,11 +879,13 @@ class DebertaModel(DebertaPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -968,7 +973,7 @@ class DebertaForSequenceClassification(DebertaPreTrainedModel):
         drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
         self.dropout = StableDropout(drop_out)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.deberta.get_input_embeddings()

src/transformers/modeling_distilbert.py

@@ -306,9 +306,8 @@ class Transformer(nn.Module):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_state,)
 
-            layer_outputs = layer_module(
-                x=hidden_state, attn_mask=attn_mask, head_mask=head_mask[i], output_attentions=output_attentions
-            )
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            layer_outputs = layer_module(hidden_state, attn_mask, layer_head_mask, output_attentions)
             hidden_state = layer_outputs[-1]
 
             if output_attentions:
@@ -420,7 +419,7 @@ class DistilBertModel(DistilBertPreTrainedModel):
         self.embeddings = Embeddings(config)  # Embeddings
         self.transformer = Transformer(config)  # Encoder
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -428,6 +427,9 @@ class DistilBertModel(DistilBertPreTrainedModel):
     def set_input_embeddings(self, new_embeddings):
         self.embeddings.word_embeddings = new_embeddings
 
+    def get_layers(self):
+        return self.transformer.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -454,11 +456,13 @@ class DistilBertModel(DistilBertPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -502,7 +506,7 @@ class DistilBertForMaskedLM(DistilBertPreTrainedModel):
         self.vocab_layer_norm = nn.LayerNorm(config.dim, eps=1e-12)
         self.vocab_projector = nn.Linear(config.dim, config.vocab_size)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
         self.mlm_loss_fct = nn.CrossEntropyLoss()
 
@@ -593,7 +597,7 @@ class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
         self.classifier = nn.Linear(config.dim, config.num_labels)
         self.dropout = nn.Dropout(config.seq_classif_dropout)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices"))
     @add_code_sample_docstrings(
@@ -674,7 +678,7 @@ class DistilBertForQuestionAnswering(DistilBertPreTrainedModel):
         assert config.num_labels == 2
         self.dropout = nn.Dropout(config.qa_dropout)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices"))
     @add_code_sample_docstrings(
@@ -770,7 +774,7 @@ class DistilBertForTokenClassification(DistilBertPreTrainedModel):
         self.dropout = nn.Dropout(config.dropout)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(DISTILBERT_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(
@@ -854,7 +858,7 @@ class DistilBertForMultipleChoice(DistilBertPreTrainedModel):
         self.classifier = nn.Linear(config.dim, 1)
         self.dropout = nn.Dropout(config.seq_classif_dropout)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(
         DISTILBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")

src/transformers/modeling_dpr.py

@@ -157,7 +157,7 @@ class DPREncoder(PreTrainedModel):
         self.projection_dim = config.projection_dim
         if self.projection_dim > 0:
             self.encode_proj = nn.Linear(self.bert_model.config.hidden_size, config.projection_dim)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def forward(
         self,
@@ -199,8 +199,8 @@ class DPREncoder(PreTrainedModel):
             return self.encode_proj.out_features
         return self.bert_model.config.hidden_size
 
-    def init_weights(self):
-        self.bert_model.init_weights()
+    def init_weights_and_layers(self):
+        self.bert_model.init_weights_and_layers()
         if self.projection_dim > 0:
             self.encode_proj.apply(self.bert_model._init_weights)
 
@@ -214,7 +214,7 @@ class DPRSpanPredictor(PreTrainedModel):
         self.encoder = DPREncoder(config)
         self.qa_outputs = nn.Linear(self.encoder.embeddings_size, 2)
         self.qa_classifier = nn.Linear(self.encoder.embeddings_size, 1)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def forward(
         self,
@@ -261,8 +261,8 @@ class DPRSpanPredictor(PreTrainedModel):
             attentions=outputs.attentions,
         )
 
-    def init_weights(self):
-        self.encoder.init_weights()
+    def init_weights_and_layers(self):
+        self.encoder.init_weights_and_layers()
 
 
 ##################
@@ -281,8 +281,8 @@ class DPRPretrainedContextEncoder(PreTrainedModel):
     base_model_prefix = "ctx_encoder"
     authorized_missing_keys = [r"position_ids"]
 
-    def init_weights(self):
-        self.ctx_encoder.init_weights()
+    def init_weights_and_layers(self):
+        self.ctx_encoder.init_weights_and_layers()
 
 
 class DPRPretrainedQuestionEncoder(PreTrainedModel):
@@ -296,8 +296,8 @@ class DPRPretrainedQuestionEncoder(PreTrainedModel):
     base_model_prefix = "question_encoder"
     authorized_missing_keys = [r"position_ids"]
 
-    def init_weights(self):
-        self.question_encoder.init_weights()
+    def init_weights_and_layers(self):
+        self.question_encoder.init_weights_and_layers()
 
 
 class DPRPretrainedReader(PreTrainedModel):
@@ -311,8 +311,8 @@ class DPRPretrainedReader(PreTrainedModel):
     base_model_prefix = "span_predictor"
     authorized_missing_keys = [r"position_ids"]
 
-    def init_weights(self):
-        self.span_predictor.encoder.init_weights()
+    def init_weights_and_layers(self):
+        self.span_predictor.encoder.init_weights_and_layers()
         self.span_predictor.qa_classifier.apply(self.span_predictor.encoder.bert_model._init_weights)
         self.span_predictor.qa_outputs.apply(self.span_predictor.encoder.bert_model._init_weights)
 
@@ -434,7 +434,7 @@ class DPRContextEncoder(DPRPretrainedContextEncoder):
         super().__init__(config)
         self.config = config
         self.ctx_encoder = DPREncoder(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(DPR_ENCODERS_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=DPRContextEncoderOutput, config_class=_CONFIG_FOR_DOC)
@@ -460,8 +460,10 @@ class DPRContextEncoder(DPRPretrainedContextEncoder):
             >>> embeddings = model(input_ids).pooler_output
         """
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -512,7 +514,7 @@ class DPRQuestionEncoder(DPRPretrainedQuestionEncoder):
         super().__init__(config)
         self.config = config
         self.question_encoder = DPREncoder(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(DPR_ENCODERS_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=DPRQuestionEncoderOutput, config_class=_CONFIG_FOR_DOC)
@@ -537,8 +539,10 @@ class DPRQuestionEncoder(DPRPretrainedQuestionEncoder):
             >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors='pt')["input_ids"]
             >>> embeddings = model(input_ids).pooler_output
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -589,7 +593,7 @@ class DPRReader(DPRPretrainedReader):
         super().__init__(config)
         self.config = config
         self.span_predictor = DPRSpanPredictor(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(DPR_READER_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=DPRReaderOutput, config_class=_CONFIG_FOR_DOC)
@@ -622,8 +626,10 @@ class DPRReader(DPRPretrainedReader):
             >>> relevance_logits = outputs.relevance_logits
 
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict

src/transformers/modeling_electra.py

@@ -451,32 +451,14 @@ class ElectraEncoder(nn.Module):
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
             layer_head_mask = head_mask[i] if head_mask is not None else None
-
-            if getattr(self.config, "gradient_checkpointing", False):
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs, output_attentions)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(layer_module),
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                )
-            else:
-                layer_outputs = layer_module(
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                    output_attentions,
-                )
+            layer_outputs = layer_module(
+                hidden_states,
+                attention_mask,
+                layer_head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
+            )
             hidden_states = layer_outputs[0]
             if output_attentions:
                 all_attentions = all_attentions + (layer_outputs[1],)
@@ -677,7 +659,7 @@ class ElectraModel(ElectraPreTrainedModel):
 
         self.encoder = ElectraEncoder(config)
         self.config = config
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -685,6 +667,9 @@ class ElectraModel(ElectraPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def get_layers(self):
+        return self.encoder.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -712,11 +697,13 @@ class ElectraModel(ElectraPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -789,7 +776,7 @@ class ElectraForSequenceClassification(ElectraPreTrainedModel):
         self.electra = ElectraModel(config)
         self.classifier = ElectraClassificationHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -871,7 +858,7 @@ class ElectraForPreTraining(ElectraPreTrainedModel):
 
         self.electra = ElectraModel(config)
         self.discriminator_predictions = ElectraDiscriminatorPredictions(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @replace_return_docstrings(output_type=ElectraForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
@@ -966,7 +953,7 @@ class ElectraForMaskedLM(ElectraPreTrainedModel):
         self.generator_predictions = ElectraGeneratorPredictions(config)
 
         self.generator_lm_head = nn.Linear(config.embedding_size, config.vocab_size)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.generator_lm_head
@@ -1058,7 +1045,7 @@ class ElectraForTokenClassification(ElectraPreTrainedModel):
         self.electra = ElectraModel(config)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1145,7 +1132,7 @@ class ElectraForQuestionAnswering(ElectraPreTrainedModel):
         self.electra = ElectraModel(config)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1246,7 +1233,7 @@ class ElectraForMultipleChoice(ElectraPreTrainedModel):
         self.sequence_summary = SequenceSummary(config)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ELECTRA_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
     @add_code_sample_docstrings(

src/transformers/modeling_flaubert.py

@@ -162,8 +162,10 @@ class FlaubertModel(XLMModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -324,7 +326,7 @@ class FlaubertWithLMHeadModel(XLMWithLMHeadModel):
     def __init__(self, config):
         super().__init__(config)
         self.transformer = FlaubertModel(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
 
 @add_start_docstrings(
@@ -345,7 +347,7 @@ class FlaubertForSequenceClassification(XLMForSequenceClassification):
     def __init__(self, config):
         super().__init__(config)
         self.transformer = FlaubertModel(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
 
 @add_start_docstrings(
@@ -366,7 +368,7 @@ class FlaubertForTokenClassification(XLMForTokenClassification):
     def __init__(self, config):
         super().__init__(config)
         self.transformer = FlaubertModel(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
 
 @add_start_docstrings(
@@ -387,7 +389,7 @@ class FlaubertForQuestionAnsweringSimple(XLMForQuestionAnsweringSimple):
     def __init__(self, config):
         super().__init__(config)
         self.transformer = FlaubertModel(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
 
 @add_start_docstrings(
@@ -408,7 +410,7 @@ class FlaubertForQuestionAnswering(XLMForQuestionAnswering):
     def __init__(self, config):
         super().__init__(config)
         self.transformer = FlaubertModel(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
 
 @add_start_docstrings(
@@ -429,4 +431,4 @@ class FlaubertForMultipleChoice(XLMForMultipleChoice):
     def __init__(self, config):
         super().__init__(config)
         self.transformer = FlaubertModel(config)
-        self.init_weights()
+        self.init_weights_and_layers()

src/transformers/modeling_fsmt.py

@@ -897,7 +897,7 @@ class FSMTModel(PretrainedFSMTModel):
         self.encoder = FSMTEncoder(config, encoder_embed_tokens)
         self.decoder = FSMTDecoder(config, decoder_embed_tokens)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(FSMT_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(
@@ -930,8 +930,10 @@ class FSMTModel(PretrainedFSMTModel):
         if decoder_input_ids is None:
             use_cache = False
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         use_cache = use_cache if use_cache is not None else self.config.use_cache

src/transformers/modeling_funnel.py

@@ -648,7 +648,7 @@ class FunnelEncoder(nn.Module):
                         key = value = hidden if self.config.pool_q_only else pooled_hidden
                     else:
                         query = key = value = hidden
-                    layer_output = layer(query, key, value, attention_inputs, output_attentions=output_attentions)
+                    layer_output = layer(query, key, value, attention_inputs, output_attentions)
                     hidden = layer_output[0]
                     if do_pooling:
                         attention_inputs = self.attention_structure.post_attention_pooling(attention_inputs)
@@ -898,7 +898,7 @@ class FunnelBaseModel(FunnelPreTrainedModel):
         self.embeddings = FunnelEmbeddings(config)
         self.encoder = FunnelEncoder(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -906,6 +906,9 @@ class FunnelBaseModel(FunnelPreTrainedModel):
     def set_input_embeddings(self, new_embeddings):
         self.embeddings.word_embeddings = new_embeddings
 
+    def get_layers(self):
+        return [layer for block in self.encoder.blocks for layer in block]
+
     @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
         tokenizer_class=_TOKENIZER_FOR_DOC,
@@ -925,8 +928,10 @@ class FunnelBaseModel(FunnelPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -975,7 +980,7 @@ class FunnelModel(FunnelPreTrainedModel):
         self.encoder = FunnelEncoder(config)
         self.decoder = FunnelDecoder(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -1000,7 +1005,9 @@ class FunnelModel(FunnelPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
@@ -1080,7 +1087,7 @@ class FunnelForPreTraining(FunnelPreTrainedModel):
 
         self.funnel = FunnelModel(config)
         self.discriminator_predictions = FunnelDiscriminatorPredictions(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @replace_return_docstrings(output_type=FunnelForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
@@ -1162,7 +1169,7 @@ class FunnelForMaskedLM(FunnelPreTrainedModel):
         self.funnel = FunnelModel(config)
         self.lm_head = nn.Linear(config.d_model, config.vocab_size)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -1238,7 +1245,7 @@ class FunnelForSequenceClassification(FunnelPreTrainedModel):
 
         self.funnel = FunnelBaseModel(config)
         self.classifier = FunnelClassificationHead(config, config.num_labels)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1315,7 +1322,7 @@ class FunnelForMultipleChoice(FunnelPreTrainedModel):
 
         self.funnel = FunnelBaseModel(config)
         self.classifier = FunnelClassificationHead(config, 1)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
     @add_code_sample_docstrings(
@@ -1401,7 +1408,7 @@ class FunnelForTokenClassification(FunnelPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1483,7 +1490,7 @@ class FunnelForQuestionAnswering(FunnelPreTrainedModel):
         self.funnel = FunnelModel(config)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(FUNNEL_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(

src/transformers/modeling_gpt2.py

@@ -227,7 +227,7 @@ class Attention(nn.Module):
             key = torch.cat((past_key, key), dim=-1)
             value = torch.cat((past_value, value), dim=-2)
 
-        if use_cache is True:
+        if use_cache:
             present = torch.stack((key.transpose(-2, -1), value))  # transpose to have same shapes for stacking
         else:
             present = (None,)
@@ -317,7 +317,7 @@ class Block(nn.Module):
         # residual connection
         hidden_states = hidden_states + feed_forward_hidden_states
 
-        outputs = [hidden_states] + outputs
+        outputs = (hidden_states,) + tuple(outputs)
         return outputs  # hidden_states, present, (cross_attentions, attentions)
 
 
@@ -487,7 +487,7 @@ class GPT2Model(GPT2PreTrainedModel):
         self.h = nn.ModuleList([Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)])
         self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.wte
@@ -495,6 +495,9 @@ class GPT2Model(GPT2PreTrainedModel):
     def set_input_embeddings(self, new_embeddings):
         self.wte = new_embeddings
 
+    def get_layers(self):
+        return self.h
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
@@ -534,12 +537,14 @@ class GPT2Model(GPT2PreTrainedModel):
             past_key_values = kwargs.pop("past")
         assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        use_cache = torch.tensor(use_cache if use_cache is not None else self.config.use_cache)
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -624,38 +629,20 @@ class GPT2Model(GPT2PreTrainedModel):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
 
-            if getattr(self.config, "gradient_checkpointing", False):
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # checkpointing only works with tuple returns, not with lists
-                        return tuple(output for output in module(*inputs, use_cache, output_attentions))
-
-                    return custom_forward
-
-                outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    hidden_states,
-                    layer_past,
-                    attention_mask,
-                    head_mask[i],
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                )
-            else:
-                outputs = block(
-                    hidden_states,
-                    layer_past=layer_past,
-                    attention_mask=attention_mask,
-                    head_mask=head_mask[i],
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=encoder_attention_mask,
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                )
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            outputs = block(
+                hidden_states,
+                layer_past,
+                attention_mask,
+                layer_head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                use_cache,
+                output_attentions,
+            )
 
             hidden_states, present = outputs[:2]
-            if use_cache is True:
+            if use_cache:
                 presents = presents + (present,)
 
             if output_attentions:
@@ -664,6 +651,7 @@ class GPT2Model(GPT2PreTrainedModel):
         hidden_states = self.ln_f(hidden_states)
 
         hidden_states = hidden_states.view(*output_shape)
+
         # Add last hidden state
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)
@@ -694,7 +682,7 @@ class GPT2LMHeadModel(GPT2PreTrainedModel):
         self.transformer = GPT2Model(config)
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -821,7 +809,7 @@ class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
         self.multiple_choice_head = SequenceSummary(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -985,7 +973,7 @@ class GPT2ForSequenceClassification(GPT2PreTrainedModel):
         self.transformer = GPT2Model(config)
         self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(

src/transformers/modeling_layoutlm.py

@@ -380,32 +380,14 @@ class LayoutLMEncoder(nn.Module):
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
             layer_head_mask = head_mask[i] if head_mask is not None else None
-
-            if getattr(self.config, "gradient_checkpointing", False):
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs, output_attentions)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(layer_module),
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                )
-            else:
-                layer_outputs = layer_module(
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                    output_attentions,
-                )
+            layer_outputs = layer_module(
+                hidden_states,
+                attention_mask,
+                layer_head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
+            )
             hidden_states = layer_outputs[0]
             if output_attentions:
                 all_attentions = all_attentions + (layer_outputs[1],)
@@ -591,7 +573,7 @@ class LayoutLMModel(LayoutLMPreTrainedModel):
         self.encoder = LayoutLMEncoder(config)
         self.pooler = LayoutLMPooler(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -599,6 +581,9 @@ class LayoutLMModel(LayoutLMPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def get_layers(self):
+        return self.encoder.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -655,11 +640,13 @@ class LayoutLMModel(LayoutLMPreTrainedModel):
         return_dict (bool, optional):
             If set to True, the model will return a ModelOutput instead of a plain tuple.
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -736,7 +723,7 @@ class LayoutLMForMaskedLM(LayoutLMPreTrainedModel):
         self.layoutlm = LayoutLMModel(config)
         self.cls = LayoutLMOnlyMLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.layoutlm.embeddings.word_embeddings
@@ -827,7 +814,7 @@ class LayoutLMForTokenClassification(LayoutLMPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.layoutlm.embeddings.word_embeddings

src/transformers/modeling_longformer.py

@@ -1106,7 +1106,7 @@ class LongformerModel(LongformerPreTrainedModel):
         self.encoder = LongformerEncoder(config)
         self.pooler = LongformerPooler(config) if add_pooling_layer else None
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -1223,8 +1223,10 @@ class LongformerModel(LongformerPreTrainedModel):
             >>> pooled_output = outputs.pooler_output
         """
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -1303,7 +1305,7 @@ class LongformerForMaskedLM(LongformerPreTrainedModel):
         self.longformer = LongformerModel(config, add_pooling_layer=False)
         self.lm_head = LongformerLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head.decoder
@@ -1410,7 +1412,7 @@ class LongformerForSequenceClassification(LongformerPreTrainedModel):
         self.longformer = LongformerModel(config, add_pooling_layer=False)
         self.classifier = LongformerClassificationHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1519,7 +1521,7 @@ class LongformerForQuestionAnswering(LongformerPreTrainedModel):
         self.longformer = LongformerModel(config, add_pooling_layer=False)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @replace_return_docstrings(output_type=QuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC)
@@ -1653,7 +1655,7 @@ class LongformerForTokenClassification(LongformerPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(LONGFORMER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1740,7 +1742,7 @@ class LongformerForMultipleChoice(LongformerPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(
         LONGFORMER_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")

src/transformers/modeling_lxmert.py

@@ -885,7 +885,7 @@ class LxmertModel(LxmertPreTrainedModel):
         self.embeddings = LxmertEmbeddings(config)
         self.encoder = LxmertEncoder(config)
         self.pooler = LxmertPooler(config)
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -914,8 +914,10 @@ class LxmertModel(LxmertPreTrainedModel):
         return_dict=None,
     ):
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -1040,7 +1042,7 @@ class LxmertForPreTraining(LxmertPreTrainedModel):
             self.answer_head = LxmertVisualAnswerHead(config, self.num_qa_labels)
 
         # Weight initialization
-        self.init_weights()
+        self.init_weights_and_layers()
 
         # Loss functions
         self.loss_fcts = {
@@ -1285,7 +1287,7 @@ class LxmertForQuestionAnswering(LxmertPreTrainedModel):
         self.answer_head = LxmertVisualAnswerHead(config, self.num_qa_labels)
 
         # Weight initialization
-        self.init_weights()
+        self.init_weights_and_layers()
 
         # Loss function
         self.loss = CrossEntropyLoss()

src/transformers/modeling_mmbt.py

@@ -218,8 +218,10 @@ class MMBTModel(nn.Module, ModuleUtilsMixin):
             encoder = ImageEncoder(args)
             mmbt = MMBTModel(config, transformer, encoder)
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict

src/transformers/modeling_mobilebert.py

@@ -566,10 +566,11 @@ class MobileBertEncoder(nn.Module):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
+            layer_head_mask = head_mask[i] if head_mask is not None else None
             layer_outputs = layer_module(
                 hidden_states,
                 attention_mask,
-                head_mask[i],
+                layer_head_mask,
                 encoder_hidden_states,
                 encoder_attention_mask,
                 output_attentions,
@@ -821,7 +822,7 @@ class MobileBertModel(MobileBertPreTrainedModel):
 
         self.pooler = MobileBertPooler(config) if add_pooling_layer else None
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -829,6 +830,9 @@ class MobileBertModel(MobileBertPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def get_layers(self):
+        return self.encoder.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -858,11 +862,13 @@ class MobileBertModel(MobileBertPreTrainedModel):
         output_attentions=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -944,7 +950,7 @@ class MobileBertForPreTraining(MobileBertPreTrainedModel):
         self.mobilebert = MobileBertModel(config)
         self.cls = MobileBertPreTrainingHeads(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.cls.predictions.decoder
@@ -1062,7 +1068,7 @@ class MobileBertForMaskedLM(MobileBertPreTrainedModel):
         self.cls = MobileBertOnlyMLMHead(config)
         self.config = config
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.cls.predictions.decoder
@@ -1182,7 +1188,7 @@ class MobileBertForNextSentencePrediction(MobileBertPreTrainedModel):
         self.mobilebert = MobileBertModel(config)
         self.cls = MobileBertOnlyNSPHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MOBILEBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @replace_return_docstrings(output_type=NextSentencePredictorOutput, config_class=_CONFIG_FOR_DOC)
@@ -1274,7 +1280,7 @@ class MobileBertForSequenceClassification(MobileBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, self.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MOBILEBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1359,7 +1365,7 @@ class MobileBertForQuestionAnswering(MobileBertPreTrainedModel):
         self.mobilebert = MobileBertModel(config, add_pooling_layer=False)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MOBILEBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1458,7 +1464,7 @@ class MobileBertForMultipleChoice(MobileBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(
         MOBILEBERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")
@@ -1555,7 +1561,7 @@ class MobileBertForTokenClassification(MobileBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(MOBILEBERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(

src/transformers/modeling_openai.py

@@ -268,7 +268,7 @@ class Block(nn.Module):
         h = self.ln_2(n + m)
 
         outputs = [h] + attn_outputs[1:]
-        return outputs
+        return tuple(outputs)
 
 
 class OpenAIGPTPreTrainedModel(PreTrainedModel):
@@ -412,7 +412,7 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
         self.h = nn.ModuleList([Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)])
 
         self.register_buffer("position_ids", torch.arange(config.n_positions))
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.tokens_embed
@@ -420,6 +420,9 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
     def set_input_embeddings(self, new_embeddings):
         self.tokens_embed = new_embeddings
 
+    def get_layers(self):
+        return self.h
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
@@ -446,11 +449,13 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -505,7 +510,8 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
 
-            outputs = block(hidden_states, attention_mask, head_mask[i], output_attentions=output_attentions)
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            outputs = block(hidden_states, attention_mask, layer_head_mask, output_attentions)
             hidden_states = outputs[0]
             if output_attentions:
                 all_attentions = all_attentions + (outputs[1],)
@@ -538,7 +544,7 @@ class OpenAIGPTLMHeadModel(OpenAIGPTPreTrainedModel):
         self.transformer = OpenAIGPTModel(config)
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -624,7 +630,7 @@ class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
         self.multiple_choice_head = SequenceSummary(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -752,7 +758,7 @@ class OpenAIGPTForSequenceClassification(OpenAIGPTPreTrainedModel):
         self.transformer = OpenAIGPTModel(config)
         self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(OPENAI_GPT_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(

src/transformers/modeling_prophetnet.py

@@ -1135,7 +1135,7 @@ class ProphetNetEncoder(ProphetNetPreTrainedModel):
 
         self.layers = nn.ModuleList([ProphetNetEncoderLayer(config) for _ in range(config.num_encoder_layers)])
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.word_embeddings
@@ -1170,8 +1170,10 @@ class ProphetNetEncoder(ProphetNetPreTrainedModel):
             >>> last_hidden_states = outputs.last_hidden_state
         """
 
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -1251,7 +1253,7 @@ class ProphetNetDecoder(ProphetNetPreTrainedModel):
         self.layers = nn.ModuleList([ProphetNetDecoderLayer(config) for _ in range(config.num_decoder_layers)])
         self.embeddings_layer_norm = ProphetNetLayerNorm(config.hidden_size)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.word_embeddings
@@ -1310,8 +1312,10 @@ class ProphetNetDecoder(ProphetNetPreTrainedModel):
             >>> last_hidden_states = outputs.last_hidden_state
         """
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -1559,7 +1563,7 @@ class ProphetNetModel(ProphetNetPreTrainedModel):
         decoder_config.is_encoder_decoder = False
         self.decoder = ProphetNetDecoder(decoder_config, self.word_embeddings)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.word_embeddings
@@ -1611,8 +1615,10 @@ class ProphetNetModel(ProphetNetPreTrainedModel):
         """
 
         use_cache == use_cache if use_cache is not None else self.config.use_cache
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -1671,7 +1677,7 @@ class ProphetNetForConditionalGeneration(ProphetNetPreTrainedModel):
 
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -1859,7 +1865,7 @@ class ProphetNetForCausalLM(ProphetNetPreTrainedModel):
 
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.decoder.word_embeddings

src/transformers/modeling_rag.py

@@ -542,8 +542,10 @@ class RagModel(RagPreTrainedModel):
         """
         n_docs = n_docs if n_docs is not None else self.config.n_docs
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         output_retrieved = output_retrieved if output_retrieved is not None else self.config.output_retrieved

src/transformers/modeling_reformer.py

@@ -1977,7 +1977,7 @@ class ReformerModel(ReformerPreTrainedModel):
         self.embeddings = ReformerEmbeddings(config)
         self.encoder = ReformerEncoder(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -2015,8 +2015,10 @@ class ReformerModel(ReformerPreTrainedModel):
         return_dict=None,
     ):
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -2192,7 +2194,7 @@ class ReformerModelWithLMHead(ReformerPreTrainedModel):
         self.reformer = ReformerModel(config)
         self.lm_head = ReformerOnlyLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head.decoder
@@ -2304,7 +2306,7 @@ class ReformerForMaskedLM(ReformerPreTrainedModel):
         self.reformer = ReformerModel(config)
         self.lm_head = ReformerOnlyLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head.decoder
@@ -2387,7 +2389,7 @@ class ReformerForSequenceClassification(ReformerPreTrainedModel):
         if config.is_decoder is True:
             logger.warning("You might want to disable causal masking for sequence classification")
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(REFORMER_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(
@@ -2489,7 +2491,7 @@ class ReformerForQuestionAnswering(ReformerPreTrainedModel):
         # 2 * config.hidden_size because we use reversible residual layers
         self.qa_outputs = nn.Linear(2 * config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(REFORMER_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(

src/transformers/modeling_retribert.py

@@ -95,7 +95,7 @@ class RetriBertModel(RetriBertPreTrainedModel):
 
         self.ce_loss = nn.CrossEntropyLoss(reduction="mean")
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def embed_sentences_checkpointed(
         self,

src/transformers/modeling_roberta.py

@@ -399,32 +399,14 @@ class RobertaEncoder(nn.Module):
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
             layer_head_mask = head_mask[i] if head_mask is not None else None
-
-            if getattr(self.config, "gradient_checkpointing", False):
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return module(*inputs, output_attentions)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(layer_module),
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                )
-            else:
-                layer_outputs = layer_module(
-                    hidden_states,
-                    attention_mask,
-                    layer_head_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                    output_attentions,
-                )
+            layer_outputs = layer_module(
+                hidden_states,
+                attention_mask,
+                layer_head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
+            )
             hidden_states = layer_outputs[0]
             if output_attentions:
                 all_attentions = all_attentions + (layer_outputs[1],)
@@ -579,7 +561,7 @@ class RobertaModel(RobertaPreTrainedModel):
 
         self.pooler = RobertaPooler(config) if add_pooling_layer else None
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -587,6 +569,9 @@ class RobertaModel(RobertaPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def get_layers(self):
+        return self.encoder.layer
+
     def _prune_heads(self, heads_to_prune):
         """
         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
@@ -626,11 +611,13 @@ class RobertaModel(RobertaPreTrainedModel):
             the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``: ``1`` for
             tokens that are NOT MASKED, ``0`` for MASKED tokens.
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = torch.tensor(return_dict if return_dict is not None else self.config.use_return_dict)
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -713,7 +700,7 @@ class RobertaForCausalLM(RobertaPreTrainedModel):
         self.roberta = RobertaModel(config, add_pooling_layer=False)
         self.lm_head = RobertaLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head.decoder
@@ -833,7 +820,7 @@ class RobertaForMaskedLM(RobertaPreTrainedModel):
         self.roberta = RobertaModel(config, add_pooling_layer=False)
         self.lm_head = RobertaLMHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head.decoder
@@ -954,7 +941,7 @@ class RobertaForSequenceClassification(RobertaPreTrainedModel):
         self.roberta = RobertaModel(config, add_pooling_layer=False)
         self.classifier = RobertaClassificationHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1037,7 +1024,7 @@ class RobertaForMultipleChoice(RobertaPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
     @add_code_sample_docstrings(
@@ -1131,7 +1118,7 @@ class RobertaForTokenClassification(RobertaPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1240,7 +1227,7 @@ class RobertaForQuestionAnswering(RobertaPreTrainedModel):
         self.roberta = RobertaModel(config, add_pooling_layer=False)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(

src/transformers/modeling_squeezebert.py

@@ -502,7 +502,7 @@ class SqueezeBertModel(SqueezeBertPreTrainedModel):
         self.encoder = SqueezeBertEncoder(config)
         self.pooler = SqueezeBertPooler(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -537,8 +537,10 @@ class SqueezeBertModel(SqueezeBertPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -600,7 +602,7 @@ class SqueezeBertForMaskedLM(SqueezeBertPreTrainedModel):
         self.transformer = SqueezeBertModel(config)
         self.lm_head = nn.Linear(config.embedding_size, config.vocab_size)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -681,7 +683,7 @@ class SqueezeBertForSequenceClassification(SqueezeBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(SQUEEZEBERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
     @add_code_sample_docstrings(
@@ -765,7 +767,7 @@ class SqueezeBertForMultipleChoice(SqueezeBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(
         SQUEEZEBERT_INPUTS_DOCSTRING.format("(batch_size, num_choices, sequence_length)")
@@ -859,7 +861,7 @@ class SqueezeBertForTokenClassification(SqueezeBertPreTrainedModel):
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(SQUEEZEBERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
     @add_code_sample_docstrings(
@@ -946,7 +948,7 @@ class SqueezeBertForQuestionAnswering(SqueezeBertPreTrainedModel):
         self.transformer = SqueezeBertModel(config)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(SQUEEZEBERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
     @add_code_sample_docstrings(

src/transformers/modeling_t5.py

@@ -657,7 +657,7 @@ class T5Stack(T5PreTrainedModel):
         self.final_layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon)
         self.dropout = nn.Dropout(config.dropout_rate)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.embed_tokens
@@ -684,8 +684,10 @@ class T5Stack(T5PreTrainedModel):
     ):
 
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -919,7 +921,7 @@ class T5Model(T5PreTrainedModel):
         decoder_config.num_layers = config.num_decoder_layers
         self.decoder = T5Stack(decoder_config, self.shared)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.shared
@@ -1067,7 +1069,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
 
         self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.shared

src/transformers/modeling_transfo_xl.py

@@ -785,7 +785,7 @@ class TransfoXLModel(TransfoXLPreTrainedModel):
         else:  # learnable embeddings and absolute embeddings
             raise NotImplementedError  # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.word_emb
@@ -852,8 +852,10 @@ class TransfoXLModel(TransfoXLPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -987,7 +989,7 @@ class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
             config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val
         )
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def tie_weights(self):
         """

src/transformers/modeling_utils.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import inspect
 import os
 import re
@@ -665,9 +666,28 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
 
         return new_embeddings
 
-    def init_weights(self):
+    def get_layers(self):
         """
-        Initializes and prunes weights if needed.
+        Returns the model's transformer layers.
+
+        Returns:
+            :obj:`List[nn.Module]` or :obj:`nn.ModuleList`: A list or :obj:`nn.ModuleList` containing all transformer
+            layers.
+        """
+        base_model = getattr(self, self.base_model_prefix, self)
+
+        if base_model is not self:
+            return base_model.get_layers()
+        else:
+            raise NotImplementedError
+
+    def init_weights(self):
+        # Backwards compatibility
+        self.init_weights_and_layers()
+
+    def init_weights_and_layers(self):
+        """
+        Initializes and prunes weights if needed. Sets up gradient checkpointing if in the configuration.
         """
         # Initialize weights
         self.apply(self._init_weights)
@@ -679,6 +699,11 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
         # Tie weights if needed
         self.tie_weights()
 
+        # Gradient checkpointing if needed
+        if self.config.gradient_checkpointing:
+            for layer in self.get_layers():
+                layer.forward = functools.partial(torch.utils.checkpoint.checkpoint, layer.forward)
+
     def prune_heads(self, heads_to_prune: Dict[int, List[int]]):
         """
         Prunes heads of the base model.

src/transformers/modeling_xlm.py

@@ -469,7 +469,7 @@ class XLMModel(XLMPreTrainedModel):
                 if self.attentions[int(layer)].n_heads == config.n_heads:
                     self.prune_heads({int(layer): list(map(int, heads))})
 
-        self.init_weights()
+        self.init_weights_and_layers()
         self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
 
     def get_input_embeddings(self):
@@ -508,8 +508,10 @@ class XLMModel(XLMPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -685,7 +687,7 @@ class XLMWithLMHeadModel(XLMPreTrainedModel):
         self.transformer = XLMModel(config)
         self.pred_layer = XLMPredLayer(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.pred_layer.proj
@@ -779,7 +781,7 @@ class XLMForSequenceClassification(XLMPreTrainedModel):
         self.transformer = XLMModel(config)
         self.sequence_summary = SequenceSummary(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -866,7 +868,7 @@ class XLMForQuestionAnsweringSimple(XLMPreTrainedModel):
         self.transformer = XLMModel(config)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -970,7 +972,7 @@ class XLMForQuestionAnswering(XLMPreTrainedModel):
         self.transformer = XLMModel(config)
         self.qa_outputs = SQuADHead(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @replace_return_docstrings(output_type=XLMForQuestionAnsweringOutput, config_class=_CONFIG_FOR_DOC)
@@ -1089,7 +1091,7 @@ class XLMForTokenClassification(XLMPreTrainedModel):
         self.dropout = nn.Dropout(config.dropout)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1182,7 +1184,7 @@ class XLMForMultipleChoice(XLMPreTrainedModel):
         self.sequence_summary = SequenceSummary(config)
         self.logits_proj = nn.Linear(config.num_labels, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLM_INPUTS_DOCSTRING.format("batch_size, num_choicec, sequence_length"))
     @add_code_sample_docstrings(

src/transformers/modeling_xlnet.py

@@ -952,7 +952,7 @@ class XLNetModel(XLNetPreTrainedModel):
         self.layer = nn.ModuleList([XLNetLayer(config) for _ in range(config.n_layer)])
         self.dropout = nn.Dropout(config.dropout)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_input_embeddings(self):
         return self.word_embedding
@@ -1087,8 +1087,10 @@ class XLNetModel(XLNetPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -1295,7 +1297,7 @@ class XLNetLMHeadModel(XLNetPreTrainedModel):
         self.transformer = XLNetModel(config)
         self.lm_loss = nn.Linear(config.d_model, config.vocab_size, bias=True)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     def get_output_embeddings(self):
         return self.lm_loss
@@ -1463,7 +1465,7 @@ class XLNetForSequenceClassification(XLNetPreTrainedModel):
         self.sequence_summary = SequenceSummary(config)
         self.logits_proj = nn.Linear(config.d_model, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1556,7 +1558,7 @@ class XLNetForTokenClassification(XLNetPreTrainedModel):
         self.transformer = XLNetModel(config)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1653,7 +1655,7 @@ class XLNetForMultipleChoice(XLNetPreTrainedModel):
         self.sequence_summary = SequenceSummary(config)
         self.logits_proj = nn.Linear(config.d_model, 1)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
     @add_code_sample_docstrings(
@@ -1754,7 +1756,7 @@ class XLNetForQuestionAnsweringSimple(XLNetPreTrainedModel):
         self.transformer = XLNetModel(config)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @add_code_sample_docstrings(
@@ -1866,7 +1868,7 @@ class XLNetForQuestionAnswering(XLNetPreTrainedModel):
         self.end_logits = PoolerEndLogits(config)
         self.answer_class = PoolerAnswerClass(config)
 
-        self.init_weights()
+        self.init_weights_and_layers()
 
     @add_start_docstrings_to_model_forward(XLNET_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     @replace_return_docstrings(output_type=XLNetForQuestionAnsweringOutput, config_class=_CONFIG_FOR_DOC)

src/transformers/pipelines.py

@@ -22,6 +22,7 @@ import sys
 import uuid
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from contextlib import contextmanager
 from os.path import abspath, exists
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
@@ -1324,6 +1325,29 @@ class FillMaskPipeline(Pipeline):
         return results
 
 
+class TokenClassificationArgumentHandler(ArgumentHandler):
+    """
+    Handles arguments for token classification.
+    """
+
+    def __call__(self, *args, **kwargs):
+
+        if args is not None and len(args) > 0:
+            if isinstance(args, str):
+                inputs = [args]
+            else:
+                inputs = args
+            batch_size = len(inputs)
+
+        offset_mapping = kwargs.get("offset_mapping", None)
+        if offset_mapping:
+            if isinstance(offset_mapping, list) and isinstance(offset_mapping[0], tuple):
+                offset_mapping = [offset_mapping]
+            if len(offset_mapping) != batch_size:
+                raise ("offset_mapping should have the same batch size as the input")
+        return inputs, offset_mapping
+
+
 @add_end_docstrings(
     PIPELINE_INIT_ARGS,
     r"""
@@ -1361,13 +1385,14 @@ class TokenClassificationPipeline(Pipeline):
         ignore_labels=["O"],
         task: str = "",
         grouped_entities: bool = False,
+        ignore_subwords: bool = True,
     ):
         super().__init__(
             model=model,
             tokenizer=tokenizer,
             modelcard=modelcard,
             framework=framework,
-            args_parser=args_parser,
+            args_parser=TokenClassificationArgumentHandler(),
             device=device,
             binary_output=binary_output,
             task=task,
@@ -1382,6 +1407,7 @@ class TokenClassificationPipeline(Pipeline):
         self._basic_tokenizer = BasicTokenizer(do_lower_case=False)
         self.ignore_labels = ignore_labels
         self.grouped_entities = grouped_entities
+        self.ignore_subwords = ignore_subwords
 
     def __call__(self, inputs: Union[str, List[str]], **kwargs):
         """
@@ -1402,10 +1428,15 @@ class TokenClassificationPipeline(Pipeline):
             - **index** (:obj:`int`, only present when ``self.grouped_entities=False``) -- The index of the
               corresponding token in the sentence.
         """
+
         if isinstance(inputs, str):
             inputs = [inputs]
+
+        offset_mappings = kwargs.get("offset_mappings")
+
         answers = []
-        for sentence in inputs:
+
+        for i, sentence in enumerate(inputs):
 
             # Manage correct placement of the tensors
             with self.device_placement():
@@ -1415,7 +1446,18 @@ class TokenClassificationPipeline(Pipeline):
                     return_attention_mask=False,
                     return_tensors=self.framework,
                     truncation=True,
+                    return_special_tokens_mask=True,
+                    return_offsets_mapping=self.tokenizer.is_fast,
                 )
+                if self.tokenizer.is_fast:
+                    offset_mapping = tokens["offset_mapping"].cpu().numpy()[0]
+                    del tokens["offset_mapping"]
+                elif offset_mappings:
+                    offset_mapping = offset_mappings[i]
+                else:
+                    raise Exception("To decode [UNK] tokens use a fast tokenizer or provide offset_mapping parameter")
+                special_tokens_mask = tokens["special_tokens_mask"].cpu().numpy()[0]
+                del tokens["special_tokens_mask"]
 
                 # Forward
                 if self.framework == "tf":
@@ -1432,24 +1474,35 @@ class TokenClassificationPipeline(Pipeline):
 
             entities = []
             # Filter to labels not in `self.ignore_labels`
+            # Filter special_tokens
             filtered_labels_idx = [
                 (idx, label_idx)
                 for idx, label_idx in enumerate(labels_idx)
-                if self.model.config.id2label[label_idx] not in self.ignore_labels
+                if (self.model.config.id2label[label_idx] not in self.ignore_labels) and not special_tokens_mask[idx]
             ]
 
             for idx, label_idx in filtered_labels_idx:
+                start_ind, end_ind = offset_mapping[idx]
+                word_ref = sentence[start_ind:end_ind]
+                word = self.tokenizer.convert_ids_to_tokens([int(input_ids[idx])])[0]
+                is_subword = len(word_ref) != len(word)
+
+                if int(input_ids[idx]) == self.tokenizer.unk_token_id:
+                    word = word_ref
+                    is_subword = False
 
                 entity = {
-                    "word": self.tokenizer.convert_ids_to_tokens(int(input_ids[idx])),
+                    "word": word,
                     "score": score[idx][label_idx].item(),
                     "entity": self.model.config.id2label[label_idx],
                     "index": idx,
                 }
 
+                if self.grouped_entities and self.ignore_subwords:
+                    entity["is_subword"] = is_subword
+
                 entities += [entity]
 
-            # Append grouped entities
             if self.grouped_entities:
                 answers += [self.group_entities(entities)]
             # Append ungrouped entities
@@ -1468,8 +1521,8 @@ class TokenClassificationPipeline(Pipeline):
             entities (:obj:`dict`): The entities predicted by the pipeline.
         """
         # Get the first entity in the entity group
-        entity = entities[0]["entity"]
-        scores = np.mean([entity["score"] for entity in entities])
+        entity = entities[0]["entity"].split("-")[-1]
+        scores = np.nanmean([entity["score"] for entity in entities])
         tokens = [entity["word"] for entity in entities]
 
         entity_group = {
@@ -1494,7 +1547,9 @@ class TokenClassificationPipeline(Pipeline):
             last_idx = entities[-1]["index"]
 
         for entity in entities:
+
             is_last_idx = entity["index"] == last_idx
+            is_subword = self.ignore_subwords and entity["is_subword"]
             if not entity_group_disagg:
                 entity_group_disagg += [entity]
                 if is_last_idx:
@@ -1503,10 +1558,19 @@ class TokenClassificationPipeline(Pipeline):
 
             # If the current entity is similar and adjacent to the previous entity, append it to the disaggregated entity group
             # The split is meant to account for the "B" and "I" suffixes
+            # Shouldn't merge if both entities are B-type
             if (
-                entity["entity"].split("-")[-1] == entity_group_disagg[-1]["entity"].split("-")[-1]
+                (
+                    entity["entity"].split("-")[-1] == entity_group_disagg[-1]["entity"].split("-")[-1]
+                    and entity["entity"].split("-")[0] != "B"
+                )
                 and entity["index"] == entity_group_disagg[-1]["index"] + 1
-            ):
+            ) or is_subword:
+                # Modify subword type to be previous_type
+                if is_subword:
+                    entity["entity"] = entity_group_disagg[-1]["entity"].split("-")[-1]
+                    entity["score"] = np.nan  # set ignored scores to nan and use np.nanmean
+
                 entity_group_disagg += [entity]
                 # Group the entities at the last entity
                 if is_last_idx:
@@ -1534,55 +1598,52 @@ class QuestionAnsweringArgumentHandler(ArgumentHandler):
     command-line supplied arguments.
     """
 
+    def normalize(self, item):
+        if isinstance(item, SquadExample):
+            return item
+        elif isinstance(item, dict):
+            for k in ["question", "context"]:
+                if k not in item:
+                    raise KeyError("You need to provide a dictionary with keys {question:..., context:...}")
+                elif item[k] is None:
+                    raise ValueError("`{}` cannot be None".format(k))
+                elif isinstance(item[k], str) and len(item[k]) == 0:
+                    raise ValueError("`{}` cannot be empty".format(k))
+
+            return QuestionAnsweringPipeline.create_sample(**item)
+        raise ValueError("{} argument needs to be of type (SquadExample, dict)".format(item))
+
     def __call__(self, *args, **kwargs):
-        # Position args, handling is sensibly the same as X and data, so forwarding to avoid duplicating
+        # Detect where the actual inputs are
         if args is not None and len(args) > 0:
             if len(args) == 1:
-                kwargs["X"] = args[0]
+                inputs = args[0]
+            elif len(args) == 2 and {type(el) for el in args} == {str}:
+                inputs = [{"question": args[0], "context": args[1]}]
             else:
-                kwargs["X"] = list(args)
-
+                inputs = list(args)
         # Generic compatibility with sklearn and Keras
         # Batched data
-        if "X" in kwargs or "data" in kwargs:
-            inputs = kwargs["X"] if "X" in kwargs else kwargs["data"]
-
-            if isinstance(inputs, dict):
-                inputs = [inputs]
-            else:
-                # Copy to avoid overriding arguments
-                inputs = [i for i in inputs]
-
-            for i, item in enumerate(inputs):
-                if isinstance(item, dict):
-                    if any(k not in item for k in ["question", "context"]):
-                        raise KeyError("You need to provide a dictionary with keys {question:..., context:...}")
-
-                    inputs[i] = QuestionAnsweringPipeline.create_sample(**item)
-
-                elif not isinstance(item, SquadExample):
-                    raise ValueError(
-                        "{} argument needs to be of type (list[SquadExample | dict], SquadExample, dict)".format(
-                            "X" if "X" in kwargs else "data"
-                        )
-                    )
-
-            # Tabular input
+        elif "X" in kwargs:
+            inputs = kwargs["X"]
+        elif "data" in kwargs:
+            inputs = kwargs["data"]
         elif "question" in kwargs and "context" in kwargs:
-            if isinstance(kwargs["question"], str):
-                kwargs["question"] = [kwargs["question"]]
-
-            if isinstance(kwargs["context"], str):
-                kwargs["context"] = [kwargs["context"]]
-
-            inputs = [
-                QuestionAnsweringPipeline.create_sample(q, c) for q, c in zip(kwargs["question"], kwargs["context"])
-            ]
+            inputs = [{"question": kwargs["question"], "context": kwargs["context"]}]
         else:
             raise ValueError("Unknown arguments {}".format(kwargs))
 
-        if not isinstance(inputs, list):
+        # Normalize inputs
+        if isinstance(inputs, dict):
             inputs = [inputs]
+        elif isinstance(inputs, Iterable):
+            # Copy to avoid overriding arguments
+            inputs = [i for i in inputs]
+        else:
+            raise ValueError("Invalid arguments {}".format(inputs))
+
+        for i, item in enumerate(inputs):
+            inputs[i] = self.normalize(item)
 
         return inputs
 

src/transformers/testing_utils.py

@@ -695,6 +695,31 @@ def mockenv(**kwargs):
     return unittest.mock.patch.dict(os.environ, kwargs)
 
 
+# --- pytest conf functions --- #
+
+# to avoid multiple invocation from tests/conftest.py and examples/conftest.py - make sure it's called only once
+pytest_opt_registered = {}
+
+
+def pytest_addoption_shared(parser):
+    """
+    This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there.
+
+    It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest`
+    option.
+
+    """
+    option = "--make-reports"
+    if option not in pytest_opt_registered:
+        parser.addoption(
+            option,
+            action="store",
+            default=False,
+            help="generate report files. The value of this option is used as a prefix to report names",
+        )
+        pytest_opt_registered[option] = 1
+
+
 def pytest_terminal_summary_main(tr, id):
     """
     Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current
@@ -728,7 +753,7 @@ def pytest_terminal_summary_main(tr, id):
     dir = "reports"
     Path(dir).mkdir(parents=True, exist_ok=True)
     report_files = {
-        k: f"{dir}/report_{id}_{k}.txt"
+        k: f"{dir}/{id}_{k}.txt"
         for k in [
             "durations",
             "errors",
@@ -824,7 +849,7 @@ def pytest_terminal_summary_main(tr, id):
     config.option.tbstyle = orig_tbstyle
 
 
-# the following code deals with async io between processes
+# --- distributed testing functions --- #
 
 # adapted from https://stackoverflow.com/a/59041913/9201239
 import asyncio  # noqa

src/transformers/tokenization_blenderbot.py

@@ -160,7 +160,7 @@ class BlenderbotSmallTokenizer(PreTrainedTokenizer):
             return self.cache[token]
         token = re.sub("([.,!?()])", r" \1", token)
         token = re.sub("(')", r" \1 ", token)
-        token = re.sub("\s{2,}", " ", token)
+        token = re.sub(r"\s{2,}", " ", token)
         if "\n" in token:
             token = token.replace("\n", " __newln__")
 

src/transformers/training_args.py

@@ -48,12 +48,19 @@ class TrainingArguments:
             If :obj:`True`, overwrite the content of the output directory. Use this to continue training if
             :obj:`output_dir` points to a checkpoint directory.
         do_train (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to run training or not.
+            Whether to run training or not. This argument is not directly used by :class:`~transformers.Trainer`, it's
+            intended to be used by your training/evaluation scripts instead. See the `example scripts
+            <https://github.com/huggingface/transformers/tree/master/examples>`__ for more details.
         do_eval (:obj:`bool`, `optional`):
-            Whether to run evaluation on the dev set or not. Will default to :obj:`evaluation_strategy` different from
-            :obj:`"no"`.
+            Whether to run evaluation on the dev set or not. Will be set to :obj:`True` if :obj:`evaluation_strategy`
+            is different from :obj:`"no"`. This argument is not directly used by :class:`~transformers.Trainer`, it's
+            intended to be used by your training/evaluation scripts instead. See the `example scripts
+            <https://github.com/huggingface/transformers/tree/master/examples>`__ for more details.
         do_predict (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to run predictions on the test set or not.
+            Whether to run predictions on the test set or not. This argument is not directly used by
+            :class:`~transformers.Trainer`, it's intended to be used by your training/evaluation scripts instead. See
+            the `example scripts <https://github.com/huggingface/transformers/tree/master/examples>`__ for more
+            details.
         evaluation_strategy (:obj:`str` or :class:`~transformers.trainer_utils.EvaluationStrategy`, `optional`, defaults to :obj:`"no"`):
             The evaluation strategy to adopt during training. Possible values are:
 

src/transformers/utils/dummy_pt_objects.py

@@ -45,6 +45,15 @@ class DataCollatorForSOP:
         requires_pytorch(self)
 
 
+class DataCollatorForTokenClassification:
+    def __init__(self, *args, **kwargs):
+        requires_pytorch(self)
+
+    @classmethod
+    def from_pretrained(self, *args, **kwargs):
+        requires_pytorch(self)
+
+
 class DataCollatorForWholeWordMask:
     def __init__(self, *args, **kwargs):
         requires_pytorch(self)

templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py

@@ -205,7 +205,7 @@ def main():
         if data_args.train_file is not None:
             data_files["train"] = data_args.train_file
         if data_args.validation_file is not None:
-            data_files["validation"] = data_args.train_file
+            data_files["validation"] = data_args.validation_file
         extension = data_args.train_file.split(".")[-1]
         if extension == "txt":
             extension = "text"

templates/adding_a_new_model/modeling_xxx.py

@@ -328,8 +328,10 @@ class XxxModel(XxxPreTrainedModel):
         output_hidden_states=None,
         return_dict=None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
+        output_attentions = torch.tensor(
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = torch.tensor(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict

tests/conftest.py

@@ -24,17 +24,14 @@ def pytest_configure(config):
 
 
 def pytest_addoption(parser):
-    parser.addoption(
-        "--make_reports",
-        action="store",
-        default=False,
-        help="generate report files - the value will be used as a `report_`+val+`reportname.txt`",
-    )
+    from transformers.testing_utils import pytest_addoption_shared
+
+    pytest_addoption_shared(parser)
 
 
 def pytest_terminal_summary(terminalreporter):
     from transformers.testing_utils import pytest_terminal_summary_main
 
-    make_reports = terminalreporter.config.getoption("--make_reports")
+    make_reports = terminalreporter.config.getoption("--make-reports")
     if make_reports:
         pytest_terminal_summary_main(terminalreporter, id=make_reports)

tests/test_data_collator.py

@@ -1,7 +1,10 @@
+import os
+import shutil
+import tempfile
 import unittest
 
-from transformers import AutoTokenizer, is_torch_available
-from transformers.testing_utils import require_torch, slow
+from transformers import BertTokenizer, is_torch_available, set_seed
+from transformers.testing_utils import require_torch
 
 
 if is_torch_available():
@@ -12,22 +15,25 @@ if is_torch_available():
         DataCollatorForNextSentencePrediction,
         DataCollatorForPermutationLanguageModeling,
         DataCollatorForSOP,
-        GlueDataset,
-        GlueDataTrainingArguments,
-        LineByLineTextDataset,
-        LineByLineWithSOPTextDataset,
-        TextDataset,
-        TextDatasetForNextSentencePrediction,
+        DataCollatorForTokenClassification,
+        DataCollatorWithPadding,
         default_data_collator,
     )
 
 
-PATH_SAMPLE_TEXT = "./tests/fixtures/sample_text.txt"
-PATH_SAMPLE_TEXT_DIR = "./tests/fixtures/tests_samples/wiki_text"
-
-
 @require_torch
 class DataCollatorIntegrationTest(unittest.TestCase):
+    def setUp(self):
+        self.tmpdirname = tempfile.mkdtemp()
+
+        vocab_tokens = ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
+        self.vocab_file = os.path.join(self.tmpdirname, "vocab.txt")
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdirname)
+
     def test_default_with_dict(self):
         features = [{"label": i, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
         batch = default_data_collator(features)
@@ -57,6 +63,17 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         self.assertEqual(batch["labels"].dtype, torch.long)
         self.assertEqual(batch["inputs"].shape, torch.Size([8, 10]))
 
+    def test_default_classification_and_regression(self):
+        data_collator = default_data_collator
+
+        features = [{"input_ids": [0, 1, 2, 3, 4], "label": i} for i in range(4)]
+        batch = data_collator(features)
+        self.assertEqual(batch["labels"].dtype, torch.long)
+
+        features = [{"input_ids": [0, 1, 2, 3, 4], "label": float(i)} for i in range(4)]
+        batch = data_collator(features)
+        self.assertEqual(batch["labels"].dtype, torch.float)
+
     def test_default_with_no_labels(self):
         features = [{"label": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
         batch = default_data_collator(features)
@@ -69,128 +86,144 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         self.assertTrue("labels" not in batch)
         self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))
 
-    @slow
-    def test_default_classification(self):
-        MODEL_ID = "bert-base-cased-finetuned-mrpc"
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        data_args = GlueDataTrainingArguments(
-            task_name="mrpc", data_dir="./tests/fixtures/tests_samples/MRPC", overwrite_cache=True
-        )
-        dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev")
-        data_collator = default_data_collator
-        batch = data_collator(dataset.features)
-        self.assertEqual(batch["labels"].dtype, torch.long)
+    def test_data_collator_with_padding(self):
+        tokenizer = BertTokenizer(self.vocab_file)
+        features = [{"input_ids": [0, 1, 2]}, {"input_ids": [0, 1, 2, 3, 4, 5]}]
 
-    @slow
-    def test_default_regression(self):
-        MODEL_ID = "distilroberta-base"
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        data_args = GlueDataTrainingArguments(
-            task_name="sts-b", data_dir="./tests/fixtures/tests_samples/STS-B", overwrite_cache=True
-        )
-        dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev")
-        data_collator = default_data_collator
-        batch = data_collator(dataset.features)
-        self.assertEqual(batch["labels"].dtype, torch.float)
+        data_collator = DataCollatorWithPadding(tokenizer)
+        batch = data_collator(features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size([2, 6]))
+        self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
+
+        data_collator = DataCollatorWithPadding(tokenizer, padding="max_length", max_length=10)
+        batch = data_collator(features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size([2, 10]))
+
+        data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
+        batch = data_collator(features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size([2, 8]))
+
+    def test_data_collator_for_token_classification(self):
+        tokenizer = BertTokenizer(self.vocab_file)
+        features = [
+            {"input_ids": [0, 1, 2], "labels": [0, 1, 2]},
+            {"input_ids": [0, 1, 2, 3, 4, 5], "labels": [0, 1, 2, 3, 4, 5]},
+        ]
+
+        data_collator = DataCollatorForTokenClassification(tokenizer)
+        batch = data_collator(features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size([2, 6]))
+        self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
+        self.assertEqual(batch["labels"].shape, torch.Size([2, 6]))
+        self.assertEqual(batch["labels"][0].tolist(), [0, 1, 2] + [-100] * 3)
+
+        data_collator = DataCollatorForTokenClassification(tokenizer, padding="max_length", max_length=10)
+        batch = data_collator(features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size([2, 10]))
+        self.assertEqual(batch["labels"].shape, torch.Size([2, 10]))
+
+        data_collator = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=8)
+        batch = data_collator(features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size([2, 8]))
+        self.assertEqual(batch["labels"].shape, torch.Size([2, 8]))
+
+        data_collator = DataCollatorForTokenClassification(tokenizer, label_pad_token_id=-1)
+        batch = data_collator(features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size([2, 6]))
+        self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
+        self.assertEqual(batch["labels"].shape, torch.Size([2, 6]))
+        self.assertEqual(batch["labels"][0].tolist(), [0, 1, 2] + [-1] * 3)
+
+    def test_data_collator_for_language_modeling(self):
+        tokenizer = BertTokenizer(self.vocab_file)
+        no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
+        pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
 
-    @slow
-    def test_lm_tokenizer_without_padding(self):
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
         data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
-        # ^ causal lm
+        batch = data_collator(no_pad_features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
 
-        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
+        batch = data_collator(pad_features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
+
+        tokenizer._pad_token = None
+        data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
         with self.assertRaises(ValueError):
-            # Expect error due to padding token missing on gpt2:
-            data_collator(examples)
+            # Expect error due to padding token missing
+            data_collator(pad_features)
 
-        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(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
-
-    @slow
-    def test_lm_tokenizer_with_padding(self):
-        tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
+        set_seed(42)  # For reproducibility
+        tokenizer = BertTokenizer(self.vocab_file)
         data_collator = DataCollatorForLanguageModeling(tokenizer)
-        # ^ masked lm
+        batch = data_collator(no_pad_features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
 
-        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((31, 107)))
-        self.assertEqual(batch["labels"].shape, torch.Size((31, 107)))
+        masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
+        self.assertTrue(torch.any(masked_tokens))
+        self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
 
-        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(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
+        batch = data_collator(pad_features)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
+
+        masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
+        self.assertTrue(torch.any(masked_tokens))
+        self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
 
-    @slow
     def test_plm(self):
-        tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased")
+        tokenizer = BertTokenizer(self.vocab_file)
+        no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
+        pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
+
         data_collator = DataCollatorForPermutationLanguageModeling(tokenizer)
-        # ^ permutation lm
 
-        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
+        batch = data_collator(pad_features)
         self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((31, 112)))
-        self.assertEqual(batch["perm_mask"].shape, torch.Size((31, 112, 112)))
-        self.assertEqual(batch["target_mapping"].shape, torch.Size((31, 112, 112)))
-        self.assertEqual(batch["labels"].shape, torch.Size((31, 112)))
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
+        self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 10, 10)))
+        self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 10, 10)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
 
-        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(examples)
+        batch = data_collator(no_pad_features)
         self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
-        self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 512, 512)))
-        self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 512, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
+        self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 10, 10)))
+        self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 10, 10)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
 
         example = [torch.randint(5, [5])]
         with self.assertRaises(ValueError):
             # Expect error due to odd sequence length
             data_collator(example)
 
-    @slow
     def test_nsp(self):
-        tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+        tokenizer = BertTokenizer(self.vocab_file)
+        features = [{"tokens_a": [0, 1, 2, 3, 4], "tokens_b": [0, 1, 2, 3, 4], "is_random_next": i} for i in range(2)]
         data_collator = DataCollatorForNextSentencePrediction(tokenizer)
+        batch = data_collator(features)
 
-        dataset = TextDatasetForNextSentencePrediction(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
+        self.assertEqual(batch["token_type_ids"].shape, torch.Size((2, 512)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
+        self.assertEqual(batch["next_sentence_label"].shape, torch.Size((2,)))
 
-        # Since there are randomly generated false samples, the total number of samples is not fixed.
-        total_samples = batch["input_ids"].shape[0]
-        self.assertEqual(batch["input_ids"].shape, torch.Size((total_samples, 512)))
-        self.assertEqual(batch["token_type_ids"].shape, torch.Size((total_samples, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((total_samples, 512)))
-        self.assertEqual(batch["next_sentence_label"].shape, torch.Size((total_samples,)))
-
-    @slow
     def test_sop(self):
-        tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")
+        tokenizer = BertTokenizer(self.vocab_file)
+        features = [
+            {
+                "input_ids": torch.tensor([0, 1, 2, 3, 4]),
+                "token_type_ids": torch.tensor([0, 1, 2, 3, 4]),
+                "sentence_order_label": torch.tensor(i),
+            }
+            for i in range(2)
+        ]
         data_collator = DataCollatorForSOP(tokenizer)
+        batch = data_collator(features)
 
-        dataset = LineByLineWithSOPTextDataset(tokenizer, file_dir=PATH_SAMPLE_TEXT_DIR, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
-
-        # Since there are randomly generated false samples, the total number of samples is not fixed.
-        total_samples = batch["input_ids"].shape[0]
-        self.assertEqual(batch["input_ids"].shape, torch.Size((total_samples, 512)))
-        self.assertEqual(batch["token_type_ids"].shape, torch.Size((total_samples, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((total_samples, 512)))
-        self.assertEqual(batch["sentence_order_label"].shape, torch.Size((total_samples,)))
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 5)))
+        self.assertEqual(batch["token_type_ids"].shape, torch.Size((2, 5)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 5)))
+        self.assertEqual(batch["sentence_order_label"].shape, torch.Size((2,)))

tests/test_generation_utils.py

@@ -140,10 +140,6 @@ class GenerationTesterMixin:
             # check `generate()` and `greedy_search()` are equal
             kwargs = {}
             if model.config.is_encoder_decoder:
-                encoder_outputs, input_ids, attention_mask = self._get_encoder_outputs(
-                    model, input_ids, attention_mask
-                )
-                kwargs["encoder_outputs"] = encoder_outputs
                 max_length = 4
 
             output_ids_generate = model.generate(
@@ -154,6 +150,13 @@ class GenerationTesterMixin:
                 max_length=max_length,
                 **logits_process_kwargs,
             )
+
+            if model.config.is_encoder_decoder:
+                encoder_outputs, input_ids, attention_mask = self._get_encoder_outputs(
+                    model, input_ids, attention_mask
+                )
+                kwargs["encoder_outputs"] = encoder_outputs
+
             with torch.no_grad():
                 output_ids_greedy = model.greedy_search(
                     input_ids,

tests/test_modeling_albert.py

@@ -24,6 +24,8 @@ from .test_modeling_common import ModelTesterMixin, ids_tensor, random_attention
 
 
 if is_torch_available():
+    import torch
+
     from transformers import (
         AlbertConfig,
         AlbertForMaskedLM,
@@ -213,6 +215,8 @@ class AlbertModelTester:
 @require_torch
 class AlbertModelTest(ModelTesterMixin, unittest.TestCase):
 
+    test_gradient_checkpointing = True
+
     all_model_classes = (
         (
             AlbertModel,
@@ -263,3 +267,28 @@ class AlbertModelTest(ModelTesterMixin, unittest.TestCase):
         for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             model = AlbertModel.from_pretrained(model_name)
             self.assertIsNotNone(model)
+
+    def test_model_gradient_checkpointing(self):
+        if not self.test_gradient_checkpointing:
+            return
+
+        for model_class in self.all_model_classes:
+            print(model_class)
+            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+            model = model_class(config)
+            model.to(torch_device)
+            model.eval()
+            outputs_no_checkpointing = model(**self._prepare_for_class(inputs_dict, model_class))
+
+            config.gradient_checkpointing = True
+            model_with_gc = model_class(config)
+            model_with_gc.load_state_dict(model.state_dict())
+            model_with_gc.to(torch_device)
+            model_with_gc.eval()
+            outputs_with_checkpointing = model_with_gc(**self._prepare_for_class(inputs_dict, model_class))
+
+            for output_no_checkpointing, output_with_checkpointing in zip(
+                outputs_no_checkpointing, outputs_with_checkpointing
+            ):
+                if isinstance(output_with_checkpointing, torch.Tensor):
+                    self.assertTrue(torch.allclose(output_no_checkpointing, output_with_checkpointing))

tests/test_modeling_bart.py

@@ -140,6 +140,7 @@ class BARTModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
     test_pruning = False
     test_head_masking = False
     test_missing_keys = False
+    test_gradient_checkpointing = False
 
     def setUp(self):
         self.model_tester = ModelTester(self)

tests/test_modeling_bert.py

@@ -360,6 +360,8 @@ class BertModelTester:
 @require_torch
 class BertModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
 
+    test_gradient_checkpointing = True
+
     all_model_classes = (
         (
             BertModel,

tests/test_modeling_bert_generation.py

@@ -189,6 +189,9 @@ class BertGenerationEncoderTest(ModelTesterMixin, GenerationTesterMixin, unittes
     all_model_classes = (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
     all_generative_model_classes = (BertGenerationDecoder,) if is_torch_available() else ()
 
+    # Should have a class in `AutoModelWithLMHead` to be tested on Gradient Checkpointing.
+    test_gradient_checkpointing = False
+
     def setUp(self):
         self.model_tester = BertGenerationEncoderTester(self)
         self.config_tester = ConfigTester(self, config_class=BertGenerationConfig, hidden_size=37)

tests/test_modeling_blenderbot.py

@@ -100,6 +100,7 @@ class BlenderbotTesterMixin(ModelTesterMixin, unittest.TestCase):
     test_pruning = False
     test_missing_keys = False
     test_torchscript = False
+    test_gradient_checkpointing = False
 
     def setUp(self):
         self.model_tester = BlenderbotModelTester(self)

tests/test_modeling_common.py

@@ -67,6 +67,7 @@ class ModelTesterMixin:
     test_head_masking = True
     test_missing_keys = True
     is_encoder_decoder = False
+    test_gradient_checkpointing = True
 
     def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
         inputs_dict = copy.deepcopy(inputs_dict)
@@ -881,6 +882,71 @@ class ModelTesterMixin:
             with torch.no_grad():
                 model(**inputs)[0]
 
+    def test_model_gradient_checkpointing_equivalent_results(self):
+        if not self.test_gradient_checkpointing:
+            return
+
+        for model_class in self.all_model_classes:
+            print(model_class)
+            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+            model = model_class(config)
+            model.to(torch_device)
+            model.eval()
+            outputs_no_checkpointing = model(**self._prepare_for_class(inputs_dict, model_class))
+
+            config.gradient_checkpointing = True
+            model_with_gc = model_class(config)
+            model_with_gc.load_state_dict(model.state_dict())
+            model_with_gc.to(torch_device)
+            model_with_gc.eval()
+            outputs_with_checkpointing = model_with_gc(**self._prepare_for_class(inputs_dict, model_class))
+
+            for output_no_checkpointing, output_with_checkpointing in zip(
+                outputs_no_checkpointing, outputs_with_checkpointing
+            ):
+                if isinstance(output_with_checkpointing, torch.Tensor):
+                    self.assertTrue(torch.allclose(output_no_checkpointing, output_with_checkpointing))
+
+    def test_model_gradient_checkpointing_memory(self):
+        if not self.test_gradient_checkpointing:
+            return
+
+        from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments
+
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        gc_config = config
+        gc_config.gradient_checkpointing = True
+
+        ngc_config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+        assert gc_config.gradient_checkpointing and not ngc_config.gradient_checkpointing
+
+        # Gradient Checkpointing becomes really valuable for large input sizes.
+        batch_size = int(self.model_tester.batch_size / 2)
+        sequence_length = int(self.model_tester.max_position_embeddings / 2)
+
+        # Run the benchmark
+        args = PyTorchBenchmarkArguments(
+            models=["gc", "ngc"],
+            batch_sizes=[batch_size],
+            sequence_lengths=[sequence_length],
+            inference=False,
+            training=True,
+            speed=False,
+        )
+        benchmark = PyTorchBenchmark(args, configs=[gc_config, ngc_config])
+
+        # Parse results
+        result = benchmark.run()
+        memory_result = result.memory_train_result
+        gc_memory_result = memory_result["gc"]["result"][batch_size][sequence_length]
+        ngc_memory_result = memory_result["ngc"]["result"][batch_size][sequence_length]
+
+        self.assertTrue(
+            ngc_memory_result > gc_memory_result * 1.2,
+            "Assert that gradient checkpointing requires a lot less memory.",
+        )
+
     @require_torch_multigpu
     def test_multigpu_data_parallel_forward(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

tests/test_modeling_deberta.py

@@ -54,6 +54,9 @@ class DebertaModelTest(ModelTesterMixin, unittest.TestCase):
     test_head_masking = False
     is_encoder_decoder = False
 
+    # Needs to have DebertaForMaskedLM to test this.
+    test_gradient_checkpointing = False
+
     class DebertaModelTester(object):
         def __init__(
             self,

tests/test_modeling_distilbert.py

@@ -211,6 +211,7 @@ class DistilBertModelTest(ModelTesterMixin, unittest.TestCase):
     test_torchscript = True
     test_resize_embeddings = True
     test_head_masking = True
+    test_gradient_checkpointing = True
 
     def setUp(self):
         self.model_tester = DistilBertModelTester(self)

tests/test_modeling_flaubert.py

@@ -329,6 +329,8 @@ class FlaubertModelTester(object):
 @require_torch
 class FlaubertModelTest(ModelTesterMixin, unittest.TestCase):
 
+    test_gradient_checkpointing = False
+
     all_model_classes = (
         (
             FlaubertModel,

tests/test_modeling_fsmt.py

@@ -128,6 +128,7 @@ class FSMTModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
     test_pruning = False
     test_head_masking = False
     test_missing_keys = False
+    test_gradient_checkpointing = False
 
     def setUp(self):
         self.model_tester = ModelTester(self)

tests/test_modeling_funnel.py

@@ -454,3 +454,28 @@ class FunnelModelIntegrationTest(unittest.TestCase):
         expected_output_mean = torch.tensor(0.0256)
         self.assertTrue(torch.allclose(output.sum(), expected_output_sum, atol=1e-4))
         self.assertTrue(torch.allclose(output.mean(), expected_output_mean, atol=1e-4))
+
+    def test_model_gradient_checkpointing_equivalent_results(self):
+        if not self.test_gradient_checkpointing:
+            return
+
+        for model_class in self.all_model_classes:
+            print(model_class)
+            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+            model = model_class(config)
+            model.to(torch_device)
+            model.eval()
+            outputs_no_checkpointing = model(**self._prepare_for_class(inputs_dict, model_class))
+
+            config.gradient_checkpointing = True
+            model_with_gc = model_class(config)
+            model_with_gc.load_state_dict(model.state_dict())
+            model_with_gc.to(torch_device)
+            model_with_gc.eval()
+            outputs_with_checkpointing = model_with_gc(**self._prepare_for_class(inputs_dict, model_class))
+
+            for output_no_checkpointing, output_with_checkpointing in zip(
+                outputs_no_checkpointing, outputs_with_checkpointing
+            ):
+                if isinstance(output_with_checkpointing, torch.Tensor):
+                    self.assertTrue(torch.allclose(output_no_checkpointing, output_with_checkpointing))

tests/test_modeling_longformer.py

@@ -274,6 +274,7 @@ class LongformerModelTest(ModelTesterMixin, unittest.TestCase):
     test_pruning = False  # pruning is not supported
     test_headmasking = False  # head masking is not supported
     test_torchscript = False
+    test_gradient_checkpointing = False
 
     all_model_classes = (
         (

tests/test_modeling_lxmert.py

@@ -528,10 +528,7 @@ class LxmertModelTest(ModelTesterMixin, unittest.TestCase):
     test_head_masking = False
     test_pruning = False
     test_torchscript = False
-
-    test_head_masking = False
-    test_pruning = False
-    test_torchscript = False
+    test_gradient_checkpointing = False
 
     def setUp(self):
         self.model_tester = LxmertModelTester(self)

tests/test_modeling_prophetnet.py

@@ -862,6 +862,7 @@ class ProphetNetModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.Test
     test_resize_embeddings = False
     test_headmasking = False
     is_encoder_decoder = True
+    test_gradient_checkpointing = False
 
     def setUp(self):
         self.model_tester = ProphetNetModelTester(self)

tests/test_modeling_reformer.py

@@ -583,6 +583,7 @@ class ReformerLocalAttnModelTest(ReformerTesterMixin, GenerationTesterMixin, Mod
     test_pruning = False
     test_headmasking = False
     test_torchscript = False
+    test_gradient_checkpointing = False
 
     def prepare_kwargs(self):
         return {

tests/test_modeling_squeezebert.py

@@ -232,6 +232,7 @@ class SqueezeBertModelTest(ModelTesterMixin, unittest.TestCase):
     test_torchscript = True
     test_resize_embeddings = True
     test_head_masking = False
+    test_gradient_checkpointing = False
 
     def setUp(self):
         self.model_tester = SqueezeBertModelTester(self)

tests/test_pipelines_ner.py

@@ -1,8 +1,8 @@
 import unittest
 
-from transformers import pipeline
+from transformers import AutoTokenizer, pipeline
 from transformers.pipelines import Pipeline
-from transformers.testing_utils import require_tf
+from transformers.testing_utils import require_tf, require_torch
 
 from .test_pipelines_common import CustomInputPipelineCommonMixin
 
@@ -19,38 +19,54 @@ class NerPipelineTests(CustomInputPipelineCommonMixin, unittest.TestCase):
 
     def _test_pipeline(self, nlp: Pipeline):
         output_keys = {"entity", "word", "score"}
+        if nlp.grouped_entities:
+            output_keys = {"entity_group", "word", "score"}
 
         ungrouped_ner_inputs = [
             [
-                {"entity": "B-PER", "index": 1, "score": 0.9994944930076599, "word": "Cons"},
-                {"entity": "B-PER", "index": 2, "score": 0.8025449514389038, "word": "##uelo"},
-                {"entity": "I-PER", "index": 3, "score": 0.9993102550506592, "word": "Ara"},
-                {"entity": "I-PER", "index": 4, "score": 0.9993743896484375, "word": "##új"},
-                {"entity": "I-PER", "index": 5, "score": 0.9992871880531311, "word": "##o"},
-                {"entity": "I-PER", "index": 6, "score": 0.9993029236793518, "word": "No"},
-                {"entity": "I-PER", "index": 7, "score": 0.9981776475906372, "word": "##guera"},
-                {"entity": "B-PER", "index": 15, "score": 0.9998136162757874, "word": "Andrés"},
-                {"entity": "I-PER", "index": 16, "score": 0.999740719795227, "word": "Pas"},
-                {"entity": "I-PER", "index": 17, "score": 0.9997414350509644, "word": "##tran"},
-                {"entity": "I-PER", "index": 18, "score": 0.9996136426925659, "word": "##a"},
-                {"entity": "B-ORG", "index": 28, "score": 0.9989739060401917, "word": "Far"},
-                {"entity": "I-ORG", "index": 29, "score": 0.7188422083854675, "word": "##c"},
+                {"entity": "B-PER", "index": 1, "score": 0.9994944930076599, "is_subword": False, "word": "Cons"},
+                {"entity": "B-PER", "index": 2, "score": 0.8025449514389038, "is_subword": True, "word": "##uelo"},
+                {"entity": "I-PER", "index": 3, "score": 0.9993102550506592, "is_subword": False, "word": "Ara"},
+                {"entity": "I-PER", "index": 4, "score": 0.9993743896484375, "is_subword": True, "word": "##új"},
+                {"entity": "I-PER", "index": 5, "score": 0.9992871880531311, "is_subword": True, "word": "##o"},
+                {"entity": "I-PER", "index": 6, "score": 0.9993029236793518, "is_subword": False, "word": "No"},
+                {"entity": "I-PER", "index": 7, "score": 0.9981776475906372, "is_subword": True, "word": "##guera"},
+                {"entity": "B-PER", "index": 15, "score": 0.9998136162757874, "is_subword": False, "word": "Andrés"},
+                {"entity": "I-PER", "index": 16, "score": 0.999740719795227, "is_subword": False, "word": "Pas"},
+                {"entity": "I-PER", "index": 17, "score": 0.9997414350509644, "is_subword": True, "word": "##tran"},
+                {"entity": "I-PER", "index": 18, "score": 0.9996136426925659, "is_subword": True, "word": "##a"},
+                {"entity": "B-ORG", "index": 28, "score": 0.9989739060401917, "is_subword": False, "word": "Far"},
+                {"entity": "I-ORG", "index": 29, "score": 0.7188422083854675, "is_subword": True, "word": "##c"},
             ],
             [
-                {"entity": "I-PER", "index": 1, "score": 0.9968166351318359, "word": "En"},
-                {"entity": "I-PER", "index": 2, "score": 0.9957635998725891, "word": "##zo"},
-                {"entity": "I-ORG", "index": 7, "score": 0.9986497163772583, "word": "UN"},
+                {"entity": "I-PER", "index": 1, "score": 0.9968166351318359, "is_subword": False, "word": "En"},
+                {"entity": "I-PER", "index": 2, "score": 0.9957635998725891, "is_subword": True, "word": "##zo"},
+                {"entity": "I-ORG", "index": 7, "score": 0.9986497163772583, "is_subword": False, "word": "UN"},
             ],
         ]
+
         expected_grouped_ner_results = [
             [
-                {"entity_group": "B-PER", "score": 0.9710702640669686, "word": "Consuelo Araújo Noguera"},
-                {"entity_group": "B-PER", "score": 0.9997273534536362, "word": "Andrés Pastrana"},
-                {"entity_group": "B-ORG", "score": 0.8589080572128296, "word": "Farc"},
+                {"entity_group": "PER", "score": 0.999369223912557, "word": "Consuelo Araújo Noguera"},
+                {"entity_group": "PER", "score": 0.9997771680355072, "word": "Andrés Pastrana"},
+                {"entity_group": "ORG", "score": 0.9989739060401917, "word": "Farc"},
             ],
             [
-                {"entity_group": "I-PER", "score": 0.9962901175022125, "word": "Enzo"},
-                {"entity_group": "I-ORG", "score": 0.9986497163772583, "word": "UN"},
+                {"entity_group": "PER", "score": 0.9968166351318359, "word": "Enzo"},
+                {"entity_group": "ORG", "score": 0.9986497163772583, "word": "UN"},
+            ],
+        ]
+
+        expected_grouped_ner_results_w_subword = [
+            [
+                {"entity_group": "PER", "score": 0.9994944930076599, "word": "Cons"},
+                {"entity_group": "PER", "score": 0.9663328925768534, "word": "##uelo Araújo Noguera"},
+                {"entity_group": "PER", "score": 0.9997273534536362, "word": "Andrés Pastrana"},
+                {"entity_group": "ORG", "score": 0.8589080572128296, "word": "Farc"},
+            ],
+            [
+                {"entity_group": "PER", "score": 0.9962901175022125, "word": "Enzo"},
+                {"entity_group": "ORG", "score": 0.9986497163772583, "word": "UN"},
             ],
         ]
 
@@ -77,12 +93,80 @@ class NerPipelineTests(CustomInputPipelineCommonMixin, unittest.TestCase):
             for key in output_keys:
                 self.assertIn(key, result)
 
-        for ungrouped_input, grouped_result in zip(ungrouped_ner_inputs, expected_grouped_ner_results):
-            self.assertEqual(nlp.group_entities(ungrouped_input), grouped_result)
+        if nlp.grouped_entities:
+            if nlp.ignore_subwords:
+                for ungrouped_input, grouped_result in zip(ungrouped_ner_inputs, expected_grouped_ner_results):
+                    self.assertEqual(nlp.group_entities(ungrouped_input), grouped_result)
+            else:
+                for ungrouped_input, grouped_result in zip(
+                    ungrouped_ner_inputs, expected_grouped_ner_results_w_subword
+                ):
+                    self.assertEqual(nlp.group_entities(ungrouped_input), grouped_result)
 
     @require_tf
     def test_tf_only(self):
         model_name = "Narsil/small"  # This model only has a TensorFlow version
         # We test that if we don't specificy framework='tf', it gets detected automatically
-        nlp = pipeline(task="ner", model=model_name, tokenizer=model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+        nlp = pipeline(task="ner", model=model_name, tokenizer=tokenizer)
         self._test_pipeline(nlp)
+
+    #         offset=tokenizer(VALID_INPUTS[0],return_offsets_mapping=True)['offset_mapping']
+    #         pipeline_running_kwargs = {"offset_mapping"}  # Additional kwargs to run the pipeline with
+
+    @require_tf
+    def test_tf_defaults(self):
+        for model_name in self.small_models:
+            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+            nlp = pipeline(task="ner", model=model_name, tokenizer=tokenizer, framework="tf")
+        self._test_pipeline(nlp)
+
+    @require_tf
+    def test_tf_small(self):
+        for model_name in self.small_models:
+            print(model_name)
+            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+            nlp = pipeline(
+                task="ner",
+                model=model_name,
+                tokenizer=tokenizer,
+                framework="tf",
+                grouped_entities=True,
+                ignore_subwords=True,
+            )
+            self._test_pipeline(nlp)
+
+            for model_name in self.small_models:
+                tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+                nlp = pipeline(
+                    task="ner",
+                    model=model_name,
+                    tokenizer=tokenizer,
+                    framework="tf",
+                    grouped_entities=True,
+                    ignore_subwords=False,
+                )
+                self._test_pipeline(nlp)
+
+    @require_torch
+    def test_pt_defaults(self):
+        for model_name in self.small_models:
+            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+            nlp = pipeline(task="ner", model=model_name, tokenizer=tokenizer)
+            self._test_pipeline(nlp)
+
+    @require_torch
+    def test_torch_small(self):
+        for model_name in self.small_models:
+            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+            nlp = pipeline(
+                task="ner", model=model_name, tokenizer=tokenizer, grouped_entities=True, ignore_subwords=True
+            )
+            self._test_pipeline(nlp)
+
+        for model_name in self.small_models:
+            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+            nlp = pipeline(
+                task="ner", model=model_name, tokenizer=tokenizer, grouped_entities=True, ignore_subwords=False
+            )
+            self._test_pipeline(nlp)

tests/test_pipelines_question_answering.py

@@ -1,6 +1,7 @@
 import unittest
 
-from transformers.pipelines import Pipeline
+from transformers.data.processors.squad import SquadExample
+from transformers.pipelines import Pipeline, QuestionAnsweringArgumentHandler
 
 from .test_pipelines_common import CustomInputPipelineCommonMixin
 
@@ -43,5 +44,116 @@ class QAPipelineTests(CustomInputPipelineCommonMixin, unittest.TestCase):
             for key in output_keys:
                 self.assertIn(key, result)
         for bad_input in invalid_inputs:
-            self.assertRaises(Exception, nlp, bad_input)
-        self.assertRaises(Exception, nlp, invalid_inputs)
+            self.assertRaises(ValueError, nlp, bad_input)
+        self.assertRaises(ValueError, nlp, invalid_inputs)
+
+    def test_argument_handler(self):
+        qa = QuestionAnsweringArgumentHandler()
+
+        Q = "Where was HuggingFace founded ?"
+        C = "HuggingFace was founded in Paris"
+
+        normalized = qa(Q, C)
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa(question=Q, context=C)
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa(question=Q, context=C)
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa({"question": Q, "context": C})
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa([{"question": Q, "context": C}])
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa([{"question": Q, "context": C}, {"question": Q, "context": C}])
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 2)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa(X={"question": Q, "context": C})
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa(X=[{"question": Q, "context": C}])
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+        normalized = qa(data={"question": Q, "context": C})
+        self.assertEqual(type(normalized), list)
+        self.assertEqual(len(normalized), 1)
+        self.assertEqual({type(el) for el in normalized}, {SquadExample})
+
+    def test_argument_handler_error_handling(self):
+        qa = QuestionAnsweringArgumentHandler()
+
+        Q = "Where was HuggingFace founded ?"
+        C = "HuggingFace was founded in Paris"
+
+        with self.assertRaises(KeyError):
+            qa({"context": C})
+        with self.assertRaises(KeyError):
+            qa({"question": Q})
+        with self.assertRaises(KeyError):
+            qa([{"context": C}])
+        with self.assertRaises(ValueError):
+            qa(None, C)
+        with self.assertRaises(ValueError):
+            qa("", C)
+        with self.assertRaises(ValueError):
+            qa(Q, None)
+        with self.assertRaises(ValueError):
+            qa(Q, "")
+
+        with self.assertRaises(ValueError):
+            qa(question=None, context=C)
+        with self.assertRaises(ValueError):
+            qa(question="", context=C)
+        with self.assertRaises(ValueError):
+            qa(question=Q, context=None)
+        with self.assertRaises(ValueError):
+            qa(question=Q, context="")
+
+        with self.assertRaises(ValueError):
+            qa({"question": None, "context": C})
+        with self.assertRaises(ValueError):
+            qa({"question": "", "context": C})
+        with self.assertRaises(ValueError):
+            qa({"question": Q, "context": None})
+        with self.assertRaises(ValueError):
+            qa({"question": Q, "context": ""})
+
+        with self.assertRaises(ValueError):
+            qa([{"question": Q, "context": C}, {"question": None, "context": C}])
+        with self.assertRaises(ValueError):
+            qa([{"question": Q, "context": C}, {"question": "", "context": C}])
+
+        with self.assertRaises(ValueError):
+            qa([{"question": Q, "context": C}, {"question": Q, "context": None}])
+        with self.assertRaises(ValueError):
+            qa([{"question": Q, "context": C}, {"question": Q, "context": ""}])
+
+    def test_argument_handler_error_handling_odd(self):
+        qa = QuestionAnsweringArgumentHandler()
+        with self.assertRaises(ValueError):
+            qa(None)
+
+        with self.assertRaises(ValueError):
+            qa(Y=None)
+
+        with self.assertRaises(ValueError):
+            qa(1)

tests/test_trainer.py

@@ -18,13 +18,13 @@ import os
 import tempfile
 import unittest
 
-import datasets
 import numpy as np
 
 from transformers import AutoTokenizer, EvaluationStrategy, PretrainedConfig, TrainingArguments, is_torch_available
 from transformers.file_utils import WEIGHTS_NAME
 from transformers.testing_utils import (
     get_tests_dir,
+    require_datasets,
     require_optuna,
     require_sentencepiece,
     require_tokenizers,
@@ -340,7 +340,10 @@ class TrainerIntegrationTest(unittest.TestCase):
         self.assertTrue(np.array_equal(labels[0], trainer.eval_dataset.ys[0]))
         self.assertTrue(np.array_equal(labels[1], trainer.eval_dataset.ys[1]))
 
+    @require_datasets
     def test_trainer_with_datasets(self):
+        import datasets
+
         np.random.seed(42)
         x = np.random.normal(size=(64,)).astype(np.float32)
         y = 2.0 * x + 3.0 + np.random.normal(scale=0.1, size=(64,))
@@ -658,15 +661,17 @@ class TrainerHyperParameterIntegrationTest(unittest.TestCase):
         def hp_name(trial):
             return MyTrialShortNamer.shortname(trial.params)
 
-        trainer = get_regression_trainer(
-            learning_rate=0.1,
-            logging_steps=1,
-            evaluation_strategy=EvaluationStrategy.EPOCH,
-            num_train_epochs=4,
-            disable_tqdm=True,
-            load_best_model_at_end=True,
-            logging_dir="runs",
-            run_name="test",
-            model_init=model_init,
-        )
-        trainer.hyperparameter_search(direction="minimize", hp_space=hp_space, hp_name=hp_name, n_trials=4)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            trainer = get_regression_trainer(
+                output_dir=tmp_dir,
+                learning_rate=0.1,
+                logging_steps=1,
+                evaluation_strategy=EvaluationStrategy.EPOCH,
+                num_train_epochs=4,
+                disable_tqdm=True,
+                load_best_model_at_end=True,
+                logging_dir="runs",
+                run_name="test",
+                model_init=model_init,
+            )
+            trainer.hyperparameter_search(direction="minimize", hp_space=hp_space, hp_name=hp_name, n_trials=4)

tests/test_trainer_distributed.py

@@ -46,8 +46,6 @@ if is_torch_available():
 class TestTrainerDistributed(TestCasePlus):
     @require_torch_multigpu
     def test_trainer(self):
-        import os
-        print(os.system("nvidia-smi"))
 
         distributed_args = f"""
             -m torch.distributed.launch