Handle rebasing

Update transformers path
Cookiecutter path
2020-11-03 10:23:01 -05:00 · 2020-11-02 18:48:26 -05:00 · 2020-11-02 18:44:40 -05:00 · 2020-11-02 18:39:24 -05:00 · 2020-11-02 18:29:56 -05:00 · 2020-11-02 18:20:15 -05:00
151 changed files with 10305 additions and 656 deletions
@@ -4,16 +4,18 @@ on:
  push:
    branches:
      - master
+      - model-templates
    paths:
      - "src/**"
      - "tests/**"
      - ".github/**"
+      - "templates/**"
  # pull_request:
  repository_dispatch:


 jobs:
-  run_tests_torch_and_tf_gpu:
+  run_tests_torch_gpu:
    runs-on: [self-hosted, single-gpu]
    steps:
      - uses: actions/checkout@v2
@@ -32,7 +34,7 @@ jobs:
        id: cache
        with:
          path: .env
-          key: v1-tests_tf_torch_gpu-${{ hashFiles('setup.py') }}
+          key: v1.1-tests_torch_gpu-${{ hashFiles('setup.py') }}

      - name: Create new python env (on self-hosted runners we have to handle isolation ourselves)
        run: |
@@ -46,8 +48,7 @@ jobs:
        run: |
          source .env/bin/activate
          pip install --upgrade pip
-          pip install torch!=1.6.0
-          pip install .[sklearn,testing,onnxruntime]
+          pip install .[torch,sklearn,testing,onnxruntime]
          pip install git+https://github.com/huggingface/datasets

      - name: Are GPUs recognized by our DL frameworks
@@ -56,18 +57,25 @@ jobs:
          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 non-slow tests on GPU
-        env:
-          TF_FORCE_GPU_ALLOW_GROWTH: "true"
-          # TF_GPU_MEMORY_LIMIT: 4096
-          OMP_NUM_THREADS: 1
+      - name: Create model files
        run: |
          source .env/bin/activate
-          python -m pytest -n 2 --dist=loadfile -s tests
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/encoder-bert-tokenizer.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/pt-encoder-bert-tokenizer.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/seq-to-seq-roberta-tokenizer.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/standalone.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/tf-encoder-bert-tokenizer.json --path=templates/cookiecutter

+      - name: Run all non-slow tests on GPU
+        env:
+          OMP_NUM_THREADS: 1
+          CUDA_VISIBLE_DEVICES: 0
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 2 --dist=loadfile -s ./tests/

-  run_tests_torch_and_tf_multiple_gpu:
-    runs-on: [self-hosted, multi-gpu]
+  run_tests_tf_gpu:
+    runs-on: [self-hosted, single-gpu]
    steps:
      - uses: actions/checkout@v2
      - name: Python version
@@ -84,7 +92,66 @@ jobs:
        id: cache
        with:
          path: .env
-          key: v1-tests_tf_torch_multiple_gpu-${{ hashFiles('setup.py') }}
+          key: v1.1-tests_tf_gpu-${{ hashFiles('setup.py') }}
+
+      - name: Create new python env (on self-hosted runners we have to handle isolation ourselves)
+        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
+
+      - 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: Create model files
+        run: |
+          source .env/bin/activate
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/encoder-bert-tokenizer.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/pt-encoder-bert-tokenizer.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/seq-to-seq-roberta-tokenizer.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/standalone.json --path=templates/cookiecutter
+          transformers-cli add-new-model --testing --testing_file=templates/cookiecutter/tests/tf-encoder-bert-tokenizer.json --path=templates/cookiecutter
+
+      - name: Run all non-slow tests on GPU
+        env:
+          OMP_NUM_THREADS: 1
+          CUDA_VISIBLE_DEVICES: 0
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 2 --dist=loadfile -s ./tests/
+
+  run_tests_torch_multiple_gpu:
+    runs-on: [self-hosted, multi-gpu]
+    steps:
+      - uses: actions/checkout@v2
+      - name: Python version
+        run: |
+          which python
+          python --version
+          pip --version
+
+      - name: Current dir
+        run: pwd
+      - run: nvidia-smi
+
+      - name: Loading cache.
+        uses: actions/cache@v2
+        id: cache
+        with:
+          path: .env
+          key: v1.1-tests_torch_multiple_gpu-${{ hashFiles('setup.py') }}

      - name: Create new python env (on self-hosted runners we have to handle isolation ourselves)
        run: |
@@ -97,8 +164,7 @@ jobs:
        run: |
          source .env/bin/activate
          pip install --upgrade pip
-          pip install torch!=1.6.0
-          pip install .[sklearn,testing,onnxruntime]
+          pip install .[torch,sklearn,testing,onnxruntime]
          pip install git+https://github.com/huggingface/datasets

      - name: Are GPUs recognized by our DL frameworks
@@ -109,8 +175,54 @@ jobs:

      - name: Run all non-slow tests on GPU
        env:
-          TF_FORCE_GPU_ALLOW_GROWTH: "true"
-          # TF_GPU_MEMORY_LIMIT: 4096
+          OMP_NUM_THREADS: 1
+        run: |
+          source .env/bin/activate
+          python -m pytest -n 2 --dist=loadfile -s ./tests/
+
+  run_tests_tf_multiple_gpu:
+    runs-on: [self-hosted, multi-gpu]
+    steps:
+      - uses: actions/checkout@v2
+      - name: Python version
+        run: |
+          which python
+          python --version
+          pip --version
+
+      - name: Current dir
+        run: pwd
+      - run: nvidia-smi
+
+      - name: Loading cache.
+        uses: actions/cache@v2
+        id: cache
+        with:
+          path: .env
+          key: v1.1-tests_tf_multiple_gpu-${{ hashFiles('setup.py') }}
+
+      - name: Create new python env (on self-hosted runners we have to handle isolation ourselves)
+        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
+
+      - 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 non-slow tests on GPU
+        env:
          OMP_NUM_THREADS: 1
        run: |
          source .env/bin/activate
@@ -9,7 +9,7 @@ on:
    - cron: "0 0 * * *"

 jobs:
-  run_all_tests_torch_and_tf_gpu:
+  run_all_tests_torch_gpu:
    runs-on: [self-hosted, single-gpu]
    steps:
      - uses: actions/checkout@v2
@@ -19,7 +19,7 @@ jobs:
        id: cache
        with:
          path: .env
-          key: v1-slow_tests_tf_torch_gpu-${{ hashFiles('setup.py') }}
+          key: v  1.1-slow_tests_torch_gpu-${{ hashFiles('setup.py') }}

      - name: Python version
        run: |
@@ -44,9 +44,9 @@ jobs:
        run: |
          source .env/bin/activate
          pip install --upgrade pip
-          pip install torch!=1.6.0
-          pip install .[sklearn,testing,onnxruntime]
+          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: |
@@ -56,31 +56,29 @@ jobs:

      - name: Run all tests on GPU
        env:
-          TF_FORCE_GPU_ALLOW_GROWTH: "true"
          OMP_NUM_THREADS: 1
          RUN_SLOW: yes
        run: |
          source .env/bin/activate
-          python -m pytest -n 1 --dist=loadfile -s --make_reports=tests tests
+          python -m pytest -n 1 --dist=loadfile -s --make_reports=tests_torch tests

      - name: Failure short reports
        if: ${{ always() }}
-        run: cat reports/report_tests_failures_short.txt
+        run: cat reports/report_tests_torch_failures_short.txt
        
      - name: Run examples tests on GPU
        if: ${{ always() }}
        env:
-          TF_FORCE_GPU_ALLOW_GROWTH: "true"
          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 examples
+          python -m pytest -n 1 --dist=loadfile -s --make_reports=examples_torch examples

      - name: Failure short reports
        if: ${{ always() }}
-        run: cat reports/report_examples_failures_short.txt
+        run: cat reports/report_examples_torch_failures_short.txt

      - name: Run all pipeline tests on GPU
        if: ${{ always() }}
@@ -91,21 +89,99 @@ 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_pipeline tests
+          python -m pytest -n 1 --dist=loadfile -s -m is_pipeline_test --make_reports=tests_torch_pipeline tests

      - name: Failure short reports
        if: ${{ always() }}
-        run: cat reports/report_tests_pipeline_failures_short.txt
+        run: cat reports/report_tests_torch_pipeline_failures_short.txt

      - name: Test suite reports artifacts
        if: ${{ always() }}
        uses: actions/upload-artifact@v2
        with:
-          name: run_all_tests_torch_and_tf_gpu_test_reports
+          name: run_all_tests_torch_gpu_test_reports
          path: reports


-  run_all_tests_torch_and_tf_multiple_gpu:
+  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 tests
+          
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/report_tests_tf_failures_short.txt
+
+      - name: Run all pipeline tests on GPU
+        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 tests -m is_pipeline_test --make_reports=tests_tf_pipelines tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/report_tests_tf_pipelines_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]
    steps:
      - uses: actions/checkout@v2
@@ -115,16 +191,18 @@ jobs:
        id: cache
        with:
          path: .env
-          key: v1-slow_tests_tf_torch_multi_gpu-${{ hashFiles('setup.py') }}
+          key: v1.1-slow_tests_torch_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: |
@@ -133,13 +211,14 @@ jobs:
          which python
          python --version
          pip --version
+
      - name: Install dependencies
        run: |
          source .env/bin/activate
          pip install --upgrade pip
-          pip install torch!=1.6.0
-          pip install .[sklearn,testing,onnxruntime]
+          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: |
@@ -149,22 +228,94 @@ jobs:

      - name: Run all tests on GPU
        env:
-          TF_FORCE_GPU_ALLOW_GROWTH: "true"
          OMP_NUM_THREADS: 1
          RUN_SLOW: yes
        run: |
          source .env/bin/activate
-          python -m pytest -n 1 --dist=loadfile -s ./tests/ --durations=50
+          python -m pytest -n 1 --dist=loadfile -s --make_reports=tests_torch tests

-      - name: Run examples tests on GPU
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/report_tests_torch_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
-          pip install -r examples/requirements.txt
-          python -m pytest -n 1 --dist=loadfile -s examples --durations=50
+          python -m pytest -n 1 --dist=loadfile -s -m is_pipeline_test --make_reports=tests_torch_pipeline tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/report_tests_torch_pipeline_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 tests --make_reports=tests_tf tests
+
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/report_tests_tf_failures_short.txt

      - name: Run all pipeline tests on GPU
        env:
@@ -174,4 +325,16 @@ jobs:
          RUN_PIPELINE_TESTS: yes
        run: |
          source .env/bin/activate
-          python -m pytest -n 1 --dist=loadfile -s ./tests/ -m is_pipeline_test --durations=50
+          python -m pytest -n 1 --dist=loadfile -s tests -m is_pipeline_test --make_reports=tests_tf_pipelines tests
+          
+      - name: Failure short reports
+        if: ${{ always() }}
+        run: cat reports/report_tests_tf_pipelines_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
+          
@@ -95,3 +95,12 @@ See :obj:`transformers.BartForConditionalGeneration` for arguments to `forward`

 .. autoclass:: transformers.BlenderbotForConditionalGeneration
    :members:
+
+
+TFBlenderbotForConditionalGeneration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+See :obj:`transformers.TFBartForConditionalGeneration` for arguments to `forward` and `generate`
+
+.. autoclass:: transformers.TFBlenderbotForConditionalGeneration
+    :members:
@@ -129,3 +129,9 @@ MarianMTModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.MarianMTModel
+
+
+TFMarianMTModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFMarianMTModel
@@ -79,4 +79,11 @@ MBartForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.MBartForConditionalGeneration
-    :members: forward
+    :members:
+
+
+TFMBartForConditionalGeneration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFMBartForConditionalGeneration
+    :members:
@@ -95,3 +95,9 @@ PegasusForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.PegasusForConditionalGeneration
+
+
+TFPegasusForConditionalGeneration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFPegasusForConditionalGeneration
@@ -5,7 +5,9 @@ Running the examples requires PyTorch 1.3.1+ or TensorFlow 2.2+.

 Here is the list of all our examples:
 - **grouped by task** (all official examples work for multiple models)
- with information on whether they are **built on top of `Trainer`/`TFTrainer`** (if not, they still work, they might just lack some features),
+- with information on whether they are **built on top of `Trainer`/`TFTrainer`** (if not, they still work, they might
+  just lack some features),
+- whether or not they leverage the [🤗 Datasets](https://github.com/huggingface/datasets) library.
 - links to **Colab notebooks** to walk through the scripts and run them easily,
 - links to **Cloud deployments** to be able to deploy large-scale trainings in the Cloud with little to no setup.

@@ -31,19 +33,19 @@ git checkout tags/v3.4.0

 ## The Big Table of Tasks

-| Task | Example datasets | Trainer support | TFTrainer support | Colab
-|---|---|:---:|:---:|:---:|
-| [**`language-modeling`**](https://github.com/huggingface/transformers/tree/master/examples/language-modeling)       | Raw text        | ✅ | -  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)
-| [**`text-classification`**](https://github.com/huggingface/transformers/tree/master/examples/text-classification)   | GLUE, XNLI      | ✅ | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/trainer/01_text_classification.ipynb)
-| [**`token-classification`**](https://github.com/huggingface/transformers/tree/master/examples/token-classification) | CoNLL NER       | ✅ | ✅ | -
-| [**`multiple-choice`**](https://github.com/huggingface/transformers/tree/master/examples/multiple-choice)           | SWAG, RACE, ARC | ✅ | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ViktorAlm/notebooks/blob/master/MPC_GPU_Demo_for_TF_and_PT.ipynb)
-| [**`question-answering`**](https://github.com/huggingface/transformers/tree/master/examples/question-answering)     | SQuAD           | ✅ | ✅ | -
-| [**`text-generation`**](https://github.com/huggingface/transformers/tree/master/examples/text-generation)           | -               | n/a | n/a | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)
-| [**`distillation`**](https://github.com/huggingface/transformers/tree/master/examples/distillation)                 | All             | - | -  | -
-| [**`summarization`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                     | CNN/Daily Mail  | ✅  | - | -
-| [**`translation`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                       | WMT             | ✅  | - | -
-| [**`bertology`**](https://github.com/huggingface/transformers/tree/master/examples/bertology)                       | -               | - | - | -
-| [**`adversarial`**](https://github.com/huggingface/transformers/tree/master/examples/adversarial)                   | HANS            | ✅ | - | -
+| Task | Example datasets | Trainer support | TFTrainer support | 🤗 Datasets | Colab
+|---|---|:---:|:---:|:---:|:---:|
+| [**`language-modeling`**](https://github.com/huggingface/transformers/tree/master/examples/language-modeling)       | Raw text        | ✅ | -  | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)
+| [**`text-classification`**](https://github.com/huggingface/transformers/tree/master/examples/text-classification)   | GLUE, XNLI      | ✅ | ✅ | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/huggingface/notebooks/blob/master/examples/text_classification.ipynb)
+| [**`token-classification`**](https://github.com/huggingface/transformers/tree/master/examples/token-classification) | CoNLL NER       | ✅ | ✅ | - | -
+| [**`multiple-choice`**](https://github.com/huggingface/transformers/tree/master/examples/multiple-choice)           | SWAG, RACE, ARC | ✅ | ✅ | - | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ViktorAlm/notebooks/blob/master/MPC_GPU_Demo_for_TF_and_PT.ipynb)
+| [**`question-answering`**](https://github.com/huggingface/transformers/tree/master/examples/question-answering)     | SQuAD           | ✅ | ✅ | - | -
+| [**`text-generation`**](https://github.com/huggingface/transformers/tree/master/examples/text-generation)           | -               | n/a | n/a | - | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)
+| [**`distillation`**](https://github.com/huggingface/transformers/tree/master/examples/distillation)                 | All             | - | -  | - | -
+| [**`summarization`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                     | CNN/Daily Mail  | ✅  | - | - | -
+| [**`translation`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                       | WMT             | ✅  | - | - | -
+| [**`bertology`**](https://github.com/huggingface/transformers/tree/master/examples/bertology)                       | -               | - | - | - | -
+| [**`adversarial`**](https://github.com/huggingface/transformers/tree/master/examples/adversarial)                   | HANS            | ✅ | - | - | -


 <br>
@@ -1,16 +1,19 @@
-
 ## Language model training

-Based on the script [`run_language_modeling.py`](https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_language_modeling.py).
+Fine-tuning (or training from scratch) the library models for language modeling on a text dataset for GPT, GPT-2,
+ALBERT, BERT, DistilBERT, RoBERTa, XLNet... GPT and GPT-2 are trained or fine-tuned using a causal language modeling
+(CLM) loss while ALBERT, BERT, DistilBERT and RoBERTa are trained or fine-tuned using a masked language modeling (MLM)
+loss. XLNet uses permutation language modeling (PLM), you can find more information about the differences between those
+objectives in our [model summary](https://huggingface.co/transformers/model_summary.html).

-Fine-tuning (or training from scratch) the library models for language modeling on a text dataset for GPT, GPT-2, BERT, DistilBERT and RoBERTa. GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT, DistilBERT and RoBERTa
-are fine-tuned using a masked language modeling (MLM) loss.
+These scripts leverage the 🤗 Datasets library and the Trainer API. You can easily customize them to your needs if you
+need extra processing on your datasets.

-Before running the following example, you should get a file that contains text on which the language model will be
-trained or fine-tuned. A good example of such text is the [WikiText-2 dataset](https://blog.einstein.ai/the-wikitext-long-term-dependency-language-modeling-dataset/).
+**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).

-We will refer to two different files: `$TRAIN_FILE`, which contains text for training, and `$TEST_FILE`, which contains
-text that will be used for evaluation.
+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.

 ### GPT-2/GPT and causal language modeling

@@ -18,66 +21,105 @@ The following example fine-tunes GPT-2 on WikiText-2. We're using the raw WikiTe
 the tokenization). The loss here is that of causal language modeling.

 ```bash
-export TRAIN_FILE=/path/to/dataset/wiki.train.raw
-export TEST_FILE=/path/to/dataset/wiki.test.raw
-
-python run_language_modeling.py \
-    --output_dir=output \
-    --model_type=gpt2 \
-    --model_name_or_path=gpt2 \
+python run_clm.py \
+    --model_name_or_path gpt2 \
+    --dataset_name wikitext \
+    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
-    --train_data_file=$TRAIN_FILE \
    --do_eval \
-    --eval_data_file=$TEST_FILE
+    --output_dir /tmp/test-clm
 ```

 This takes about half an hour to train on a single K80 GPU and about one minute for the evaluation to run. It reaches
 a score of ~20 perplexity once fine-tuned on the dataset.

+To run on your own training and validation files, use the following command:
+
+```bash
+python run_clm.py \
+    --model_name_or_path gpt2 \
+    --train_file path_to_train_file \
+    --validation_file path_to_validation_file \
+    --do_train \
+    --do_eval \
+    --output_dir /tmp/test-clm
+```
+
+
 ### RoBERTa/BERT/DistilBERT and masked language modeling

 The following example fine-tunes RoBERTa on WikiText-2. Here too, we're using the raw WikiText-2. The loss is different
 as BERT/RoBERTa have a bidirectional mechanism; we're therefore using the same loss that was used during their
 pre-training: masked language modeling.

-In accordance to the RoBERTa paper, we use dynamic masking rather than static masking. The model may, therefore, converge
-slightly slower (over-fitting takes more epochs).
-
-We use the `--mlm` flag so that the script may change its loss function.
-
-If using whole-word masking, use both the`--mlm` and `--wwm` flags.
+In accordance to the RoBERTa paper, we use dynamic masking rather than static masking. The model may, therefore,
+converge slightly slower (over-fitting takes more epochs).

 ```bash
-export TRAIN_FILE=/path/to/dataset/wiki.train.raw
-export TEST_FILE=/path/to/dataset/wiki.test.raw
-
-python run_language_modeling.py \
-    --output_dir=output \
-    --model_type=roberta \
-    --model_name_or_path=roberta-base \
+python run_mlm.py \
+    --model_name_or_path roberta-base \
+    --dataset_name wikitext \
+    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
-    --train_data_file=$TRAIN_FILE \
    --do_eval \
-    --eval_data_file=$TEST_FILE \
-    --mlm \
-    --whole_word_mask
+    --output_dir /tmp/test-mlm
 ```

-For Chinese models, it's same with English model with only `--mlm`. If using whole-word masking, we need to generate a reference files, because it's char level.
+To run on your own training and validation files, use the following command:

-**Q :** Why ref file ?
+```bash
+python run_clm.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
+```

-**A :** Suppose we have a Chinese sentence like : `我喜欢你` The original Chinese-BERT will tokenize it as `['我','喜','欢','你']` in char level.
-Actually, `喜欢` is a whole word. For whole word mask proxy, We need res like `['我','喜','##欢','你']`.
-So we need a ref file to tell model which pos of BERT original token should be added `##`.
+If your dataset is organized with one sample per line, you can use the `--line_by_line` flag (otherwise the script
+concatenates all texts and then splits them in blocks of the same length).
+
+**Note:** On TPU, you should use the flag `--pad_to_max_length` in conjunction with the `--line_by_line` flag to make
+sure all your batches have the same length.
+
+### Whole word masking
+
+The BERT authors released a new version of BERT using Whole Word Masking in May 2019. Instead of masking randomly
+selected tokens (which may be part of words), they mask randomly selected words (masking all the tokens corresponding
+to that word). This technique has been refined for Chinese in [this paper](https://arxiv.org/abs/1906.08101).
+
+To fine-tune a model using whole word masking, use the following script:
+
+python run_mlm_wwm.py \
+    --model_name_or_path roberta-base \
+    --dataset_name wikitext \
+    --dataset_config_name wikitext-2-raw-v1 \
+    --do_train \
+    --do_eval \
+    --output_dir /tmp/test-mlm-wwm
+```
+
+For Chinese models, we need to generate a reference files (which requires the ltp library), because it's tokenized at
+the character level.
+
+**Q :** Why a reference file?
+
+**A :** Suppose we have a Chinese sentence like: `我喜欢你` The original Chinese-BERT will tokenize it as
+`['我','喜','欢','你']` (character level). But `喜欢` is a whole word. For whole word masking proxy, we need a result
+like `['我','喜','##欢','你']`, so we need a reference file to tell the model which position of the BERT original token
+should be added `##`.

 **Q :** Why LTP ?

-**A :** Cause the best known Chinese WWM BERT is [Chinese-BERT-wwm](https://github.com/ymcui/Chinese-BERT-wwm) by HIT. It works well on so many Chines Task like CLUE (Chinese GLUE).
-They use LTP, so if we want to fine-tune their model, we need LTP.
+**A :** Cause the best known Chinese WWM BERT is [Chinese-BERT-wwm](https://github.com/ymcui/Chinese-BERT-wwm) by HIT.
+It works well on so many Chines Task like CLUE (Chinese GLUE). They use LTP, so if we want to fine-tune their model,
+we need LTP.

 Now LTP only only works well on `transformers==3.2.0`. So we don't add it to requirements.txt.
-You need to check to `3.2.0` for `run_chinese_ref.py`. And the code  could be found in `examples/contrib`.
+You need to create a separate environment with this version of Transformers to run the `run_chinese_ref.py` script that
+will create the reference files. The script is in `examples/contrib`. Once in the proper environment, run the
+following:


 ```bash
@@ -87,33 +129,29 @@ export BERT_RESOURCE=/path/to/bert/tokenizer
 export SAVE_PATH=/path/to/data/ref.txt

 python examples/contrib/run_chinese_ref.py \
-    --file_name=$TRAIN_FILE \
-    --ltp=$LTP_RESOURCE \
-    --bert=$BERT_RESOURCE \
-    --save_path=$SAVE_PATH 
+    --file_name=path_to_train_or_eval_file \
+    --ltp=path_to_ltp_tokenizer \
+    --bert=path_to_bert_tokenizer \
+    --save_path=path_to_reference_file
 ```
-Now Chinese Ref is only supported by `LineByLineWithRefDataset` Class, so we need add `line_by_line` flag: 
+
+Then you can run the script like this: 


 ```bash
-export TRAIN_FILE=/path/to/dataset/wiki.train.raw
-export TEST_FILE=/path/to/dataset/wiki.test.raw
-export REF_FILE=/path/to/ref.txt
-
-python run_language_modeling.py \
-    --output_dir=output \
-    --model_type=roberta \
-    --model_name_or_path=roberta-base \
+python run_mlm_wwm.py \
+    --model_name_or_path roberta-base \
+    --train_file path_to_train_file \
+    --validation_file path_to_validation_file \
+    --train_ref_file path_to_train_chinese_ref_file \
+    --validation_ref_file path_to_validation_chinese_ref_file \
    --do_train \
-    --train_data_file=$TRAIN_FILE \
-    --chinese_ref_file=$REF_FILE \
    --do_eval \
-    --eval_data_file=$TEST_FILE \
-    --mlm \
-    --line_by_line \
-    --whole_word_mask
+    --output_dir /tmp/test-mlm-wwm
 ```

+**Note:** On TPU, you should the flag `--pad_to_max_length` to make sure all your batches have the same length.
+
 ### XLNet and permutation language modeling

 XLNet uses a different training objective, which is permutation language modeling. It is an autoregressive method 
@@ -126,15 +164,32 @@ context length for permutation language modeling.
 The `--max_span_length` flag may also be used to limit the length of a span of masked tokens used 
 for permutation language modeling.

-```bash
-export TRAIN_FILE=/path/to/dataset/wiki.train.raw
-export TEST_FILE=/path/to/dataset/wiki.test.raw
+Here is how to fine-tun XLNet on wikitext-2:

-python run_language_modeling.py \
-    --output_dir=output \
+```bash
+python run_plm.py \
    --model_name_or_path=xlnet-base-cased \
+    --dataset_name wikitext \
+    --dataset_config_name wikitext-2-raw-v1 \
    --do_train \
-    --train_data_file=$TRAIN_FILE \
    --do_eval \
-    --eval_data_file=$TEST_FILE \
+    --output_dir /tmp/test-plm
 ```
+
+To fine-tune it on your own training and validation file, run:
+
+```bash
+python run_plm.py \
+    --model_name_or_path=xlnet-base-cased \
+    --train_file path_to_train_file \
+    --validation_file path_to_validation_file \
+    --do_train \
+    --do_eval \
+    --output_dir /tmp/test-plm
+```
+
+If your dataset is organized with one sample per line, you can use the `--line_by_line` flag (otherwise the script
+concatenates all texts and then splits them in blocks of the same length).
+
+**Note:** On TPU, you should use the flag `--pad_to_max_length` in conjunction with the `--line_by_line` flag to make
+sure all your batches have the same length.
@@ -175,10 +175,10 @@ def main():

    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
-    # (the dataset will be downloaded automatically from the datasets Hub
+    # (the dataset will be downloaded automatically from the datasets Hub).
    #
-    # For CSV/JSON files, this script will use the column called 'text' or the first column. You can easily tweak this
-    # behavior (see below)
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
@@ -259,14 +259,14 @@ def main():
    )

    if data_args.block_size <= 0:
-        block_size = tokenizer.max_len
+        block_size = tokenizer.model_max_length
    else:
-        if data_args.block_size > tokenizer.max_len:
+        if data_args.block_size > tokenizer.model_max_length:
            logger.warn(
                f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model"
-                f"({tokenizer.max_len}). Using block_size={tokenizer.max_len}."
+                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
            )
-        block_size = min(data_args.block_size, tokenizer.max_len)
+        block_size = min(data_args.block_size, tokenizer.model_max_length)

    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
    def group_texts(examples):
@@ -116,6 +116,17 @@ class DataTrainingArguments:
    mlm_probability: float = field(
        default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}
    )
+    line_by_line: bool = field(
+        default=False,
+        metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad all samples to `max_seq_length`. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch."
+        },
+    )

    def __post_init__(self):
        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
@@ -246,18 +257,73 @@ def main():
        column_names = datasets["validation"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

-    def tokenize_function(examples):
-        # Remove empty lines
-        examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
-        return tokenizer(examples["text"], truncation=True, max_length=data_args.max_seq_length)
+    if data_args.line_by_line:
+        # When using line_by_line, we just tokenize each nonempty line.
+        padding = "max_length" if data_args.pad_to_max_length else False

-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        num_proc=data_args.preprocessing_num_workers,
-        remove_columns=[text_column_name],
-        load_from_cache_file=not data_args.overwrite_cache,
-    )
+        def tokenize_function(examples):
+            # Remove empty lines
+            examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
+            return tokenizer(examples["text"], padding=padding, truncation=True, max_length=data_args.max_seq_length)
+
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=[text_column_name],
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+    else:
+        # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.
+        def tokenize_function(examples):
+            return tokenizer(examples[text_column_name])
+
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=[text_column_name],
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
+        if data_args.max_seq_length is None:
+            max_seq_length = tokenizer.model_max_length
+        else:
+            if data_args.max_seq_length > tokenizer.model_max_length:
+                logger.warn(
+                    f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
+                    f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+                )
+            max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
+
+        # Main data processing function that will concatenate all texts from our dataset and generate chunks of
+        # max_seq_length.
+        def group_texts(examples):
+            # Concatenate all texts.
+            concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+            total_length = len(concatenated_examples[list(examples.keys())[0]])
+            # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+            # customize this part to your needs.
+            total_length = (total_length // max_seq_length) * max_seq_length
+            # Split by chunks of max_len.
+            result = {
+                k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
+                for k, t in concatenated_examples.items()
+            }
+            return result
+
+        # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
+        # remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
+        # might be slower to preprocess.
+        #
+        # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+        # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+        tokenized_datasets = tokenized_datasets.map(
+            group_texts,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )

    # Data collator
    # This one will take care of randomly masking the tokens.
@@ -0,0 +1,334 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Team All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for masked language modeling (BERT, ALBERT, RoBERTa...) with whole word masking on a
+text file or a dataset.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=masked-lm
+"""
+# You can also adapt this script on your own masked language modeling task. Pointers for this are left as comments.
+
+import json
+import logging
+import math
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Optional
+
+from datasets import Dataset, load_dataset
+
+import transformers
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_FOR_MASKED_LM_MAPPING,
+    AutoConfig,
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    DataCollatorForWholeWordMask,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+from transformers.trainer_utils import is_main_process
+
+
+logger = logging.getLogger(__name__)
+MODEL_CONFIG_CLASSES = list(MODEL_FOR_MASKED_LM_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
+    """
+
+    model_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The model checkpoint for weights initialization."
+            "Don't set if you want to train a model from scratch."
+        },
+    )
+    model_type: Optional[str] = field(
+        default=None,
+        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
+    )
+    train_ref_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input train ref data file for whole word masking in Chinese."},
+    )
+    validation_ref_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input validation ref data file for whole word masking in Chinese."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    max_seq_length: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated. Default to the max input length of the model."
+        },
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    mlm_probability: float = field(
+        default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad all samples to `max_seq_length`. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch."
+        },
+    )
+
+    def __post_init__(self):
+        if self.train_file is not None:
+            extension = self.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
+        if self.validation_file is not None:
+            extension = self.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
+
+
+def add_chinese_references(dataset, ref_file):
+    with open(ref_file, "r", encoding="utf-8") as f:
+        refs = [json.loads(line) for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]
+    assert len(dataset) == len(refs)
+
+    dataset_dict = {c: dataset[c] for c in dataset.column_names}
+    dataset_dict["chinese_ref"] = refs
+    return Dataset.from_dict(dataset_dict)
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    if (
+        os.path.exists(training_args.output_dir)
+        and os.listdir(training_args.output_dir)
+        and training_args.do_train
+        and not training_args.overwrite_output_dir
+    ):
+        raise ValueError(
+            f"Output directory ({training_args.output_dir}) already exists and is not empty."
+            "Use --overwrite_output_dir to overcome."
+        )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO if is_main_process(training_args.local_rank) else logging.WARN,
+    )
+
+    # Log on each process the small summary:
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+    # Set the verbosity to info of the Transformers logger (on main process only):
+    if is_main_process(training_args.local_rank):
+        transformers.utils.logging.set_verbosity_info()
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    data_files = {}
+    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
+    extension = data_args.train_file.split(".")[-1]
+    if extension == "txt":
+        extension = "text"
+    datasets = load_dataset(extension, data_files=data_files)
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    if model_args.config_name:
+        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
+    elif model_args.model_name_or_path:
+        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
+    else:
+        config = CONFIG_MAPPING[model_args.model_type]()
+        logger.warning("You are instantiating a new config instance from scratch.")
+
+    if model_args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+        )
+    elif model_args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+        )
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    if model_args.model_name_or_path:
+        model = AutoModelForMaskedLM.from_pretrained(
+            model_args.model_name_or_path,
+            from_tf=bool(".ckpt" in model_args.model_name_or_path),
+            config=config,
+            cache_dir=model_args.cache_dir,
+        )
+    else:
+        logger.info("Training new model from scratch")
+        model = AutoModelForMaskedLM.from_config(config)
+
+    model.resize_token_embeddings(len(tokenizer))
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    if training_args.do_train:
+        column_names = datasets["train"].column_names
+    else:
+        column_names = datasets["validation"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    padding = "max_length" if data_args.pad_to_max_length else False
+
+    def tokenize_function(examples):
+        # Remove empty lines
+        examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
+        return tokenizer(examples["text"], padding=padding, truncation=True, max_length=data_args.max_seq_length)
+
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        num_proc=data_args.preprocessing_num_workers,
+        remove_columns=[text_column_name],
+        load_from_cache_file=not data_args.overwrite_cache,
+    )
+
+    # Add the chinese references if provided
+    if data_args.train_ref_file is not None:
+        tokenized_datasets["train"] = add_chinese_references(tokenized_datasets["train"], data_args.train_ref_file)
+    if data_args.valid_ref_file is not None:
+        tokenized_datasets["validation"] = add_chinese_references(
+            tokenized_datasets["validation"], data_args.validation_ref_file
+        )
+
+    # Data collator
+    # This one will take care of randomly masking the tokens.
+    data_collator = DataCollatorForWholeWordMask(tokenizer=tokenizer, mlm_probability=data_args.mlm_probability)
+
+    # Initialize our Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"] if training_args.do_train else None,
+        eval_dataset=tokenized_datasets["validation"] if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # Training
+    if training_args.do_train:
+        trainer.train(
+            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
+        )
+        trainer.save_model()  # Saves the tokenizer too for easy upload
+
+    # Evaluation
+    results = {}
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+
+        eval_output = trainer.evaluate()
+
+        perplexity = math.exp(eval_output["eval_loss"])
+        results["perplexity"] = perplexity
+
+        output_eval_file = os.path.join(training_args.output_dir, "eval_results_mlm_wwm.txt")
+        if trainer.is_world_process_zero():
+            with open(output_eval_file, "w") as writer:
+                logger.info("***** Eval results *****")
+                for key, value in results.items():
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+    return results
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()
@@ -96,7 +96,7 @@ class DataTrainingArguments:
        default=None,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
-            "than this will be truncated."
+            "than this will be truncated. Default to the max input length of the model."
        },
    )
    preprocessing_num_workers: Optional[int] = field(
@@ -113,6 +113,17 @@ class DataTrainingArguments:
    max_span_length: int = field(
        default=5, metadata={"help": "Maximum length of a span of masked tokens for permutation language modeling."}
    )
+    line_by_line: bool = field(
+        default=False,
+        metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad all samples to `max_seq_length`. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch."
+        },
+    )

    def __post_init__(self):
        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
@@ -172,10 +183,10 @@ def main():

    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
-    # (the dataset will be downloaded automatically from the datasets Hub
+    # (the dataset will be downloaded automatically from the datasets Hub).
    #
-    # For CSV/JSON files, this script will use the column called 'text' or the first column. You can easily tweak this
-    # behavior (see below)
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
@@ -243,18 +254,73 @@ def main():
        column_names = datasets["validation"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

-    def tokenize_function(examples):
-        # Remove empty lines
-        examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
-        return tokenizer(examples["text"], truncation=True, max_length=data_args.max_seq_length)
+    if data_args.line_by_line:
+        # When using line_by_line, we just tokenize each nonempty line.
+        padding = "max_length" if data_args.pad_to_max_length else False

-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        num_proc=data_args.preprocessing_num_workers,
-        remove_columns=[text_column_name],
-        load_from_cache_file=not data_args.overwrite_cache,
-    )
+        def tokenize_function(examples):
+            # Remove empty lines
+            examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
+            return tokenizer(examples["text"], padding=padding, truncation=True, max_length=data_args.max_seq_length)
+
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=[text_column_name],
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+    else:
+        # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.
+        def tokenize_function(examples):
+            return tokenizer(examples[text_column_name])
+
+        tokenized_datasets = datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=[text_column_name],
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
+        if data_args.max_seq_length is None:
+            max_seq_length = tokenizer.model_max_length
+        else:
+            if data_args.max_seq_length > tokenizer.model_max_length:
+                logger.warn(
+                    f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
+                    f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+                )
+            max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
+
+        # Main data processing function that will concatenate all texts from our dataset and generate chunks of
+        # max_seq_length.
+        def group_texts(examples):
+            # Concatenate all texts.
+            concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+            total_length = len(concatenated_examples[list(examples.keys())[0]])
+            # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+            # customize this part to your needs.
+            total_length = (total_length // max_seq_length) * max_seq_length
+            # Split by chunks of max_len.
+            result = {
+                k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
+                for k, t in concatenated_examples.items()
+            }
+            return result
+
+        # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
+        # remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
+        # might be slower to preprocess.
+        #
+        # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+        # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+        tokenized_datasets = tokenized_datasets.map(
+            group_texts,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )

    # Data collator
    data_collator = DataCollatorForPermutationLanguageModeling(
@@ -20,12 +20,6 @@ from transformers.optimization import (
 from transformers.trainer_pt_utils import get_tpu_sampler


-try:
-    from .utils import label_smoothed_nll_loss
-except ImportError:
-    from utils import label_smoothed_nll_loss
-
-
 logger = logging.get_logger(__name__)

 arg_to_scheduler = {
@@ -64,6 +58,17 @@ class Seq2SeqTrainer(Trainer):
                f"The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for padding.."
            )

+        if self.args.label_smoothing == 0:
+            self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
+        else:
+            # dynamically import label_smoothed_nll_loss
+            try:
+                from .utils import label_smoothed_nll_loss
+            except ImportError:
+                from utils import label_smoothed_nll_loss
+
+            self.loss_fn = label_smoothed_nll_loss
+
    def create_optimizer_and_scheduler(self, num_training_steps: int):
        """
        Setup the optimizer and the learning rate scheduler.
@@ -135,9 +140,7 @@ class Seq2SeqTrainer(Trainer):
            if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
                # force training to ignore pad token
                logits = model(**inputs, use_cache=False)[0]
-
-                loss_fct = torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
-                loss = loss_fct(logits.view(-1, logits.shape[-1]), labels.view(-1))
+                loss = self.loss_fn(logits.view(-1, logits.shape[-1]), labels.view(-1))
            else:
                # compute usual loss via models
                loss, logits = model(**inputs, labels=labels, use_cache=False)[:2]
@@ -145,9 +148,7 @@ class Seq2SeqTrainer(Trainer):
            # compute label smoothed loss
            logits = model(**inputs, use_cache=False)[0]
            lprobs = torch.nn.functional.log_softmax(logits, dim=-1)
-            loss, _ = label_smoothed_nll_loss(
-                lprobs, labels, self.args.label_smoothing, ignore_index=self.config.pad_token_id
-            )
+            loss, _ = self.loss_fn(lprobs, labels, self.args.label_smoothing, ignore_index=self.config.pad_token_id)
        return loss, logits

    def compute_loss(self, model, inputs):
@@ -177,7 +177,7 @@ def main():
    set_seed(training_args.seed)

    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
-    # or specify a GLUE benchmark task (the dataset will be downloaded automatically from the datasets Hub
+    # or specify a GLUE benchmark task (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For CSV/JSON files, this script will use as labels the column called 'label' and as pair of sentences the
    # sentences in columns called 'sentence1' and 'sentence2' if such column exists or the first two columns not named
@@ -8,13 +8,13 @@ allowing both adaptive width and depth using knowledge distillation.
 * This code is modified based on the repository developed by Hugging Face: [Transformers v2.1.1](https://github.com/huggingface/transformers/tree/v2.1.1), and is released in [GitHub](https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/DynaBERT).

 ### Reference
-Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Qun Liu.
+Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, Qun Liu.
 [DynaBERT: Dynamic BERT with Adaptive Width and Depth](https://arxiv.org/abs/2004.04037).
 ```
@inproceedings{hou2020dynabert,
  title = {DynaBERT: Dynamic BERT with Adaptive Width and Depth},
-  author = {Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Qun Liu},  
-  booktitle = {NeurIPS},
+  author = {Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, Qun Liu},  
+  booktitle = {Advances in Neural Information Processing Systems},
  year = {2020}
 }
 ```
@@ -8,13 +8,13 @@ allowing both adaptive width and depth using knowledge distillation.
 * This code is modified based on the repository developed by Hugging Face: [Transformers v2.1.1](https://github.com/huggingface/transformers/tree/v2.1.1), and is released in [GitHub](https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/DynaBERT).

 ### Reference
-Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Qun Liu.
+Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, Qun Liu.
 [DynaBERT: Dynamic BERT with Adaptive Width and Depth](https://arxiv.org/abs/2004.04037).
 ```
@inproceedings{hou2020dynabert,
  title = {DynaBERT: Dynamic BERT with Adaptive Width and Depth},
-  author = {Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Qun Liu},  
-  booktitle = {NeurIPS},
+  author = {Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, Qun Liu},  
+  booktitle = {Advances in Neural Information Processing Systems},
  year = {2020}
 }
 ```
@@ -0,0 +1,25 @@
+---
+language: 
+- bn 
+tags:
+- MaskedLM
+- Bengali
+---
+# Indic-Transformers Bengali BERT
+## Model description
+This is a BERT language model pre-trained on ~3 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-bn-bert')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-bn-bert')
+text = "আপনি কেমন আছেন?"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 6, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- bn 
+tags:
+- MaskedLM
+- Bengali
+- DistilBERT
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Bengali DistilBERT
+## Model description
+This is a DistilBERT language model pre-trained on ~6 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-bn-distilbert')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-bn-distilbert')
+text = "আপনি কেমন আছেন?"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- bn 
+tags:
+- MaskedLM
+- Bengali
+- RoBERTa
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Bengali RoBERTa
+## Model description
+This is a RoBERTa language model pre-trained on ~6 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-bn-roberta')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-bn-roberta')
+text = "আপনি কেমন আছেন?"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 10, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- bn 
+tags:
+- MaskedLM
+- Bengali
+- XLMRoBERTa
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Bengali XLMRoBERTa
+## Model description
+This is a XLMRoBERTa language model pre-trained on ~3 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-bn-xlmroberta')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-bn-xlmroberta')
+text = "আপনি কেমন আছেন?"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- hi 
+tags:
+- MaskedLM
+- Hindi
+- BERT
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Hindi BERT
+## Model description
+This is a BERT language model pre-trained on ~3 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-hi-bert')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-hi-bert')
+text = "आपका स्वागत हैं"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- hi 
+tags:
+- MaskedLM
+- Hindi
+- DistilBERT
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Hindi DistilBERT
+## Model description
+This is a DistilBERT language model pre-trained on ~10 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-hi-distilbert')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-hi-distilbert')
+text = "आपका स्वागत हैं"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- hi 
+tags:
+- MaskedLM
+- Hindi
+- RoBERTa
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Hindi RoBERTa
+## Model description
+This is a RoBERTa language model pre-trained on ~10 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-hi-roberta')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-hi-roberta')
+text = "आपका स्वागत हैं"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 11, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- hi 
+tags:
+- MaskedLM
+- Hindi
+- XLMRoBERTa
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Hindi XLMRoBERTa
+## Model description
+This is a XLMRoBERTa language model pre-trained on ~3 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-hi-xlmroberta')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-hi-xlmroberta')
+text = "आपका स्वागत हैं"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- te
+tags:
+- MaskedLM
+- Telugu
+- BERT
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Telugu BERT
+## Model description
+This is a BERT language model pre-trained on ~1.6 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-te-bert')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-te-bert')
+text = "మీరు ఎలా ఉన్నారు"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- te
+tags:
+- MaskedLM
+- Telugu
+- DistilBERT
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Telugu DistilBERT
+## Model description
+This is a DistilBERT language model pre-trained on ~2 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-te-distilbert')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-te-distilbert')
+text = "మీరు ఎలా ఉన్నారు"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- te
+tags:
+- MaskedLM
+- Telugu
+- RoBERTa
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Telugu RoBERTa
+## Model description
+This is a RoBERTa language model pre-trained on ~2 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-te-roberta')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-te-roberta')
+text = "మీరు ఎలా ఉన్నారు"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 14, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,29 @@
+---
+language: 
+- te
+tags:
+- MaskedLM
+- Telugu
+- XLMRoBERTa
+- Question-Answering
+- Token Classification
+- Text Classification
+---
+# Indic-Transformers Telugu XLMRoBERTa
+## Model description
+This is a XLMRoBERTa language model pre-trained on ~1.6 GB of monolingual training corpus. The pre-training data was majorly taken from [OSCAR](https://oscar-corpus.com/).
+This model can be fine-tuned on various downstream tasks like text-classification, POS-tagging, question-answering, etc. Embeddings from this model can also be used for feature-based training.
+## Intended uses & limitations
+#### How to use
+```
+from transformers import AutoTokenizer, AutoModel
+tokenizer = AutoTokenizer.from_pretrained('neuralspace-reverie/indic-transformers-te-xlmroberta')
+model = AutoModel.from_pretrained('neuralspace-reverie/indic-transformers-te-xlmroberta')
+text = "మీరు ఎలా ఉన్నారు"
+input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+out = model(input_ids)[0]
+print(out.shape)
+# out = [1, 5, 768] 
+```
+#### Limitations and bias
+The original language model has been trained using `PyTorch` and hence the use of `pytorch_model.bin` weights file is recommended. The h5 file for `Tensorflow` has been generated manually by commands suggested [here](https://huggingface.co/transformers/model_sharing.html).
@@ -0,0 +1,16 @@
+---
+language: en
+license: apache-2.0
+datasets:
+- cnn_dailymail
+tags:
+- summarization
+---
+
+Bert2Bert Summarization with 🤗EncoderDecoder Framework
+This model is a warm-started *BERT2BERT* model fine-tuned on the *CNN/Dailymail* summarization dataset.
+
+The model achieves a **18.22** ROUGE-2 score on *CNN/Dailymail*'s test dataset.
+
+For more details on how the model was fine-tuned, please refer to 
+[this](https://colab.research.google.com/drive/1Ekd5pUeCX7VOrMx94_czTkwNtLN32Uyu?usp=sharing) notebook.
@@ -0,0 +1,16 @@
+---
+language: en
+license: apache-2.0
+datasets:
+- xsum
+tags:
+- summarization
+---
+
+Shared RoBERTa2RoBERTa Summarization with 🤗EncoderDecoder Framework
+This model is a warm-started *RoBERTaShared* model fine-tuned on the *BBC XSum* summarization dataset.
+
+The model achieves a **16.89** ROUGE-2 score on *BBC XSUM*'s test dataset.
+
+For more details on how the model was fine-tuned, please refer to 
+[this](https://colab.research.google.com/drive/1Ekd5pUeCX7VOrMx94_czTkwNtLN32Uyu?usp=sharing) notebook.
@@ -26,7 +26,7 @@ Pull Request so it can be included under the Community notebooks.

 | Notebook     |      Description      |      Author      |      |
 |:----------|:-------------|:-------------|------:|
-| [Train T5 in Tensoflow 2 ](https://github.com/snapthat/TF-T5-text-to-text) | How to train T5 for any task using Tensorflow 2. This notebook demonstrates a Question & Answer task implemented in Tensorflow 2 using SQUAD | [Muhammad Harris](https://github.com/HarrisDePerceptron) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/snapthat/TF-T5-text-to-text/blob/master/snapthatT5/notebooks/TF-T5-Datasets%20Training.ipynb) |
+| [Train T5 in Tensorflow 2 ](https://github.com/snapthat/TF-T5-text-to-text) | How to train T5 for any task using Tensorflow 2. This notebook demonstrates a Question & Answer task implemented in Tensorflow 2 using SQUAD | [Muhammad Harris](https://github.com/HarrisDePerceptron) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/snapthat/TF-T5-text-to-text/blob/master/snapthatT5/notebooks/TF-T5-Datasets%20Training.ipynb) |
 | [Train T5 on TPU](https://github.com/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb)  | How to train T5 on SQUAD with Transformers and Nlp | [Suraj Patil](https://github.com/patil-suraj) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb#scrollTo=QLGiFCDqvuil) |
 | [Fine-tune T5 for Classification and Multiple Choice](https://github.com/patil-suraj/exploring-T5/blob/master/t5_fine_tuning.ipynb)  | How to fine-tune T5 for classification and multiple choice tasks using a text-to-text format with PyTorch Lightning |  [Suraj Patil](https://github.com/patil-suraj) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/t5_fine_tuning.ipynb) |
 | [Fine-tune DialoGPT on New Datasets and Languages](https://github.com/ncoop57/i-am-a-nerd/blob/master/_notebooks/2020-05-12-chatbot-part-1.ipynb)  | How to fine-tune the DialoGPT model on a new dataset for open-dialog conversational chatbots |  [Nathan Cooper](https://github.com/ncoop57) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ncoop57/i-am-a-nerd/blob/master/_notebooks/2020-05-12-chatbot-part-1.ipynb) |
@@ -51,3 +51,5 @@ Pull Request so it can be included under the Community notebooks.
 |[Fine-tune Roberta for sentiment analysis](https://github.com/DhavalTaunk08/NLP_scripts/blob/master/sentiment_analysis_using_roberta.ipynb) | How to fine-tune an Roberta model for sentiment analysis | [Dhaval Taunk](https://github.com/DhavalTaunk08) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/DhavalTaunk08/NLP_scripts/blob/master/sentiment_analysis_using_roberta.ipynb)|
 |[Evaluating Question Generation Models](https://github.com/flexudy-pipe/qugeev) | How accurate are the answers to questions generated by your seq2seq transformer model? | [Pascal Zoleko](https://github.com/zolekode) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1bpsSqCQU-iw_5nNoRm_crPq6FRuJthq_?usp=sharing)|
 |[Classify text with DistilBERT and Tensorflow](https://github.com/peterbayerle/huggingface_notebook/blob/main/distilbert_tf.ipynb) | How to fine-tune DistilBERT for text classification in TensorFlow | [Peter Bayerle](https://github.com/peterbayerle) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/peterbayerle/huggingface_notebook/blob/main/distilbert_tf.ipynb)|
+|[Leverage BERT for Encoder-Decoder Summarization on CNN/Dailymail](https://github.com/patrickvonplaten/notebooks/blob/master/BERT2BERT_for_CNN_Dailymail.ipynb) | How to warm-start a *EncoderDecoderModel* with a *bert-base-uncased* checkpoint for summarization on CNN/Dailymail | [Patrick von Platen](https://github.com/patrickvonplaten) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/patrickvonplaten/notebooks/blob/master/BERT2BERT_for_CNN_Dailymail.ipynb)|
+|[Leverage RoBERTa for Encoder-Decoder Summarization on BBC XSum](https://github.com/patrickvonplaten/notebooks/blob/master/RoBERTaShared_for_BBC_XSum.ipynb) | How to warm-start a shared *EncoderDecoderModel* with a *roberta-base* checkpoint for summarization on BBC/XSum | [Patrick von Platen](https://github.com/patrickvonplaten) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/patrickvonplaten/notebooks/blob/master/RoBERTaShared_for_BBC_XSum.ipynb)|
@@ -98,12 +98,13 @@ else:

 extras["tokenizers"] = ["tokenizers==0.9.2"]
 extras["onnxruntime"] = ["onnxruntime>=1.4.0", "onnxruntime-tools>=1.4.2"]
+extras["modelcreation"] = ["cookiecutter==1.7.2"]

 extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"]

 extras["sentencepiece"] = ["sentencepiece==0.1.91"]
 extras["retrieval"] = ["faiss-cpu", "datasets"]
-extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator", "parameterized", "psutil"] + extras["retrieval"]
+extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator", "parameterized", "psutil"] + extras["retrieval"] + extras["modelcreation"]
 # 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["quality"] = ["black >= 20.8b1", "isort >= 5.5.4", "flake8 >= 3.8.3"]
@@ -111,7 +112,7 @@ extras["quality"] = ["black >= 20.8b1", "isort >= 5.5.4", "flake8 >= 3.8.3"]

 extras["all"] = extras["tf"] + extras["torch"] + extras["flax"] + extras["sentencepiece"] + extras["tokenizers"]

-extras["dev"] = extras["all"] + extras["testing"] + extras["quality"] + extras["ja"] + extras["docs"] + extras["sklearn"]
+extras["dev"] = extras["all"] + extras["testing"] + extras["quality"] + extras["ja"] + extras["docs"] + extras["sklearn"] + extras["modelcreation"]


 setup(
@@ -670,6 +670,7 @@ if is_tf_available():
        TFBertModel,
        TFBertPreTrainedModel,
    )
+    from .modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration
    from .modeling_tf_camembert import (
        TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFCamembertForMaskedLM,
@@ -750,6 +751,8 @@ if is_tf_available():
        TFLxmertPreTrainedModel,
        TFLxmertVisualFeatureEncoder,
    )
+    from .modeling_tf_marian import TFMarianMTModel
+    from .modeling_tf_mbart import TFMBartForConditionalGeneration
    from .modeling_tf_mobilebert import (
        TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFMobileBertForMaskedLM,
@@ -771,6 +774,7 @@ if is_tf_available():
        TFOpenAIGPTModel,
        TFOpenAIGPTPreTrainedModel,
    )
+    from .modeling_tf_pegasus import TFPegasusForConditionalGeneration
    from .modeling_tf_roberta import (
        TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFRobertaForMaskedLM,
@@ -2,6 +2,7 @@ import math

 import torch
 import torch.nn.functional as F
+from packaging import version

 from .utils import logging

@@ -9,29 +10,25 @@ from .utils import logging
 logger = logging.get_logger(__name__)


-def swish(x):
-    return x * torch.sigmoid(x)
-
-
 def _gelu_python(x):
    """
-    Original Implementation of the gelu activation function in Google Bert repo when initially created. For
-    information: OpenAI GPT's gelu is slightly different (and gives slightly different results): 0.5 * x * (1 +
+    Original Implementation of the GELU activation function in Google BERT repo when initially created. For
+    information: OpenAI GPT's GELU is slightly different (and gives slightly different results): 0.5 * x * (1 +
    torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) This is now written in C in
-    torch.nn.functional Also see https://arxiv.org/abs/1606.08415
+    torch.nn.functional Also see the Gaussian Error Linear Units paper: https://arxiv.org/abs/1606.08415
    """
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))


 def gelu_new(x):
    """
-    Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT). Also see
-    https://arxiv.org/abs/1606.08415
+    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT). Also see
+    the Gaussian Error Linear Units paper: https://arxiv.org/abs/1606.08415
    """
    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))


-if torch.__version__ < "1.4.0":
+if version.parse(torch.__version__) < version.parse("1.4"):
    gelu = _gelu_python
 else:
    gelu = F.gelu
@@ -41,6 +38,23 @@ def gelu_fast(x):
    return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))


+def _silu_python(x):
+    """
+    See Gaussian Error Linear Units (Hendrycks et al., https://arxiv.org/abs/1606.08415) where the SiLU (Sigmoid Linear
+    Unit) was originally introduced and coined, and see Sigmoid-Weighted Linear Units for Neural Network Function
+    Approximation in Reinforcement Learning (Elfwing et al., https://arxiv.org/abs/1702.03118) and Swish: a Self-Gated
+    Activation Function (Ramachandran et al., https://arxiv.org/abs/1710.05941v1) where the SiLU was experimented with
+    later.
+    """
+    return x * torch.sigmoid(x)
+
+
+if version.parse(torch.__version__) < version.parse("1.7"):
+    silu = _silu_python
+else:
+    silu = F.silu
+
+
 def mish(x):
    return x * torch.tanh(torch.nn.functional.softplus(x))

@@ -51,7 +65,8 @@ def linear_act(x):

 ACT2FN = {
    "relu": F.relu,
-    "swish": swish,
+    "silu": silu,
+    "swish": silu,
    "gelu": gelu,
    "tanh": torch.tanh,
    "gelu_new": gelu_new,
@@ -52,6 +52,7 @@ ACT2FN = {
    "gelu": tf.keras.layers.Activation(gelu),
    "relu": tf.keras.activations.relu,
    "swish": tf.keras.activations.swish,
+    "silu": tf.keras.activations.swish,
    "gelu_new": tf.keras.layers.Activation(gelu_new),
    "mish": tf.keras.layers.Activation(mish),
    "tanh": tf.keras.activations.tanh,
@@ -0,0 +1,192 @@
+import json
+import os
+import shutil
+from argparse import ArgumentParser, Namespace
+from pathlib import Path
+from typing import List
+
+from cookiecutter.main import cookiecutter
+from transformers.commands import BaseTransformersCLICommand
+
+from ..utils import logging
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def add_new_model_command_factory(args: Namespace):
+    return AddNewModelCommand(args.testing, args.testing_file, path=args.path)
+
+
+class AddNewModelCommand(BaseTransformersCLICommand):
+    @staticmethod
+    def register_subcommand(parser: ArgumentParser):
+        add_new_model_parser = parser.add_parser("add-new-model")
+        add_new_model_parser.add_argument("--testing", action="store_true", help="If in testing mode.")
+        add_new_model_parser.add_argument("--testing_file", type=str, help="Configuration file on which to run.")
+        add_new_model_parser.add_argument(
+            "--path", type=str, help="Path to cookiecutter. Should only be used for testing purposes."
+        )
+        add_new_model_parser.set_defaults(func=add_new_model_command_factory)
+
+    def __init__(self, testing: bool, testing_file: str, path=None, *args):
+        self._testing = testing
+        self._testing_file = testing_file
+        self._path = path
+
+    def run(self):
+        # Ensure that there is no other `cookiecutter-template-xxx` directory in the current working directory
+        directories = [directory for directory in os.listdir() if "cookiecutter-template-" in directory[:22]]
+        if len(directories) > 0:
+            raise ValueError(
+                "Several directories starting with `cookiecutter-template-` in current working directory. "
+                "Please clean your directory by removing all folders startign with `cookiecutter-template-` or "
+                "change your working directory."
+            )
+
+        path_to_transformer_root = (
+            Path(__file__).parent.parent.parent.parent if self._path is None else Path(self._path).parent.parent
+        )
+        path_to_cookiecutter = path_to_transformer_root / "templates" / "cookiecutter"
+
+        # Execute cookiecutter
+        if not self._testing:
+            cookiecutter(str(path_to_cookiecutter))
+        else:
+            with open(self._testing_file, "r") as configuration_file:
+                testing_configuration = json.load(configuration_file)
+
+            cookiecutter(
+                str(path_to_cookiecutter if self._path is None else self._path),
+                no_input=True,
+                extra_context=testing_configuration,
+            )
+
+        directory = [directory for directory in os.listdir() if "cookiecutter-template-" in directory[:22]][0]
+
+        # Retrieve configuration
+        with open(directory + "/configuration.json", "r") as configuration_file:
+            configuration = json.load(configuration_file)
+
+        lowercase_model_name = configuration["lowercase_modelname"]
+        pytorch_or_tensorflow = configuration["generate_tensorflow_and_pytorch"]
+        os.remove(f"{directory}/configuration.json")
+
+        output_pytorch = "PyTorch" in pytorch_or_tensorflow
+        output_tensorflow = "TensorFlow" in pytorch_or_tensorflow
+
+        shutil.move(
+            f"{directory}/configuration_{lowercase_model_name}.py",
+            f"{path_to_transformer_root}/src/transformers/configuration_{lowercase_model_name}.py",
+        )
+
+        def remove_copy_lines(path):
+            with open(path, "r") as f:
+                lines = f.readlines()
+            with open(path, "w") as f:
+                for line in lines:
+                    if "# Copied from transformers." not in line:
+                        f.write(line)
+
+        if output_pytorch:
+            if not self._testing:
+                remove_copy_lines(f"{directory}/modeling_{lowercase_model_name}.py")
+
+            shutil.move(
+                f"{directory}/modeling_{lowercase_model_name}.py",
+                f"{path_to_transformer_root}/src/transformers/modeling_{lowercase_model_name}.py",
+            )
+
+            shutil.move(
+                f"{directory}/test_modeling_{lowercase_model_name}.py",
+                f"{path_to_transformer_root}/tests/test_modeling_{lowercase_model_name}.py",
+            )
+        else:
+            os.remove(f"{directory}/modeling_{lowercase_model_name}.py")
+            os.remove(f"{directory}/test_modeling_{lowercase_model_name}.py")
+
+        if output_tensorflow:
+            if not self._testing:
+                remove_copy_lines(f"{directory}/modeling_tf_{lowercase_model_name}.py")
+
+            shutil.move(
+                f"{directory}/modeling_tf_{lowercase_model_name}.py",
+                f"{path_to_transformer_root}/src/transformers/modeling_tf_{lowercase_model_name}.py",
+            )
+
+            shutil.move(
+                f"{directory}/test_modeling_tf_{lowercase_model_name}.py",
+                f"{path_to_transformer_root}/tests/test_modeling_tf_{lowercase_model_name}.py",
+            )
+        else:
+            os.remove(f"{directory}/modeling_tf_{lowercase_model_name}.py")
+            os.remove(f"{directory}/test_modeling_tf_{lowercase_model_name}.py")
+
+        shutil.move(
+            f"{directory}/{lowercase_model_name}.rst",
+            f"{path_to_transformer_root}/docs/source/model_doc/{lowercase_model_name}.rst",
+        )
+
+        shutil.move(
+            f"{directory}/tokenization_{lowercase_model_name}.py",
+            f"{path_to_transformer_root}/src/transformers/tokenization_{lowercase_model_name}.py",
+        )
+
+        from os import fdopen, remove
+        from shutil import copymode, move
+        from tempfile import mkstemp
+
+        def replace(original_file: str, line_to_copy_below: str, lines_to_copy: List[str]):
+            # Create temp file
+            fh, abs_path = mkstemp()
+            line_found = False
+            with fdopen(fh, "w") as new_file:
+                with open(original_file) as old_file:
+                    for line in old_file:
+                        new_file.write(line)
+                        if line_to_copy_below in line:
+                            line_found = True
+                            for line_to_copy in lines_to_copy:
+                                new_file.write(line_to_copy)
+
+            if not line_found:
+                raise ValueError(f"Line {line_to_copy_below} was not found in file.")
+
+            # Copy the file permissions from the old file to the new file
+            copymode(original_file, abs_path)
+            # Remove original file
+            remove(original_file)
+            # Move new file
+            move(abs_path, original_file)
+
+        def skip_units(line):
+            return ("generating PyTorch" in line and not output_pytorch) or (
+                "generating TensorFlow" in line and not output_tensorflow
+            )
+
+        def replace_in_files(path_to_datafile):
+            with open(path_to_datafile) as datafile:
+                lines_to_copy = []
+                skip_file = False
+                skip_snippet = False
+                for line in datafile:
+                    if "# To replace in: " in line and "##" not in line:
+                        file_to_replace_in = line.split('"')[1]
+                        skip_file = skip_units(line)
+                    elif "# Below: " in line and "##" not in line:
+                        line_to_copy_below = line.split('"')[1]
+                        skip_snippet = skip_units(line)
+                    elif "# End." in line and "##" not in line:
+                        if not skip_file and not skip_snippet:
+                            replace(file_to_replace_in, line_to_copy_below, lines_to_copy)
+
+                        lines_to_copy = []
+                    elif "# Replace with" in line and "##" not in line:
+                        lines_to_copy = []
+                    elif "##" not in line:
+                        lines_to_copy.append(line)
+
+            remove(path_to_datafile)
+
+        replace_in_files(f"{directory}/to_replace_{lowercase_model_name}.py")
+        os.rmdir(directory)
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 from argparse import ArgumentParser

+from transformers.commands.add_new_model import AddNewModelCommand
 from transformers.commands.convert import ConvertCommand
 from transformers.commands.download import DownloadCommand
 from transformers.commands.env import EnvironmentCommand
@@ -20,6 +21,7 @@ def main():
    RunCommand.register_subcommand(commands_parser)
    ServeCommand.register_subcommand(commands_parser)
    UserCommands.register_subcommand(commands_parser)
+    AddNewModelCommand.register_subcommand(commands_parser)

    # Let's go
    args = parser.parse_args()
@@ -61,7 +61,7 @@ class AlbertConfig(PretrainedConfig):
            The number of inner repetition of attention and ffn.
        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu_new"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0):
@@ -59,6 +59,7 @@ from .configuration_xlnet import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLNetConfi
 ALL_PRETRAINED_CONFIG_ARCHIVE_MAP = dict(
    (key, value)
    for pretrained_map in [
+        # Add archive maps here
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
        BART_PRETRAINED_CONFIG_ARCHIVE_MAP,
        BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP,
@@ -95,6 +96,7 @@ ALL_PRETRAINED_CONFIG_ARCHIVE_MAP = dict(

 CONFIG_MAPPING = OrderedDict(
    [
+        # Add configs here
        ("retribert", RetriBertConfig),
        ("t5", T5Config),
        ("mobilebert", MobileBertConfig),
@@ -136,6 +138,7 @@ CONFIG_MAPPING = OrderedDict(

 MODEL_NAMES_MAPPING = OrderedDict(
    [
+        # Add cased model names here
        ("retribert", "RetriBERT"),
        ("t5", "T5"),
        ("mobilebert", "MobileBERT"),
@@ -59,7 +59,7 @@ class BartConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        dropout (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
@@ -20,6 +20,7 @@ from .utils import logging


 logger = logging.get_logger(__name__)
+logger.setLevel(logging.INFO)

 BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
    "bert-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json",
@@ -74,7 +75,7 @@ class BertConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -40,7 +40,7 @@ class BertGenerationConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often called feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -56,7 +56,7 @@ class BlenderbotConfig(BartConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        dropout (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
@@ -52,7 +52,7 @@ class DebertaConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"`, :obj:`"gelu"`, :obj:`"tanh"`, :obj:`"gelu_fast"`,
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"`, :obj:`"gelu"`, :obj:`"tanh"`, :obj:`"gelu_fast"`,
            :obj:`"mish"`, :obj:`"linear"`, :obj:`"sigmoid"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
@@ -66,7 +66,7 @@ class DistilBertConfig(PretrainedConfig):
            The dropout ratio for the attention probabilities.
        activation (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        initializer_range (:obj:`float`, `optional`, defaults to 0.02):
            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
        qa_dropout (:obj:`float`, `optional`, defaults to 0.1):
@@ -55,7 +55,7 @@ class DPRConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -60,7 +60,7 @@ class ElectraConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -68,7 +68,7 @@ class EncoderDecoderConfig(PretrainedConfig):
        >>> encoder_decoder_config = EncoderDecoderConfig.from_pretrained('my-model')
        >>> model = EncoderDecoderModel.from_pretrained('my-model', config=encoder_decoder_config)
    """
-    model_type = "encoder_decoder"
+    model_type = "encoder-decoder"
    is_composition = True

    def __init__(self, **kwargs):
@@ -71,7 +71,7 @@ class FSMTConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
        activation_function (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"relu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        dropout (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
@@ -66,7 +66,7 @@ class FunnelConfig(PretrainedConfig):
            Inner dimension in the feed-forward blocks.
        hidden_act (:obj:`str` or :obj:`callable`, `optional`, defaults to :obj:`"gelu_new"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_dropout (:obj:`float`, `optional`, defaults to 0.1):
@@ -60,7 +60,7 @@ class GPT2Config(PretrainedConfig):
        n_inner (:obj:`int`, `optional`, defaults to None):
            Dimensionality of the inner feed-forward layers. :obj:`None` will set it to 4 times n_embd
        activation_function (:obj:`str`, `optional`, defaults to :obj:`"gelu"`):
-            Activation function, to be selected in the list :obj:`["relu", "swish", "gelu", "tanh", "gelu_new"]`.
+            Activation function, to be selected in the list :obj:`["relu", "silu", "gelu", "tanh", "gelu_new"]`.
        resid_pdrop (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        embd_pdrop (:obj:`int`, `optional`, defaults to 0.1):
@@ -52,7 +52,7 @@ class LayoutLMConfig(BertConfig):
            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -55,7 +55,7 @@ class LxmertConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -50,7 +50,7 @@ class MarianConfig(BartConfig):
            Dimensionality of the "intermediate" (i.e., feed-forward) layer in decoder.
        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        dropout (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
@@ -55,7 +55,7 @@ class MBartConfig(BartConfig):
            Dimensionality of the "intermediate" (i.e., feed-forward) layer in decoder.
        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        dropout (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
@@ -48,7 +48,7 @@ class MobileBertConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"relu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.0):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -54,7 +54,7 @@ class OpenAIGPTConfig(PretrainedConfig):
            Number of attention heads for each attention layer in the Transformer encoder.
        afn (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        resid_pdrop (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        embd_pdrop (:obj:`int`, `optional`, defaults to 0.1):
@@ -94,7 +94,7 @@ class PegasusConfig(BartConfig):
            Dimensionality of the "intermediate" (i.e., feed-forward) layer in decoder.
        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        dropout (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
@@ -39,7 +39,7 @@ class ProphetNetConfig(PretrainedConfig):
            The dropout ratio for activations inside the fully connected layer.
        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        vocab_size (:obj:`int`, `optional`, defaults to 30522):
            Vocabulary size of the ProphetNET model. Defines the number of different tokens that can be represented by
            the :obj:`inputs_ids` passed when calling :class:`~transformers.ProphetNetModel`.
@@ -80,7 +80,7 @@ class ReformerConfig(PretrainedConfig):
            :obj:`None` to ensure fully random rotations in local sensitive hashing scheme.
        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"relu"`):
            The non-linear activation function (function or string) in the feed forward layer in the residual attention
-            block. If string, :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            block. If string, :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.05):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        hidden_size (:obj:`int`, `optional`, defaults to 256):
@@ -49,7 +49,7 @@ class RetriBertConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -50,7 +50,7 @@ class SqueezeBertConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string,
-            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
@@ -54,7 +54,7 @@ class XLNetConfig(PretrainedConfig):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        ff_activation (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
            The non-linear activation function (function or string) in the If string, :obj:`"gelu"`, :obj:`"relu"`,
-            :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+            :obj:`"silu"` and :obj:`"gelu_new"` are supported.
        untie_r (:obj:`bool`, `optional`, defaults to :obj:`True`):
            Whether or not to untie relative position biases
        attn_type (:obj:`str`, `optional`, defaults to :obj:`"bi"`):
@@ -29,7 +29,8 @@ try:
    wandb.ensure_configured()
    if wandb.api.api_key is None:
        _has_wandb = False
-        wandb.termwarn("W&B installed but not logged in.  Run `wandb login` or set the WANDB_API_KEY env variable.")
+        if os.getenv("WANDB_DISABLED"):
+            logger.warning("W&B installed but not logged in. Run `wandb login` or set the WANDB_API_KEY env variable.")
    else:
        _has_wandb = False if os.getenv("WANDB_DISABLED") else True
 except (ImportError, AttributeError):
@@ -281,7 +282,9 @@ class TensorBoardCallback(TrainerCallback):
                if hasattr(model, "config") and model.config is not None:
                    model_config_json = model.config.to_json_string()
                    self.tb_writer.add_text("model_config", model_config_json)
-            self.tb_writer.add_hparams(args.to_sanitized_dict(), metric_dict={})
+            # Version of TensorBoard coming from tensorboardX does not have this method.
+            if hasattr(self.tb_writer, "add_hparams"):
+                self.tb_writer.add_hparams(args.to_sanitized_dict(), metric_dict={})

    def on_log(self, args, state, control, logs=None, **kwargs):
        if state.is_world_process_zero:
@@ -478,7 +478,7 @@ class AlbertPreTrainedModel(PreTrainedModel):
@dataclass
 class AlbertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.AlbertForPreTrainingModel`.
+    Output type of :class:`~transformers.AlbertForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`):
@@ -222,11 +222,14 @@ from .modeling_xlnet import (
 from .utils import logging


+# Add modeling imports here
+
 logger = logging.get_logger(__name__)


 MODEL_MAPPING = OrderedDict(
    [
+        # Base model mapping
        (RetriBertConfig, RetriBertModel),
        (T5Config, T5Model),
        (DistilBertConfig, DistilBertModel),
@@ -262,6 +265,7 @@ MODEL_MAPPING = OrderedDict(

 MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(
    [
+        # Model for pre-training mapping
        (LayoutLMConfig, LayoutLMForMaskedLM),
        (RetriBertConfig, RetriBertModel),
        (T5Config, T5ForConditionalGeneration),
@@ -290,6 +294,7 @@ MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(

 MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(
    [
+        # Model with LM heads mapping
        (LayoutLMConfig, LayoutLMForMaskedLM),
        (T5Config, T5ForConditionalGeneration),
        (DistilBertConfig, DistilBertForMaskedLM),
@@ -320,6 +325,7 @@ MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(

 MODEL_FOR_CAUSAL_LM_MAPPING = OrderedDict(
    [
+        # Model for Causal LM mapping
        (CamembertConfig, CamembertForCausalLM),
        (XLMRobertaConfig, XLMRobertaForCausalLM),
        (RobertaConfig, RobertaForCausalLM),
@@ -342,6 +348,7 @@ MODEL_FOR_CAUSAL_LM_MAPPING = OrderedDict(

 MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(
    [
+        # Model for Masked LM mapping
        (LayoutLMConfig, LayoutLMForMaskedLM),
        (DistilBertConfig, DistilBertForMaskedLM),
        (AlbertConfig, AlbertForMaskedLM),
@@ -363,6 +370,7 @@ MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(

 MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = OrderedDict(
    [
+        # Model for Seq2Seq Causal LM mapping
        (T5Config, T5ForConditionalGeneration),
        (PegasusConfig, PegasusForConditionalGeneration),
        (MarianConfig, MarianMTModel),
@@ -378,6 +386,7 @@ MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = OrderedDict(

 MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(
    [
+        # Model for Sequence Classification mapping
        (DistilBertConfig, DistilBertForSequenceClassification),
        (AlbertConfig, AlbertForSequenceClassification),
        (CamembertConfig, CamembertForSequenceClassification),
@@ -401,6 +410,7 @@ MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(

 MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
    [
+        # Model for Question Answering mapping
        (DistilBertConfig, DistilBertForQuestionAnswering),
        (AlbertConfig, AlbertForQuestionAnswering),
        (CamembertConfig, CamembertForQuestionAnswering),
@@ -422,6 +432,7 @@ MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(

 MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict(
    [
+        # Model for Token Classification mapping
        (LayoutLMConfig, LayoutLMForTokenClassification),
        (DistilBertConfig, DistilBertForTokenClassification),
        (CamembertConfig, CamembertForTokenClassification),
@@ -443,6 +454,7 @@ MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict(

 MODEL_FOR_MULTIPLE_CHOICE_MAPPING = OrderedDict(
    [
+        # Model for Multiple Choice mapping
        (CamembertConfig, CamembertForMultipleChoice),
        (ElectraConfig, ElectraForMultipleChoice),
        (XLMRobertaConfig, XLMRobertaForMultipleChoice),
@@ -494,8 +506,8 @@ AUTO_MODEL_PRETRAINED_DOCSTRING = r"""
                    - The model is a model provided by the library (loaded with the `shortcut name` string of a
                      pretrained model).
                    - The model was saved using :meth:`~transformers.PreTrainedModel.save_pretrained` and is reloaded
-                      by suppling the save directory.
-                    - The model is loaded by suppling a local directory as ``pretrained_model_name_or_path`` and a
+                      by supplying the save directory.
+                    - The model is loaded by supplying a local directory as ``pretrained_model_name_or_path`` and a
                      configuration JSON file named `config.json` is found in the directory.
            state_dict (`Dict[str, torch.Tensor]`, `optional`):
                A state dictionary to use instead of a state dictionary loaded from saved weights file.
@@ -427,7 +427,6 @@ class DecoderLayer(nn.Module):
        output_attentions=False,
    ):
        residual = x
-
        if layer_state is None:
            layer_state = {}
        if self.normalize_before:
@@ -447,7 +446,7 @@ class DecoderLayer(nn.Module):
        if not self.normalize_before:
            x = self.self_attn_layer_norm(x)

-        # Cross attention
+        # Cross-Attention Block
        residual = x
        assert self.encoder_attn.cache_key != self.self_attn.cache_key
        if self.normalize_before:
@@ -628,7 +627,6 @@ class BartDecoder(nn.Module):
        encoder_hidden_states = encoder_hidden_states.transpose(0, 1)

        next_cache = next_decoder_cache if use_cache else None
-
        if not return_dict:
            return tuple(v for v in [x, next_cache, all_hidden_states, all_self_attns] if v is not None)
        return BaseModelOutputWithPast(
@@ -606,7 +606,7 @@ class BertPreTrainedModel(PreTrainedModel):
@dataclass
 class BertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.BertForPreTrainingModel`.
+    Output type of :class:`~transformers.BertForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`):
@@ -555,7 +555,7 @@ class ElectraPreTrainedModel(PreTrainedModel):
@dataclass
 class ElectraForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.ElectraForPreTrainingModel`.
+    Output type of :class:`~transformers.ElectraForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`):
@@ -20,7 +20,6 @@ import numpy as np
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from flax.linen import compact

 from .configuration_bert import BertConfig
 from .file_utils import add_start_docstrings
@@ -108,13 +107,15 @@ class FlaxBertLayerNorm(nn.Module):
    """

    epsilon: float = 1e-6
-    dtype: jnp.dtype = jnp.float32
-    bias: bool = True
-    scale: bool = True
+    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    bias: bool = True  # If True, bias (beta) is added.
+    scale: bool = True  # If True, multiply by scale (gamma). When the next layer is linear
+    # (also e.g. nn.relu), this can be disabled since the scaling will be
+    # done by the next layer.
    bias_init: jnp.ndarray = nn.initializers.zeros
    scale_init: jnp.ndarray = nn.initializers.ones

-    @compact
+    @nn.compact
    def __call__(self, x):
        """
        Applies layer normalization on the input. It normalizes the activations of the layer for each given example in
@@ -123,13 +124,6 @@ class FlaxBertLayerNorm(nn.Module):

        Args:
          x: the inputs
-          epsilon: A small float added to variance to avoid dividing by zero.
-          dtype: the dtype of the computation (default: float32).
-          bias:  If True, bias (beta) is added.
-          scale: If True, multiply by scale (gamma). When the next layer is linear
-            (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.
-          bias_init: Initializer for bias, by default, zero.
-          scale_init: Initializer for scale, by default, one

        Returns:
          Normalized inputs (the same shape as inputs).
@@ -157,7 +151,7 @@ class FlaxBertEmbedding(nn.Module):
    hidden_size: int
    emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1)

-    @compact
+    @nn.compact
    def __call__(self, inputs):
        embedding = self.param("weight", self.emb_init, (self.vocab_size, self.hidden_size))
        return jnp.take(embedding, inputs, axis=0)
@@ -171,7 +165,7 @@ class FlaxBertEmbeddings(nn.Module):
    type_vocab_size: int
    max_length: int

-    @compact
+    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):

        # Embed
@@ -198,7 +192,7 @@ class FlaxBertAttention(nn.Module):
    num_heads: int
    head_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        self_att = nn.attention.SelfAttention(num_heads=self.num_heads, qkv_features=self.head_size, name="self")(
            hidden_state, attention_mask
@@ -211,7 +205,7 @@ class FlaxBertAttention(nn.Module):
 class FlaxBertIntermediate(nn.Module):
    output_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state):
        # TODO: Add ACT2FN reference to change activation function
        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)
@@ -219,7 +213,7 @@ class FlaxBertIntermediate(nn.Module):


 class FlaxBertOutput(nn.Module):
-    @compact
+    @nn.compact
    def __call__(self, intermediate_output, attention_output):
        hidden_state = nn.Dense(attention_output.shape[-1], name="dense")(intermediate_output)
        hidden_state = FlaxBertLayerNorm(name="layer_norm")(hidden_state + attention_output)
@@ -231,7 +225,7 @@ class FlaxBertLayer(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        attention = FlaxBertAttention(self.num_heads, self.head_size, name="attention")(hidden_state, attention_mask)
        intermediate = FlaxBertIntermediate(self.intermediate_size, name="intermediate")(attention)
@@ -250,7 +244,7 @@ class FlaxBertLayerCollection(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, inputs, attention_mask):
        assert self.num_layers > 0, f"num_layers should be >= 1, got ({self.num_layers})"

@@ -270,7 +264,7 @@ class FlaxBertEncoder(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        layer = FlaxBertLayerCollection(
            self.num_layers, self.num_heads, self.head_size, self.intermediate_size, name="layer"
@@ -279,7 +273,7 @@ class FlaxBertEncoder(nn.Module):


 class FlaxBertPooler(nn.Module):
-    @compact
+    @nn.compact
    def __call__(self, hidden_state):
        cls_token = hidden_state[:, 0]
        out = nn.Dense(hidden_state.shape[-1], name="dense")(cls_token)
@@ -296,7 +290,7 @@ class FlaxBertModule(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):

        # Embedding
@@ -19,7 +19,6 @@ import numpy as np
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from flax.linen import compact

 from .configuration_roberta import RobertaConfig
 from .file_utils import add_start_docstrings
@@ -108,13 +107,15 @@ class FlaxRobertaLayerNorm(nn.Module):
    """

    epsilon: float = 1e-6
-    dtype: jnp.dtype = jnp.float32
-    bias: bool = True
-    scale: bool = True
+    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    bias: bool = True  # If True, bias (beta) is added.
+    scale: bool = True  # If True, multiply by scale (gamma). When the next layer is linear
+    # (also e.g. nn.relu), this can be disabled since the scaling will be
+    # done by the next layer.
    bias_init: jnp.ndarray = nn.initializers.zeros
    scale_init: jnp.ndarray = nn.initializers.ones

-    @compact
+    @nn.compact
    def __call__(self, x):
        """
        Applies layer normalization on the input. It normalizes the activations of the layer for each given example in
@@ -123,13 +124,6 @@ class FlaxRobertaLayerNorm(nn.Module):

        Args:
          x: the inputs
-          epsilon: A small float added to variance to avoid dividing by zero.
-          dtype: the dtype of the computation (default: float32).
-          bias:  If True, bias (beta) is added.
-          scale: If True, multiply by scale (gamma). When the next layer is linear
-            (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.
-          bias_init: Initializer for bias, by default, zero.
-          scale_init: Initializer for scale, by default, one

        Returns:
          Normalized inputs (the same shape as inputs).
@@ -158,7 +152,7 @@ class FlaxRobertaEmbedding(nn.Module):
    hidden_size: int
    emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1)

-    @compact
+    @nn.compact
    def __call__(self, inputs):
        embedding = self.param("weight", self.emb_init, (self.vocab_size, self.hidden_size))
        return jnp.take(embedding, inputs, axis=0)
@@ -173,7 +167,7 @@ class FlaxRobertaEmbeddings(nn.Module):
    type_vocab_size: int
    max_length: int

-    @compact
+    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):

        # Embed
@@ -201,7 +195,7 @@ class FlaxRobertaAttention(nn.Module):
    num_heads: int
    head_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        self_att = nn.attention.SelfAttention(num_heads=self.num_heads, qkv_features=self.head_size, name="self")(
            hidden_state, attention_mask
@@ -215,7 +209,7 @@ class FlaxRobertaAttention(nn.Module):
 class FlaxRobertaIntermediate(nn.Module):
    output_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state):
        # TODO: Add ACT2FN reference to change activation function
        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)
@@ -224,7 +218,7 @@ class FlaxRobertaIntermediate(nn.Module):

 # Copied from transformers.modeling_flax_bert.FlaxBertOutput with Bert->Roberta
 class FlaxRobertaOutput(nn.Module):
-    @compact
+    @nn.compact
    def __call__(self, intermediate_output, attention_output):
        hidden_state = nn.Dense(attention_output.shape[-1], name="dense")(intermediate_output)
        hidden_state = FlaxRobertaLayerNorm(name="layer_norm")(hidden_state + attention_output)
@@ -236,7 +230,7 @@ class FlaxRobertaLayer(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        attention = FlaxRobertaAttention(self.num_heads, self.head_size, name="attention")(
            hidden_state, attention_mask
@@ -258,7 +252,7 @@ class FlaxRobertaLayerCollection(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, inputs, attention_mask):
        assert self.num_layers > 0, f"num_layers should be >= 1, got ({self.num_layers})"

@@ -279,7 +273,7 @@ class FlaxRobertaEncoder(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, hidden_state, attention_mask):
        layer = FlaxRobertaLayerCollection(
            self.num_layers, self.num_heads, self.head_size, self.intermediate_size, name="layer"
@@ -289,7 +283,7 @@ class FlaxRobertaEncoder(nn.Module):

 # Copied from transformers.modeling_flax_bert.FlaxBertPooler with Bert->Roberta
 class FlaxRobertaPooler(nn.Module):
-    @compact
+    @nn.compact
    def __call__(self, hidden_state):
        cls_token = hidden_state[:, 0]
        out = nn.Dense(hidden_state.shape[-1], name="dense")(cls_token)
@@ -307,7 +301,7 @@ class FlaxRobertaModule(nn.Module):
    head_size: int
    intermediate_size: int

-    @compact
+    @nn.compact
    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):

        # Embedding
@@ -53,6 +53,7 @@ def gelu(x):
 ACT2FN = {
    "gelu": nn.gelu,
    "relu": nn.relu,
+    "silu": nn.swish,
    "swish": nn.swish,
    "gelu_new": gelu,
 }
@@ -798,7 +798,7 @@ class FunnelClassificationHead(nn.Module):
@dataclass
 class FunnelForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.FunnelForPreTrainingModel`.
+    Output type of :class:`~transformers.FunnelForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`):
@@ -144,7 +144,7 @@ class LxmertForQuestionAnsweringOutput(ModelOutput):
@dataclass
 class LxmertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.LxmertForPreTrainingModel`.
+    Output type of :class:`~transformers.LxmertForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`):
@@ -695,7 +695,7 @@ class MobileBertPreTrainedModel(PreTrainedModel):
@dataclass
 class MobileBertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.MobileBertForPreTrainingModel`.
+    Output type of :class:`~transformers.MobileBertForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`):
@@ -27,7 +27,7 @@ import torch
 import torch.nn as nn
 from torch.nn import CrossEntropyLoss, MSELoss

-from .activations import gelu_new, swish
+from .activations import gelu_new, silu
 from .configuration_openai import OpenAIGPTConfig
 from .file_utils import (
    ModelOutput,
@@ -139,7 +139,7 @@ def load_tf_weights_in_openai_gpt(model, config, openai_checkpoint_folder_path):
    return model


-ACT_FNS = {"relu": nn.ReLU, "swish": swish, "gelu": gelu_new}
+ACT_FNS = {"relu": nn.ReLU, "silu": silu, "gelu": gelu_new, "swish": silu}


 class Attention(nn.Module):
@@ -260,7 +260,7 @@ class MaskedLMOutput(ModelOutput):

    Args:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
-            Masked languaged modeling (MLM) loss.
+            Masked language modeling (MLM) loss.
        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
@@ -289,7 +289,7 @@ class Seq2SeqLMOutput(ModelOutput):

    Args:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
-            Languaged modeling loss.
+            Language modeling loss.
        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        past_key_values (:obj:`List[torch.FloatTensor]`, `optional`, returned when ``use_cache=True`` is passed or when ``config.use_cache=True``):
@@ -225,7 +225,7 @@ class ProphetNetSeq2SeqLMOutput(ModelOutput):

    Args:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
-            Languaged modeling loss.
+            Language modeling loss.
        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, decoder_sequence_length, config.vocab_size)`):
            Prediction scores of the main stream language modeling head (scores for each vocabulary token before
            SoftMax).
@@ -438,7 +438,7 @@ class ProphetNetDecoderLMOutput(ModelOutput):

    Args:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
-            Languaged modeling loss.
+            Language modeling loss.
        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, decoder_sequence_length, config.vocab_size)`):
            Prediction scores of the main stream language modeling head (scores for each vocabulary token before
            SoftMax).
@@ -1124,7 +1124,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
            >>> model = T5ForConditionalGeneration.from_pretrained('t5-small', return_dict=True)

            >>> input_ids = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='pt').input_ids
-            labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2> </s>', return_tensors='pt').input_ids
+            >>> labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2> </s>', return_tensors='pt').input_ids
            >>> outputs = model(input_ids=input_ids, labels=labels)
            >>> loss = outputs.loss
            >>> logits = outputs.logits
@@ -79,25 +79,31 @@ class TFAlbertEmbeddings(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)

-        self.config = config
        self.vocab_size = config.vocab_size
+        self.embedding_size = config.embedding_size
+        self.initializer_range = config.initializer_range
+        self.max_position_embeddings = config.max_position_embeddings
+        self.type_vocab_size = config.type_vocab_size
+        self.layer_norm_eps = config.layer_norm_eps
+        self.hidden_dropout_prob = config.hidden_dropout_prob
+
        self.position_embeddings = tf.keras.layers.Embedding(
-            config.max_position_embeddings,
-            config.embedding_size,
-            embeddings_initializer=get_initializer(self.config.initializer_range),
+            self.max_position_embeddings,
+            self.embedding_size,
+            embeddings_initializer=get_initializer(self.initializer_range),
            name="position_embeddings",
        )
        self.token_type_embeddings = tf.keras.layers.Embedding(
-            config.type_vocab_size,
-            config.embedding_size,
-            embeddings_initializer=get_initializer(self.config.initializer_range),
+            self.type_vocab_size,
+            self.embedding_size,
+            embeddings_initializer=get_initializer(self.initializer_range),
            name="token_type_embeddings",
        )

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
-        self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
-        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
+        self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=self.layer_norm_eps, name="LayerNorm")
+        self.dropout = tf.keras.layers.Dropout(self.hidden_dropout_prob)

    def build(self, input_shape):
        """Build shared word embedding layer """
@@ -106,8 +112,8 @@ class TFAlbertEmbeddings(tf.keras.layers.Layer):
            # arbitrarily, and works well.
            self.word_embeddings = self.add_weight(
                "weight",
-                shape=[self.config.vocab_size, self.config.embedding_size],
-                initializer=get_initializer(self.config.initializer_range),
+                shape=[self.vocab_size, self.embedding_size],
+                initializer=get_initializer(self.initializer_range),
            )
        super().build(input_shape)

@@ -182,9 +188,9 @@ class TFAlbertEmbeddings(tf.keras.layers.Layer):
        """
        batch_size = shape_list(inputs)[0]
        length = shape_list(inputs)[1]
-        x = tf.reshape(inputs, [-1, self.config.embedding_size])
+        x = tf.reshape(inputs, [-1, self.embedding_size])
        logits = tf.matmul(x, self.word_embeddings, transpose_b=True)
-        return tf.reshape(logits, [batch_size, length, self.config.vocab_size])
+        return tf.reshape(logits, [batch_size, length, self.vocab_size])


 class TFAlbertSelfOutput(tf.keras.layers.Layer):
@@ -207,7 +213,7 @@ class TFAlbertAttention(tf.keras.layers.Layer):
    """ Contains the complete attention sublayer, including both dropouts and layer norm. """

    def __init__(self, config, **kwargs):
-        super().__init__(config, **kwargs)
+        super().__init__(**kwargs)

        self.hidden_size = config.hidden_size
        self.output_attentions = config.output_attentions
@@ -371,7 +377,8 @@ class TFAlbertTransformer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)

-        self.config = config
+        self.num_hidden_layers = config.num_hidden_layers
+        self.num_hidden_groups = config.num_hidden_groups
        self.embedding_hidden_mapping_in = tf.keras.layers.Dense(
            config.hidden_size,
            kernel_initializer=get_initializer(config.initializer_range),
@@ -396,12 +403,12 @@ class TFAlbertTransformer(tf.keras.layers.Layer):
        all_attentions = () if output_attentions else None
        all_hidden_states = (hidden_states,) if output_hidden_states else None

-        for i in range(self.config.num_hidden_layers):
+        for i in range(self.num_hidden_layers):
            # Number of layers in a hidden group
-            layers_per_group = int(self.config.num_hidden_layers / self.config.num_hidden_groups)
+            layers_per_group = int(self.num_hidden_layers / self.num_hidden_groups)

            # Index of the hidden group
-            group_idx = int(i / (self.config.num_hidden_layers / self.config.num_hidden_groups))
+            group_idx = int(i / (self.num_hidden_layers / self.num_hidden_groups))

            layer_group_output = self.albert_layer_groups[group_idx](
                hidden_states,
@@ -621,7 +628,7 @@ class TFAlbertMainLayer(tf.keras.layers.Layer):
@dataclass
 class TFAlbertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.TFAlbertForPreTrainingModel`.
+    Output type of :class:`~transformers.TFAlbertForPreTraining`.

    Args:
        prediction_logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
@@ -41,6 +41,10 @@ from .configuration_auto import (
    XLNetConfig,
    replace_list_option_in_docstrings,
 )
+from .configuration_blenderbot import BlenderbotConfig
+from .configuration_marian import MarianConfig
+from .configuration_mbart import MBartConfig
+from .configuration_pegasus import PegasusConfig
 from .configuration_utils import PretrainedConfig
 from .file_utils import add_start_docstrings
 from .modeling_tf_albert import (
@@ -63,6 +67,7 @@ from .modeling_tf_bert import (
    TFBertLMHeadModel,
    TFBertModel,
 )
+from .modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration
 from .modeling_tf_camembert import (
    TFCamembertForMaskedLM,
    TFCamembertForMultipleChoice,
@@ -108,6 +113,8 @@ from .modeling_tf_funnel import (
 )
 from .modeling_tf_gpt2 import TFGPT2LMHeadModel, TFGPT2Model
 from .modeling_tf_longformer import TFLongformerForMaskedLM, TFLongformerForQuestionAnswering, TFLongformerModel
+from .modeling_tf_marian import TFMarianMTModel
+from .modeling_tf_mbart import TFMBartForConditionalGeneration
 from .modeling_tf_mobilebert import (
    TFMobileBertForMaskedLM,
    TFMobileBertForMultipleChoice,
@@ -118,6 +125,7 @@ from .modeling_tf_mobilebert import (
    TFMobileBertModel,
 )
 from .modeling_tf_openai import TFOpenAIGPTLMHeadModel, TFOpenAIGPTModel
+from .modeling_tf_pegasus import TFPegasusForConditionalGeneration
 from .modeling_tf_roberta import (
    TFRobertaForMaskedLM,
    TFRobertaForMultipleChoice,
@@ -155,11 +163,14 @@ from .modeling_tf_xlnet import (
 from .utils import logging


+# Add modeling imports here
+
 logger = logging.get_logger(__name__)


 TF_MODEL_MAPPING = OrderedDict(
    [
+        # Base model mapping
        (T5Config, TFT5Model),
        (DistilBertConfig, TFDistilBertModel),
        (AlbertConfig, TFAlbertModel),
@@ -184,6 +195,7 @@ TF_MODEL_MAPPING = OrderedDict(

 TF_MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(
    [
+        # Model for pre-training mapping
        (T5Config, TFT5ForConditionalGeneration),
        (DistilBertConfig, TFDistilBertForMaskedLM),
        (AlbertConfig, TFAlbertForPreTraining),
@@ -207,9 +219,11 @@ TF_MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(

 TF_MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(
    [
+        # Model with LM heads mapping
        (T5Config, TFT5ForConditionalGeneration),
        (DistilBertConfig, TFDistilBertForMaskedLM),
        (AlbertConfig, TFAlbertForMaskedLM),
+        (MarianConfig, TFMarianMTModel),
        (BartConfig, TFBartForConditionalGeneration),
        (CamembertConfig, TFCamembertForMaskedLM),
        (XLMRobertaConfig, TFXLMRobertaForMaskedLM),
@@ -231,6 +245,7 @@ TF_MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(

 TF_MODEL_FOR_CAUSAL_LM_MAPPING = OrderedDict(
    [
+        # Model for Causal LM mapping
        (BertConfig, TFBertLMHeadModel),
        (OpenAIGPTConfig, TFOpenAIGPTLMHeadModel),
        (GPT2Config, TFGPT2LMHeadModel),
@@ -246,6 +261,7 @@ TF_MODEL_FOR_CAUSAL_LM_MAPPING = OrderedDict(

 TF_MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(
    [
+        # Model for Masked LM mapping
        (DistilBertConfig, TFDistilBertForMaskedLM),
        (AlbertConfig, TFAlbertForMaskedLM),
        (CamembertConfig, TFCamembertForMaskedLM),
@@ -261,12 +277,22 @@ TF_MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(
    ]
 )

+
 TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = OrderedDict(
-    [(T5Config, TFT5ForConditionalGeneration), (BartConfig, TFBartForConditionalGeneration)]
+    [
+        # Model for Seq2Seq Causal LM mapping
+        (T5Config, TFT5ForConditionalGeneration),
+        (MarianConfig, TFMarianMTModel),
+        (MBartConfig, TFMBartForConditionalGeneration),
+        (PegasusConfig, TFPegasusForConditionalGeneration),
+        (BlenderbotConfig, TFBlenderbotForConditionalGeneration),
+        (BartConfig, TFBartForConditionalGeneration),
+    ]
 )

 TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(
    [
+        # Model for Sequence Classification mapping
        (DistilBertConfig, TFDistilBertForSequenceClassification),
        (AlbertConfig, TFAlbertForSequenceClassification),
        (CamembertConfig, TFCamembertForSequenceClassification),
@@ -284,6 +310,7 @@ TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(

 TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
    [
+        # Model for Question Answering mapping
        (DistilBertConfig, TFDistilBertForQuestionAnswering),
        (AlbertConfig, TFAlbertForQuestionAnswering),
        (CamembertConfig, TFCamembertForQuestionAnswering),
@@ -302,6 +329,7 @@ TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(

 TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict(
    [
+        # Model for Token Classification mapping
        (DistilBertConfig, TFDistilBertForTokenClassification),
        (AlbertConfig, TFAlbertForTokenClassification),
        (CamembertConfig, TFCamembertForTokenClassification),
@@ -319,6 +347,7 @@ TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict(

 TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING = OrderedDict(
    [
+        # Model for Multiple Choice mapping
        (CamembertConfig, TFCamembertForMultipleChoice),
        (XLMConfig, TFXLMForMultipleChoice),
        (XLMRobertaConfig, TFXLMRobertaForMultipleChoice),
@@ -19,9 +19,10 @@ import random
 import warnings
 from typing import Dict, Optional, Tuple

+import numpy as np
 import tensorflow as tf
 from tensorflow import Tensor
-from tensorflow.keras.layers import Dense, LayerNormalization
+from tensorflow.keras.layers import Dense, Layer, LayerNormalization

 from .activations_tf import ACT2FN
 from .configuration_bart import BartConfig
@@ -43,7 +44,6 @@ from .utils import logging


 _CONFIG_FOR_DOC = "BartConfig"
-_TOKENIZER_FOR_DOC = "BartTokenizer"

 BART_START_DOCSTRING = r"""

@@ -218,22 +218,21 @@ PAST_KV_DEPRECATION_WARNING = (
 )


-class TFEncoderLayer(tf.keras.layers.Layer):
+class TFEncoderLayer(Layer):
    def __init__(self, config: BartConfig, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = config.d_model
        self.self_attn = TFAttention(
            self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout, name="self_attn"
        )
-
-        self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
-        self.dropout_wt = tf.keras.layers.Dropout(config.dropout)
+        self.normalize_before = config.normalize_before
+        self.self_attn_layer_norm = LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
+        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
-        self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout)
+        self.activation_dropout = config.activation_dropout
        self.fc1 = Dense(config.encoder_ffn_dim, name="fc1")
        self.fc2 = Dense(self.embed_dim, name="fc2")
        self.final_layer_norm = LayerNormalization(epsilon=1e-5, name="final_layer_norm")
-        self.normalize_before = config.normalize_before

    def call(self, x, encoder_padding_mask, training=False):
        """
@@ -251,8 +250,10 @@ class TFEncoderLayer(tf.keras.layers.Layer):
        if self.normalize_before:
            x = self.self_attn_layer_norm(x)
        x, self_attn_weights = self.self_attn(query=x, key=x, key_padding_mask=encoder_padding_mask)
-        assert x.shape == residual.shape, f"Self attn modified the shape of query {residual.shape} to {x.shape}"
-        x = self.dropout_wt(x, training=training)
+        assert shape_list(x) == shape_list(
+            residual
+        ), f"Self attn modified the shape of query {shape_list(residual)} to {shape_list(x)}"
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
        if not self.normalize_before:
            x = self.self_attn_layer_norm(x)
@@ -261,9 +262,9 @@ class TFEncoderLayer(tf.keras.layers.Layer):
        if self.normalize_before:
            x = self.final_layer_norm(x)
        x = self.activation_fn(self.fc1(x))
-        x = self.activation_dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.self.activation_dropout if training else 0)
        x = self.fc2(x)
-        x = self.dropout_wt(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
        if not self.normalize_before:
            x = self.final_layer_norm(x)
@@ -271,7 +272,7 @@ class TFEncoderLayer(tf.keras.layers.Layer):
        return x, self_attn_weights


-class TFBartEncoder(tf.keras.layers.Layer):
+class TFBartEncoder(Layer):
    # config_class = BartConfig
    """
    Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
@@ -289,26 +290,30 @@ class TFBartEncoder(tf.keras.layers.Layer):
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions

-        embed_dim = embed_tokens.vocab_size
-        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
        self.padding_idx = config.pad_token_id
        self.max_source_positions = config.max_position_embeddings

        self.embed_tokens = embed_tokens
-        self.embed_positions = TFLearnedPositionalEmbedding(
-            config.max_position_embeddings,
-            embed_tokens.hidden_size,
-            self.padding_idx,
-            config.extra_pos_embeddings,
-            name="embed_positions",
-        )
+        if config.static_position_embeddings:
+            self.embed_positions = TFSinusoidalPositionalEmbedding(
+                config.max_position_embeddings,
+                config.d_model,
+                name="embed_positions",
+            )
+        else:
+            self.embed_positions = TFLearnedPositionalEmbedding(
+                config.max_position_embeddings,
+                config.d_model,
+                self.padding_idx,
+                config.extra_pos_embeddings,
+                name="embed_positions",
+            )
        self.layers = [TFEncoderLayer(config, name=f"layers.{i}") for i in range(config.encoder_layers)]
-        self.layernorm_embedding = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding")
-        self.layer_norm = (
-            tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm")
-            if config.add_final_layer_norm
-            else None
+        self.layernorm_embedding = (
+            LayerNormalization(epsilon=1e-5, name="layernorm_embedding") if config.normalize_embedding else Layer()
        )
+        self.layer_norm = LayerNormalization(epsilon=1e-5, name="layer_norm") if config.add_final_layer_norm else None
        self.return_dict = config.return_dict

    def call(
@@ -347,7 +352,7 @@ class TFBartEncoder(tf.keras.layers.Layer):
            ), f"expected attention_mask._rank() to be a 2D tensor got {attention_mask._rank()}"
            attention_mask = tf.cast(attention_mask, dtype=tf.float32)
            attention_mask = (1.0 - attention_mask) * LARGE_NEGATIVE
-        inputs_embeds = self.embed_tokens(input_ids)
+        inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
        embed_pos = self.embed_positions(input_ids)
        x = inputs_embeds + embed_pos
        x = self.layernorm_embedding(x)
@@ -384,7 +389,7 @@ class TFBartEncoder(tf.keras.layers.Layer):
        return TFBaseModelOutput(last_hidden_state=x, hidden_states=encoder_states, attentions=all_attentions)


-class TFDecoderLayer(tf.keras.layers.Layer):
+class TFDecoderLayer(Layer):
    def __init__(self, config: BartConfig, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = config.d_model
@@ -397,8 +402,9 @@ class TFDecoderLayer(tf.keras.layers.Layer):
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout
+        self.normalize_before = config.normalize_before

-        self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
+        self.self_attn_layer_norm = LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
        self.encoder_attn = TFAttention(
            self.embed_dim,
            config.decoder_attention_heads,
@@ -406,10 +412,10 @@ class TFDecoderLayer(tf.keras.layers.Layer):
            encoder_decoder_attention=True,
            name="encoder_attn",
        )
-        self.encoder_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="encoder_attn_layer_norm")
+        self.encoder_attn_layer_norm = LayerNormalization(epsilon=1e-5, name="encoder_attn_layer_norm")
        self.fc1 = Dense(config.decoder_ffn_dim, name="fc1")
        self.fc2 = Dense(self.embed_dim, name="fc2")
-        self.final_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="final_layer_norm")
+        self.final_layer_norm = LayerNormalization(epsilon=1e-5, name="final_layer_norm")

    def call(
        self,
@@ -433,10 +439,12 @@ class TFDecoderLayer(tf.keras.layers.Layer):

            Tuple containing, encoded output of shape `(seq_len, batch, embed_dim)`, self_attn_weights, layer_state
        """
+        residual = x  # Make a copy of the input tensor to add later.
        if layer_state is None:
            layer_state = {}
+        if self.normalize_before:
+            x = self.self_attn_layer_norm(x)

-        residual = x  # Make a copy of the input tensor to add later.
        # next line mutates layer state and we need a copy of it
        x, self_attn_weights = self.self_attn(
            query=x,
@@ -447,9 +455,12 @@ class TFDecoderLayer(tf.keras.layers.Layer):
        )
        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
-        x = self.self_attn_layer_norm(x)
+        if not self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+        # Cross-Attention Block
        residual = x
-        # Cross-Attention
+        if self.normalize_before:
+            x = self.encoder_attn_layer_norm(x)
        x, _ = self.encoder_attn(
            query=x,
            key=encoder_hidden_states,
@@ -458,16 +469,19 @@ class TFDecoderLayer(tf.keras.layers.Layer):
        )
        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
-
-        x = self.encoder_attn_layer_norm(x)
-
+        if not self.normalize_before:
+            x = self.encoder_attn_layer_norm(x)
+        # Fully Connected
        residual = x
+        if self.normalize_before:
+            x = self.final_layer_norm(x)
        x = self.activation_fn(self.fc1(x))
        x = tf.nn.dropout(x, rate=self.activation_dropout if training else 0)
        x = self.fc2(x)
        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
-        x = self.final_layer_norm(x)
+        if not self.normalize_before:
+            x = self.final_layer_norm(x)
        return (
            x,
            self_attn_weights,
@@ -475,7 +489,7 @@ class TFDecoderLayer(tf.keras.layers.Layer):
        )  # just self_attn weights for now, following t5, layer_state = cache for decoding


-class TFBartDecoder(tf.keras.layers.Layer):
+class TFBartDecoder(Layer):
    """
    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a :class:`TFDecoderLayer`

@@ -491,26 +505,27 @@ class TFBartDecoder(tf.keras.layers.Layer):
        self.max_target_positions = config.max_position_embeddings
        self.embed_tokens = embed_tokens
        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
-        self.embed_positions = TFLearnedPositionalEmbedding(
-            config.max_position_embeddings,
-            config.d_model,
-            self.padding_idx,
-            config.extra_pos_embeddings,
-            name="embed_positions",
-        )
+        if config.static_position_embeddings:
+            self.embed_positions = TFSinusoidalPositionalEmbedding(
+                config.max_position_embeddings,
+                config.d_model,
+                name="embed_positions",
+            )
+        else:
+            self.embed_positions = TFLearnedPositionalEmbedding(
+                config.max_position_embeddings,
+                config.d_model,
+                self.padding_idx,
+                config.extra_pos_embeddings,
+                name="embed_positions",
+            )
        self.layers = [TFDecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)]
        self.layernorm_embedding = (
-            tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layernorm_embedding")
-            if config.normalize_embedding
-            else tf.identity
-        )
-        self.layer_norm = (
-            tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm")
-            if config.add_final_layer_norm
-            else None
+            LayerNormalization(epsilon=1e-5, name="layernorm_embedding") if config.normalize_embedding else Layer()
        )
+        self.layer_norm = LayerNormalization(epsilon=1e-5, name="layer_norm") if config.add_final_layer_norm else None

-        self.dropout = tf.keras.layers.Dropout(config.dropout)
+        self.dropout = config.dropout
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions
        self.use_cache = config.use_cache
@@ -553,11 +568,11 @@ class TFBartDecoder(tf.keras.layers.Layer):
            x = self.layernorm_embedding(x) + positions
        else:
            x = self.layernorm_embedding(x + positions)
-        x = self.dropout(x)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)

        # Convert to Bart output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim)
        x = tf.transpose(x, perm=(1, 0, 2))
-        assert len(encoder_hidden_states.shape) == 3, "encoder_hidden_states must be a 3D tensor"
+        assert len(shape_list(encoder_hidden_states)) == 3, "encoder_hidden_states must be a 3D tensor"
        encoder_hidden_states = tf.transpose(encoder_hidden_states, perm=(1, 0, 2))

        # decoder layers
@@ -623,7 +638,7 @@ def _reorder_buffer(attn_cache, new_order):
    return attn_cache


-class TFAttention(tf.keras.layers.Layer):
+class TFAttention(Layer):
    """Multi-headed attention from "Attention Is All You Need"""

    def __init__(
@@ -678,8 +693,10 @@ class TFAttention(tf.keras.layers.Layer):
                (default: None).
        """
        static_kv = self.encoder_decoder_attention  # value=key=encoder_hidden_states,
-        tgt_len, bsz, embed_dim = query.shape
-        assert embed_dim == self.embed_dim, f"query must be shaped {(tgt_len, bsz, self.embed_dim)} got {query.shape}"
+        tgt_len, bsz, embed_dim = shape_list(query)
+        assert (
+            embed_dim == self.embed_dim
+        ), f"query must be shaped {(tgt_len, bsz, self.embed_dim)} got {shape_list(query)}"
        # get here for encoder decoder cause of static_kv
        if layer_state is not None:  # get the last k and v for reuse
            saved_state = layer_state.get(self.cache_key, {})
@@ -718,7 +735,7 @@ class TFAttention(tf.keras.layers.Layer):
            )

        # Compute multi-headed attention
-        src_len = k.shape[1]
+        src_len = shape_list(k)[1]
        attn_weights = tf.matmul(q, k, transpose_b=True)  # shape (bsz * self.num_heads, tgt_len, src_len)

        if attn_mask is not None:
@@ -770,7 +787,7 @@ class TFLearnedPositionalEmbedding(TFSharedEmbeddings):

    def call(self, input_ids: tf.Tensor, use_cache=False):
        """Input is expected to be of size [bsz x seqlen]."""
-        bsz, seq_len = input_ids.shape[:2]
+        bsz, seq_len = shape_list(input_ids)[:2]

        if use_cache:
            positions = tf.fill((1, 1), seq_len - 1)
@@ -780,6 +797,56 @@ class TFLearnedPositionalEmbedding(TFSharedEmbeddings):
        return super().call(positions + self.offset)  # super object is not callable for some reason


+class TFSinusoidalPositionalEmbedding(tf.keras.layers.Embedding):
+    """This module produces sinusoidal positional embeddings of any length."""
+
+    def __init__(self, num_positions, embedding_dim, **kwargs):
+
+        if embedding_dim % 2 != 0:
+            raise NotImplementedError(f"odd embedding_dim {embedding_dim} not supported")
+        super().__init__(
+            num_positions,
+            embedding_dim,
+            **kwargs,
+        )
+
+    def build(self, input_shape):
+        """
+        Build shared token embedding layer Shared weights logic adapted from
+        https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24
+        """
+        super().build(input_shape)  # Instantiates self.weight so it can be loaded
+        weight: np.ndarray = self._init_weight(self.input_dim, self.output_dim)
+        self.set_weights([weight])  # overwrite self.weight to correct value
+
+    @staticmethod
+    def _init_weight(n_pos, dim):
+        """
+        Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in
+        the 2nd half of the vector. [dim // 2:]
+        """
+        position_enc = np.array(
+            [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)]
+        )
+        # index 0 is all zero
+        position_enc[:, 0 : dim // 2] = np.sin(position_enc[:, 0::2])
+        position_enc[:, dim // 2 :] = np.cos(position_enc[:, 1::2])
+        # convert to tensor
+        table = tf.convert_to_tensor(position_enc, dtype=tf.float32)
+        tf.stop_gradient(table)
+        return table
+
+    def call(self, input_ids, use_cache=False):
+        """Input is expected to be of size [bsz x seqlen]."""
+        bsz, seq_len = shape_list(input_ids)[:2]
+        if use_cache:
+            positions = tf.fill((1, 1), seq_len - 1)
+        else:
+            # starts at 0, ends at 1-seq_len
+            positions = tf.range(0, seq_len, delta=1, dtype=tf.int32, name="range")
+        return super().call(positions)
+
+
 # Public API


@@ -818,7 +885,7 @@ class TFBartModel(TFPretrainedBartModel):
        pad_token_id = self.config.pad_token_id
        if decoder_input_ids is None:
            decoder_input_ids = self._shift_right(inputs)
-        bsz, tgt_len = decoder_input_ids.shape[:2]
+        bsz, tgt_len = shape_list(decoder_input_ids)[:2]
        if decoder_attn_mask is None:
            decoder_padding_mask = make_padding_mask(decoder_input_ids, pad_token_id)
        else:
@@ -950,16 +1017,20 @@ class TFBartForConditionalGeneration(TFPretrainedBartModel):
    base_model_prefix = "model"
    authorized_missing_keys = [
        r"final_logits_bias",
-        r"encoder\.version",
-        r"decoder\.version",
-        "model.encoder.embed_tokens.weight",
-        "model.decoder.embed_tokens.weight",
+    ]
+    authorized_unexpected_keys = [
+        r"model.encoder.embed_tokens.weight",
+        r"model.decoder.embed_tokens.weight",
    ]

    def __init__(self, config: BartConfig, *args, **kwargs):
        super().__init__(config, *args, **kwargs)
        self.model = TFBartModel(config, name="model")
        self.use_cache = config.use_cache
+        # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency.
+        self.final_logits_bias = self.add_weight(
+            name="/final_logits_bias", shape=[1, config.vocab_size], initializer="zeros", trainable=False
+        )

    @add_start_docstrings_to_model_forward(BART_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
@@ -1050,6 +1121,7 @@ class TFBartForConditionalGeneration(TFPretrainedBartModel):
            return_dict=True,  # TODO(SS): this may need to change to support compilation
        )
        logits = self.model.shared(outputs.last_hidden_state, mode="linear")
+        logits = logits + self.final_logits_bias
        loss = None if labels is None else self.compute_loss(labels, logits)

        past = outputs.past_key_values if cast_bool_to_primitive(use_cache, self.config.use_cache) else None
@@ -1096,7 +1168,7 @@ class TFBartForConditionalGeneration(TFPretrainedBartModel):
            ), f"decoder cached states must be truthy. got {decoder_cached_states} from the 2nd element of past"
        assert isinstance(
            encoder_outputs, TFBaseModelOutput
-        ), "encoder_outputs should be a TFBaseModelOutput, Instead got "
+        ), f"encoder_outputs should be a TFBaseModelOutput, Instead got {type(encoder_outputs)}."
        return {
            "inputs": None,  # encoder_outputs is defined. input_ids not needed
            "encoder_outputs": encoder_outputs,
@@ -1113,7 +1185,6 @@ class TFBartForConditionalGeneration(TFPretrainedBartModel):
        reordered_past = []
        for layer_past in decoder_cached_states:
            # get the correct batch idx from decoder layer's batch dim for cross and self-attn
-
            layer_past_new = {
                attn_key: _reorder_buffer(attn_cache, beam_idx) for attn_key, attn_cache in layer_past.items()
            }
@@ -1124,26 +1195,13 @@ class TFBartForConditionalGeneration(TFPretrainedBartModel):

    def adjust_logits_during_generation(self, logits, cur_len, max_length):
        if cur_len == 1 and self.config.force_bos_token_to_be_generated:
-            logits = self._force_token_id_to_be_generated(logits, self.config.bos_token_id)
-        elif cur_len == max_length - 1 and self.config.eos_token_id is not None:
-            logits = self._force_token_id_to_be_generated(logits, self.config.eos_token_id)
-        return logits
-
-    @staticmethod
-    def _force_token_id_to_be_generated(scores, token_id) -> None:
-        """force one of token_ids to be generated by setting prob of all other tokens to 0 (logprob=-float("inf"))"""
-        output_list = []
-
-        # Is there a better way to do scores[:, [x for if x != token_id]] = -float("inf") in TF?
-        bs, vocab_size = scores.shape
-        for x in range(vocab_size):
-            if x != token_id:
-                output_list.append(tf.convert_to_tensor([-float("inf")] * bs, dtype=scores.dtype))
-            else:
-                output_list.append(scores[:, x])
-        scores = tf.stack(output_list, axis=1, name="scores")
-        assert scores.shape == (bs, vocab_size)
-        return scores
+            vocab_range = tf.constant(range(self.config.vocab_size))
+            return tf.where(vocab_range != self.config.bos_token_id, LARGE_NEGATIVE, logits)
+        elif cur_len == max_length - 1:
+            vocab_range = tf.constant(range(self.config.vocab_size))
+            return tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits)
+        else:
+            return logits

    def get_output_embeddings(self):
        return self.model.shared
@@ -666,7 +666,7 @@ class TFBertPreTrainedModel(TFPreTrainedModel):
@dataclass
 class TFBertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.TFBertForPreTrainingModel`.
+    Output type of :class:`~transformers.TFBertForPreTraining`.

    Args:
        prediction_logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
@@ -0,0 +1,47 @@
+# coding=utf-8
+# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TF BlenderBot model, ported from the fairseq repo."""
+from .configuration_blenderbot import BlenderbotConfig
+from .file_utils import add_start_docstrings, is_tf_available
+from .modeling_tf_bart import BART_START_DOCSTRING, LARGE_NEGATIVE, TFBartForConditionalGeneration
+from .utils import logging
+
+
+if is_tf_available():
+    import tensorflow as tf
+
+
+_CONFIG_FOR_DOC = "BlenderbotConfig"
+
+START_DOCSTRING = BART_START_DOCSTRING.replace(
+    "inherits from :class:`~transformers.TFPreTrainedModel`",
+    "inherits from :class:`~transformers.TFBartForConditionalGeneration`",
+).replace("BartConfig", _CONFIG_FOR_DOC)
+
+
+logger = logging.get_logger(__name__)
+
+
+@add_start_docstrings("Blenderbot model for open domain dialogue", START_DOCSTRING)
+class TFBlenderbotForConditionalGeneration(TFBartForConditionalGeneration):
+    config_class = BlenderbotConfig
+
+    def adjust_logits_during_generation(self, logits, cur_len, max_length):
+        """Never predict pad_token_id. Predict </s> when max_length is reached."""
+        vocab_range = tf.constant(range(self.config.vocab_size))
+        logits = tf.where(vocab_range == self.config.pad_token_id, LARGE_NEGATIVE, logits)
+        if cur_len == max_length - 1:
+            logits = tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits)
+        return logits
@@ -598,7 +598,7 @@ class TFElectraMainLayer(tf.keras.layers.Layer):
@dataclass
 class TFElectraForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.TFElectraForPreTrainingModel`.
+    Output type of :class:`~transformers.TFElectraForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``tf.Tensor`` of shape :obj:`(1,)`):
@@ -307,7 +307,7 @@ class TFFunnelAttentionStructure:
            pooled_pos = pos

        ref_point = pooled_pos[0] - pos[0]
-        num_remove = shift * len(pooled_pos)
+        num_remove = shift * pooled_pos.shape[0]
        max_dist = ref_point + num_remove * stride
        min_dist = pooled_pos[0] - pos[-1]

@@ -1031,7 +1031,7 @@ class TFFunnelPreTrainedModel(TFPreTrainedModel):
@dataclass
 class TFFunnelForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.FunnelForPreTrainingModel`.
+    Output type of :class:`~transformers.FunnelForPreTraining`.

    Args:
        logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`):
@@ -96,7 +96,7 @@ class TFLxmertModelOutput(ModelOutput):
@dataclass
 class TFLxmertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.LxmertForPreTrainingModel`.
+    Output type of :class:`~transformers.LxmertForPreTraining`.

    Args:
        loss (`optional`, returned when ``labels`` is provided, ``tf.Tensor`` of shape :obj:`(1,)`):
@@ -0,0 +1,52 @@
+# coding=utf-8
+# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TF Marian model, ported from the fairseq repo."""
+
+from .configuration_marian import MarianConfig
+from .file_utils import add_start_docstrings, is_tf_available
+from .modeling_tf_bart import BART_START_DOCSTRING, LARGE_NEGATIVE, TFBartForConditionalGeneration
+from .utils import logging
+
+
+if is_tf_available():
+    import tensorflow as tf
+
+
+_CONFIG_FOR_DOC = "MarianConfig"
+
+START_DOCSTRING = BART_START_DOCSTRING.replace(
+    "inherits from :class:`~transformers.TFPreTrainedModel`",
+    "inherits from :class:`~transformers.TFBartForConditionalGeneration`",
+).replace("BartConfig", _CONFIG_FOR_DOC)
+
+
+logger = logging.get_logger(__name__)
+
+
+@add_start_docstrings("Marian model for machine translation", START_DOCSTRING)
+class TFMarianMTModel(TFBartForConditionalGeneration):
+    authorized_missing_keys = [
+        r"model.encoder.embed_positions.weight",
+        r"model.decoder.embed_positions.weight",
+    ]
+    config_class = MarianConfig
+
+    def adjust_logits_during_generation(self, logits, cur_len, max_length):
+        """Never predict pad_token_id. Predict </s> when max_length is reached."""
+        vocab_range = tf.constant(range(self.config.vocab_size))
+        logits = tf.where(vocab_range == self.config.pad_token_id, LARGE_NEGATIVE, logits)
+        if cur_len == max_length - 1:
+            logits = tf.where(vocab_range != self.config.eos_token_id, LARGE_NEGATIVE, logits)
+        return logits
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TF mBART model, originally from fairseq."""
+from .configuration_mbart import MBartConfig
+from .file_utils import add_start_docstrings
+from .modeling_tf_bart import BART_START_DOCSTRING, TFBartForConditionalGeneration
+from .utils import logging
+
+
+_CONFIG_FOR_DOC = "MBartConfig"
+
+START_DOCSTRING = BART_START_DOCSTRING.replace(
+    "inherits from :class:`~transformers.TFPreTrainedModel`",
+    "inherits from :class:`~transformers.TFBartForConditionalGeneration`",
+).replace("BartConfig", _CONFIG_FOR_DOC)
+
+
+logger = logging.get_logger(__name__)
+
+
+@add_start_docstrings("mBART (multilingual BART) model for machine translation", START_DOCSTRING)
+class TFMBartForConditionalGeneration(TFBartForConditionalGeneration):
+    config_class = MBartConfig
+    # All the code is in src/transformers/modeling_tf_bart.py
@@ -833,7 +833,7 @@ class TFMobileBertPreTrainedModel(TFPreTrainedModel):
@dataclass
 class TFMobileBertForPreTrainingOutput(ModelOutput):
    """
-    Output type of :class:`~transformers.TFMobileBertForPreTrainingModel`.
+    Output type of :class:`~transformers.TFMobileBertForPreTraining`.

    Args:
        prediction_logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
@@ -227,7 +227,7 @@ class TFMaskedLMOutput(ModelOutput):

    Args:
        loss (:obj:`tf.Tensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
-            Masked languaged modeling (MLM) loss.
+            Masked language modeling (MLM) loss.
        logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
@@ -256,7 +256,7 @@ class TFSeq2SeqLMOutput(ModelOutput):

    Args:
        loss (:obj:`tf.Tensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
-            Languaged modeling loss.
+            Language modeling loss.
        logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        past_key_values (:obj:`List[tf.Tensor]`, `optional`, returned when ``use_cache=True`` is passed or when ``config.use_cache=True``):
@@ -0,0 +1,41 @@
+# coding=utf-8
+# Copyright 2020 The Facebook AI Research Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TF Pegasus model, ported from the fairseq repo."""
+from .configuration_pegasus import PegasusConfig
+from .file_utils import add_start_docstrings
+from .modeling_tf_bart import BART_START_DOCSTRING, TFBartForConditionalGeneration
+from .utils import logging
+
+
+_CONFIG_FOR_DOC = "PegasusConfig"
+
+START_DOCSTRING = BART_START_DOCSTRING.replace(
+    "inherits from :class:`~transformers.TFPreTrainedModel`",
+    "inherits from :class:`~transformers.TFBartForConditionalGeneration`",
+).replace("BartConfig", _CONFIG_FOR_DOC)
+
+
+logger = logging.get_logger(__name__)
+
+
+@add_start_docstrings("Pegasus model for summarization", START_DOCSTRING)
+class TFPegasusForConditionalGeneration(TFBartForConditionalGeneration):
+    authorized_missing_keys = [
+        r"final_logits_bias",
+        r"model.encoder.embed_positions.weight",
+        r"model.decoder.embed_positions.weight",
+    ]
+    config_class = PegasusConfig
+    # All the code is in src/transformers/modeling_tf_bart.py
@@ -1213,7 +1213,7 @@ class TFT5ForConditionalGeneration(TFT5PreTrainedModel, TFCausalLanguageModeling
            >>> model = TFT5ForConditionalGeneration.from_pretrained('t5-small')

            >>> inputs = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='tf').input_ids
-            labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2> </s>', return_tensors='tf').input_ids
+            >>> labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2> </s>', return_tensors='tf').input_ids
            >>> outputs = model(inputs, labels=labels)
            >>> loss = outputs.loss
            >>> logits = outputs.logits
@@ -550,8 +550,8 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
                    - The model is a model provided by the library (loaded with the `shortcut name` string of a
                      pretrained model).
                    - The model was saved using :func:`~transformers.TFPreTrainedModel.save_pretrained` and is reloaded
-                      by suppling the save directory.
-                    - The model is loaded by suppling a local directory as ``pretrained_model_name_or_path`` and a
+                      by supplying the save directory.
+                    - The model is loaded by supplying a local directory as ``pretrained_model_name_or_path`` and a
                      configuration JSON file named `config.json` is found in the directory.
            from_pt: (:obj:`bool`, `optional`, defaults to :obj:`False`):
                Load the model weights from a PyTorch state_dict save file (see docstring of
@@ -1257,17 +1257,17 @@ class TFXLNetLMHeadModel(TFXLNetPreTrainedModel, TFCausalLanguageModelingLoss):

        transformer_outputs = self.transformer(
            inputs,
-            attention_mask=None,
-            mems=None,
-            perm_mask=None,
-            target_mapping=None,
-            token_type_ids=None,
-            input_mask=None,
-            head_mask=None,
-            inputs_embeds=None,
-            use_cache=True,
-            output_attentions=None,
-            output_hidden_states=None,
+            attention_mask=attention_mask,
+            mems=mems,
+            perm_mask=perm_mask,
+            target_mapping=target_mapping,
+            token_type_ids=token_type_ids,
+            input_mask=input_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
@@ -784,8 +784,8 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
                    - The model is a model provided by the library (loaded with the `shortcut name` string of a
                      pretrained model).
                    - The model was saved using :func:`~transformers.PreTrainedModel.save_pretrained` and is reloaded
-                      by suppling the save directory.
-                    - The model is loaded by suppling a local directory as ``pretrained_model_name_or_path`` and a
+                      by supplying the save directory.
+                    - The model is loaded by supplying a local directory as ``pretrained_model_name_or_path`` and a
                      configuration JSON file named `config.json` is found in the directory.
            state_dict (:obj:`Dict[str, torch.Tensor]`, `optional`):
                A state dictionary to use instead of a state dictionary loaded from saved weights file.
@@ -23,9 +23,8 @@ import uuid
 import warnings
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from itertools import chain
 from os.path import abspath, exists
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID

 import numpy as np
@@ -185,57 +184,6 @@ class ArgumentHandler(ABC):
        raise NotImplementedError()


-class DefaultArgumentHandler(ArgumentHandler):
-    """
-    Default argument parser handling parameters for each :class:`~transformers.pipelines.Pipeline`.
-    """
-
-    @staticmethod
-    def handle_kwargs(kwargs: Dict) -> List:
-        if len(kwargs) == 1:
-            output = list(kwargs.values())
-        else:
-            output = list(chain(kwargs.values()))
-
-        return DefaultArgumentHandler.handle_args(output)
-
-    @staticmethod
-    def handle_args(args: Sequence[Any]) -> List[str]:
-
-        # Only one argument, let's do case by case
-        if len(args) == 1:
-            if isinstance(args[0], str):
-                return [args[0]]
-            elif not isinstance(args[0], list):
-                return list(args)
-            else:
-                return args[0]
-
-        # Multiple arguments (x1, x2, ...)
-        elif len(args) > 1:
-            if all([isinstance(arg, str) for arg in args]):
-                return list(args)
-
-            # If not instance of list, then it should instance of iterable
-            elif isinstance(args, Iterable):
-                return list(chain.from_iterable(chain(args)))
-            else:
-                raise ValueError(
-                    "Invalid input type {}. Pipeline supports Union[str, Iterable[str]]".format(type(args))
-                )
-        else:
-            return []
-
-    def __call__(self, *args, **kwargs):
-        if len(kwargs) > 0 and len(args) > 0:
-            raise ValueError("Pipeline cannot handle mixed args and kwargs")
-
-        if len(kwargs) > 0:
-            return DefaultArgumentHandler.handle_kwargs(kwargs)
-        else:
-            return DefaultArgumentHandler.handle_args(args)
-
-
 class PipelineDataFormat:
    """
    Base class for all the pipeline supported data format both for reading and writing. Supported data formats
@@ -574,7 +522,6 @@ class Pipeline(_ScikitCompat):
        self.framework = framework
        self.device = device if framework == "tf" else torch.device("cpu" if device < 0 else "cuda:{}".format(device))
        self.binary_output = binary_output
-        self._args_parser = args_parser or DefaultArgumentHandler()

        # Special handling
        if self.framework == "pt" and self.device.type == "cuda":
@@ -669,12 +616,11 @@ class Pipeline(_ScikitCompat):
                f"The model '{self.model.__class__.__name__}' is not supported for {self.task}. Supported models are {supported_models}",
            )

-    def _parse_and_tokenize(self, *args, padding=True, add_special_tokens=True, **kwargs):
+    def _parse_and_tokenize(self, inputs, padding=True, add_special_tokens=True, **kwargs):
        """
        Parse arguments and tokenize
        """
        # Parse arguments
-        inputs = self._args_parser(*args, **kwargs)
        inputs = self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
@@ -836,7 +782,7 @@ class TextGenerationPipeline(Pipeline):

    # overriding _parse_and_tokenize to allow for unusual language-modeling tokenizer arguments

-    def _parse_and_tokenize(self, *args, padding=True, add_special_tokens=True, **kwargs):
+    def _parse_and_tokenize(self, inputs, padding=True, add_special_tokens=True, **kwargs):
        """
        Parse arguments and tokenize
        """
@@ -845,7 +791,6 @@ class TextGenerationPipeline(Pipeline):
            tokenizer_kwargs = {"add_space_before_punct_symbol": True}
        else:
            tokenizer_kwargs = {}
-        inputs = self._args_parser(*args, **kwargs)
        inputs = self.tokenizer(
            inputs,
            add_special_tokens=add_special_tokens,
@@ -858,7 +803,7 @@ class TextGenerationPipeline(Pipeline):

    def __call__(
        self,
-        *args,
+        text_inputs,
        return_tensors=False,
        return_text=True,
        clean_up_tokenization_spaces=False,
@@ -890,7 +835,6 @@ class TextGenerationPipeline(Pipeline):
            - **generated_token_ids** (:obj:`torch.Tensor` or :obj:`tf.Tensor`, present when ``return_tensors=True``)
              -- The token ids of the generated text.
        """
-        text_inputs = self._args_parser(*args)

        results = []
        for prompt_text in text_inputs:
@@ -1094,7 +1038,8 @@ class ZeroShotClassificationPipeline(Pipeline):
    """

    def __init__(self, args_parser=ZeroShotClassificationArgumentHandler(), *args, **kwargs):
-        super().__init__(*args, args_parser=args_parser, **kwargs)
+        super().__init__(*args, **kwargs)
+        self._args_parser = args_parser
        if self.entailment_id == -1:
            logger.warning(
                "Failed to determine 'entailment' label id from the label2id mapping in the model config. Setting to "
@@ -1108,13 +1053,15 @@ class ZeroShotClassificationPipeline(Pipeline):
                return ind
        return -1

-    def _parse_and_tokenize(self, *args, padding=True, add_special_tokens=True, **kwargs):
+    def _parse_and_tokenize(
+        self, sequences, candidal_labels, hypothesis_template, padding=True, add_special_tokens=True, **kwargs
+    ):
        """
        Parse arguments and tokenize only_first so that hypothesis (label) is not truncated
        """
-        inputs = self._args_parser(*args, **kwargs)
+        sequence_pairs = self._args_parser(sequences, candidal_labels, hypothesis_template)
        inputs = self.tokenizer(
-            inputs,
+            sequence_pairs,
            add_special_tokens=add_special_tokens,
            return_tensors=self.framework,
            padding=padding,
@@ -1123,7 +1070,13 @@ class ZeroShotClassificationPipeline(Pipeline):

        return inputs

-    def __call__(self, sequences, candidate_labels, hypothesis_template="This example is {}.", multi_class=False):
+    def __call__(
+        self,
+        sequences: Union[str, List[str]],
+        candidate_labels,
+        hypothesis_template="This example is {}.",
+        multi_class=False,
+    ):
        """
        Classify the sequence(s) given as inputs. See the :obj:`~transformers.ZeroShotClassificationPipeline`
        documentation for more information.
@@ -1154,8 +1107,11 @@ class ZeroShotClassificationPipeline(Pipeline):
            - **labels** (:obj:`List[str]`) -- The labels sorted by order of likelihood.
            - **scores** (:obj:`List[float]`) -- The probabilities for each of the labels.
        """
+        if sequences and isinstance(sequences, str):
+            sequences = [sequences]
+
        outputs = super().__call__(sequences, candidate_labels, hypothesis_template)
-        num_sequences = 1 if isinstance(sequences, str) else len(sequences)
+        num_sequences = len(sequences)
        candidate_labels = self._args_parser._parse_labels(candidate_labels)
        reshaped_outputs = outputs.reshape((num_sequences, len(candidate_labels), -1))

@@ -1425,12 +1381,12 @@ class TokenClassificationPipeline(Pipeline):
        self.ignore_labels = ignore_labels
        self.grouped_entities = grouped_entities

-    def __call__(self, *args, **kwargs):
+    def __call__(self, inputs: Union[str, List[str]], **kwargs):
        """
        Classify each token of the text(s) given as inputs.

        Args:
-            args (:obj:`str` or :obj:`List[str]`):
+            inputs (:obj:`str` or :obj:`List[str]`):
                One or several texts (or one list of texts) for token classification.

        Return:
@@ -1444,7 +1400,8 @@ class TokenClassificationPipeline(Pipeline):
            - **index** (:obj:`int`, only present when ``self.grouped_entities=False``) -- The index of the
              corresponding token in the sentence.
        """
-        inputs = self._args_parser(*args, **kwargs)
+        if isinstance(inputs, str):
+            inputs = [inputs]
        answers = []
        for sentence in inputs:

@@ -1659,12 +1616,12 @@ class QuestionAnsweringPipeline(Pipeline):
            tokenizer=tokenizer,
            modelcard=modelcard,
            framework=framework,
-            args_parser=QuestionAnsweringArgumentHandler(),
            device=device,
            task=task,
            **kwargs,
        )

+        self._args_parser = QuestionAnsweringArgumentHandler()
        self.check_model_type(
            TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING if self.framework == "tf" else MODEL_FOR_QUESTION_ANSWERING_MAPPING
        )
@@ -2489,12 +2446,11 @@ class ConversationalPipeline(Pipeline):
            else:
                return output

-    def _parse_and_tokenize(self, *args, **kwargs):
+    def _parse_and_tokenize(self, inputs, **kwargs):
        """
        Parse arguments and tokenize, adding an EOS token at the end of the user input
        """
        # Parse arguments
-        inputs = self._args_parser(*args, **kwargs)
        inputs = self.tokenizer(inputs, add_special_tokens=False, padding=False).get("input_ids", [])
        for input in inputs:
            input.append(self.tokenizer.eos_token_id)
@@ -113,6 +113,7 @@ else:
    T5Tokenizer = None
    XLMRobertaTokenizer = None
    XLNetTokenizer = None
+    XLMProphetNetTokenizer = None

 if is_tokenizers_available():
    from .tokenization_albert_fast import AlbertTokenizerFast
@@ -134,7 +134,7 @@ class BlenderbotSmallTokenizer(PreTrainedTokenizer):
        bos_token="__start__",
        eos_token="__end__",
        unk_token="__unk__",
-        pad_token="__null",
+        pad_token="__null__",
        **kwargs
    ):
        super().__init__(unk_token=unk_token, bos_token=bos_token, eos_token=eos_token, pad_token=pad_token, **kwargs)
Author	SHA1	Message	Date
Lysandre	ee998fb662	Handle rebasing	2020-11-03 10:23:01 -05:00
Lysandre	68104c86db	Update `transformers` path	2020-11-02 18:48:26 -05:00
Lysandre	2f53251745	Cookiecutter path	2020-11-02 18:44:40 -05:00
Lysandre	f12bd2b3c0	Update file path	2020-11-02 18:39:24 -05:00
Lysandre	4a602db823	Add to testing suite	2020-11-02 18:29:56 -05:00
Lysandre	377d082d63	Code quality fixes	2020-11-02 18:20:15 -05:00
Lysandre	ed584d879f	Choose to skip either TF or PT	2020-11-02 18:20:15 -05:00
Lysandre	f969f98e01	Update docs	2020-11-02 18:20:15 -05:00
Lysandre	ca60113e1c	Style	2020-11-02 18:03:00 -05:00
Lysandre	e928962f97	Better names	2020-11-02 18:03:00 -05:00
Lysandre	29017001ad	Fixes	2020-11-02 18:03:00 -05:00
Lysandre	a986942519	README	2020-11-02 18:03:00 -05:00
Lysandre	73cc0aa2fc	Documentation	2020-11-02 18:02:47 -05:00
Lysandre	85b21f7e7f	Let's go testing in TF	2020-11-02 18:02:47 -05:00
Lysandre	333a037fd7	Encoder-Decoder in TF	2020-11-02 18:02:23 -05:00
Lysandre	cac1713315	Let's go testing	2020-11-02 18:02:23 -05:00
Lysandre	e77e584479	Encoder-Decoder	2020-11-02 18:02:23 -05:00
Lysandre	7eda3a1105	Tokenizer + Refactoring	2020-11-02 18:02:23 -05:00
Lysandre	ebd2c7f243	CI	2020-11-02 18:02:23 -05:00
Lysandre	efcdc5f749	Remove pooler	2020-11-02 18:02:14 -05:00
Lysandre	1cb130f7b6	TensorFlow	2020-11-02 18:02:14 -05:00
Lysandre	7ef5cd5940	Model templates	2020-11-02 18:02:14 -05:00
Patrick von Platen	f744b81572	add new notebooks (#8246 )	2020-11-02 20:21:55 +01:00
Patrick von Platen	dc26726df2	fix encoder decoder bug (#8243 )	2020-11-02 20:12:34 +01:00
Lysandre Debut	9a23af4aff	Add XLMProphetNetTokenizer to tokenization auto (#8245 )	2020-11-02 14:10:09 -05:00
Patrick von Platen	5b178f3c87	Create README.md	2020-11-02 20:03:44 +01:00
Sylvain GuggerandLysandre Debut	e1b1b614b1	Add line by line option to mlm/plm scripts (#8240 ) * Make line by line optional in run_mlm * Add option to disable dynamic padding * Add option to plm too and update README * Typos * More typos * Even more typos * Apply suggestions from code review Co-authored-by: Lysandre Debut <lysandre@huggingface.co> Co-authored-by: Lysandre Debut <lysandre@huggingface.co>	2020-11-02 12:27:04 -05:00
Patrick von Platen	ebec410c71	Create README.md	2020-11-02 17:53:22 +01:00
Sylvain Gugger	5406f31a1a	Fix TensorBoardCallback for older versions of PyTorch (#8239 )	2020-11-02 10:43:28 -05:00
Sylvain Gugger	d1ad4bff44	Fix bad import with PyTorch <= 1.4.1 (#8237 )	2020-11-02 10:26:37 -05:00
Lysandre Debut	3c8d401cf6	Patch reports (#8238 )	2020-11-02 10:26:25 -05:00
Martin Monperrus	93354bc779	doc: fix typo (#8235 )	2020-11-02 08:53:17 -05:00
Santiago Castro	0c92e7d9fa	Fix ignore list behavior in doctests (#8213 )	2020-11-02 08:47:37 -05:00
Nicolas Patry	84caa23301	Fix the behaviour of DefaultArgumentHandler (removing it). (#8180 ) * Some work to fix the behaviour of DefaultArgumentHandler by removing it. * Fixing specific pipelines argument checking.	2020-11-02 12:33:50 +01:00
Zhiqi Huang	00cc2d1df2	DynaBERT model cards update (#8192 ) * Update README.md * Update README.md	2020-11-02 13:19:38 +08:00
Kushal	aa79aa4e7d	Added 12 model cards for Indian Language Models (#8198 ) * Create README.md * added model cards	2020-11-02 13:17:43 +08:00
Patrick von Platen	9bd30f7cf4	[Seq2SeqTrainer] Move import to init to make file self-contained (#8194 ) * boom boom * reverse order	2020-11-01 23:31:55 +01:00
guillaume-be	1f12934df4	[Bug fix] Fixed value for BlenderBot pad token (#8205 )	2020-11-01 10:21:57 -05:00
Abi See	8f1c960ee7	Fix two bugs with --logging_first_step (#8193 ) * make sure that logging_first_step evaluates * fix bug with incorrect loss on logging_first_step * fix style * logging_first_step only logs, not evals	2020-10-30 16:45:38 -04:00
Avital Oliver	689ff74f99	Minor style improvements for the Flax BERT and RoBERTa examples (#8178 ) * Minor style improvements: 1. Use `@nn.compact` rather than `@compact` (as to not make it seem like compact is a standard Python decorator. 2. Move attribute docstrings from two `__call__` methods to comments on the attributes themselves. (This was probably a remnant from the pre-Linen version where the attributes were arguments to `call`.) * Use black on the Flax modeling code	2020-10-30 16:25:39 -04:00
Sylvain Gugger	9eb3a410cd	Remove deprecated arguments from new run_clm (#8197 )	2020-10-30 15:27:20 -04:00
TFUsersandTFUsers	00112c3539	Replace swish with silu (#8166 ) * Replace swish with silu * revert nn.silu to nn.swish due to older version * simplify optimized silu conditional and fix format * Update activations.py * Update activations_tf.py * Update modeling_flax_utils.py * Update modeling_openai.py * add swish testcase * add pytorch swish testcase * Add more robust python version check * more formatting fixes Co-authored-by: TFUsers <TFUsers@gmail.com>	2020-10-30 15:09:10 -04:00
Sylvain Gugger Lysandre Debut Thomas Wolf	cdc48ce92d	Finalize lm examples (#8188 ) * Finish the cleanup of the language-modeling examples * Update main README * Apply suggestions from code review Co-authored-by: Lysandre Debut <lysandre@huggingface.co> * Apply suggestions from code review Co-authored-by: Thomas Wolf <thomwolf@users.noreply.github.com> * Propagate changes Co-authored-by: Lysandre Debut <lysandre@huggingface.co> Co-authored-by: Thomas Wolf <thomwolf@users.noreply.github.com>	2020-10-30 14:20:18 -04:00
Sylvain Gugger	089cc1015e	Doc fixes and filter warning in wandb (#8189 )	2020-10-30 12:37:34 -04:00
Sam ShleiferandLysandre Debut	566b083eb1	TFMarian, TFMbart, TFPegasus, TFBlenderbot (#7987 ) * Start plumbing * Marian close * Small stubs for all children * Fixed bart * marian working * pegasus test is good, but failing * Checkin tests * More model files * Subtle marian, pegasus integration test failures * Works well * rm print * boom boom * Still failing model2doc * merge master * Equivalence test failing, all others fixed * cleanup * Fix embed_scale * Cleanup marian pipeline test * Undo extra changes * Smaller delta * Cleanup model testers * undo delta * fix tests import structure * cross test decorator * Cleaner set_weights * Respect authorized_unexpected_keys * No warnings * No warnings * style * Nest tf import * black * Apply suggestions from code review Co-authored-by: Lysandre Debut <lysandre@huggingface.co> * functional dropout * fixup * Fixup * style_doc * embs * shape list * delete slow force_token_id_to_be_generated func * fixup Co-authored-by: Lysandre Debut <lysandre@huggingface.co>	2020-10-30 11:23:16 -04:00
Santiago Castro	6279072f5f	Fix typo: s/languaged/language/ (#8165 )	2020-10-30 11:22:03 -04:00
Lysandre Debut	10f8c63620	Ci test tf super slow (#8007 ) * Test TF GPU CI * Change cache * Fix missing torch requirement * Fix some model tests Style * LXMERT * MobileBERT * Longformer skip test * XLNet * The rest of the tests * RAG goes OOM in multi gpu setup * YAML test files * Last fixes * Skip doctests * Fill mask tests * Yaml files * Last test fix * Style * Update cache * Change ONNX tests to slow + use tiny model	2020-10-30 10:25:48 -04:00
Nicolas Patry	7e36deec7a	Fixing some warnings in DeBerta (#8176 ) * Fixing some warnings in DeBerta * Fixing docs with their rewritten version.	2020-10-30 09:15:41 -04:00