fix positions ids bug

Co-authored-by: Julien Plu <plu.julien@gmail.com>

.circleci/config.yml

@@ -11,40 +11,21 @@ commands:
       - run:
           name: docs-only changes skip check
           command: |
-            if test -n "$CIRCLE_PR_NUMBER"
+            # pipeline.git.base_revision is not always defined, so only proceed if all external vars are defined
+            if test -n "<< pipeline.git.base_revision >>" && test -n "<< pipeline.git.revision >>" && test -n "$(git diff --name-only << pipeline.git.base_revision >>...<< pipeline.git.revision >>)"
             then
-                echo $CIRCLE_PR_NUMBER
-                resp=$(curl -Ls https://api.github.com/repos/huggingface/transformers/pulls/${CIRCLE_PR_NUMBER})
-                user=$(jq -r .user.login \<<< $resp)   # PR creator username
-                head_ref=$(jq -r .head.ref \<<< $resp) # PR user's branch name
-                echo head_ref=$head_ref, user=$user
-            fi
-            
-            if test -n "$user" && test -n "$head_ref"
-            then
-                git clone https://github.com/$user/transformers user-clone
-                cd user-clone
-                git checkout $head_ref
-                fork_point_sha=$(git merge-base --fork-point master)
-                cd -
-            fi
-
-            if test -n "$fork_point_sha" && test -n "$(git diff --name-only $fork_point_sha)"
-            then
-                git --no-pager diff --name-only $fork_point_sha
-                if git diff --name-only $fork_point_sha | egrep -qv '\.(md|rst)$'
+                if git diff --name-only << pipeline.git.base_revision >>...<< pipeline.git.revision >> | egrep -qv '\.(md|rst)$'
                 then
                     echo "Non-docs were modified in this PR, proceeding normally"
                 else
                     echo "Only docs were modified in this PR, quitting this job"
-                    # enable skipping once we get this sorted out
+                    # disable skipping for now, as circleCI's base_revision is inconsistent leading to invalid ranges
                     # circleci step halt
                 fi
             else
-                echo "Not enough data to perform a skipping check - continuing the job"
+                echo "Can't perform skipping check w/o base_revision defined, continuing the job"
             fi
 
-
 # TPU REFERENCES
 references:
     checkout_ml_testing: &checkout_ml_testing

docs/source/_static/js/custom.js

@@ -1,10 +1,10 @@
 // These two things need to be updated at each release for the version selector.
 // Last stable version
 const stableVersion = "v4.0.0"
-// Dictionary doc folder to label. The last stable version should have an empty key.
+// Dictionary doc folder to label
 const versionMapping = {
     "master": "master",
-    "": "v4.0.0 (stable)",
+    "v4.0.0": "v4.0.0",
     "v3.5.1": "v3.5.0/v3.5.1",
     "v3.4.0": "v3.4.0",
     "v3.3.1": "v3.3.0/v3.3.1",

docs/source/internal/generation_utils.rst

@@ -15,8 +15,7 @@ Utilities for Generation
 
 This page lists all the utility functions used by :meth:`~transformers.PretrainedModel.generate`,
 :meth:`~transformers.PretrainedModel.greedy_search`, :meth:`~transformers.PretrainedModel.sample`,
-:meth:`~transformers.PretrainedModel.beam_search`, :meth:`~transformers.PretrainedModel.beam_sample`, and
-:meth:`~transformers.PretrainedModel.group_beam_search`.
+:meth:`~transformers.PretrainedModel.beam_search`, and :meth:`~transformers.PretrainedModel.beam_sample`.
 
 Most of those are only useful if you are studying the code of the generate methods in the library.
 
@@ -53,12 +52,6 @@ generation.
 .. autoclass:: transformers.NoBadWordsLogitsProcessor
     :members: __call__
 
-.. autoclass:: transformers.PrefixConstrainedLogitsProcessor
-    :members: __call__
-
-.. autoclass:: transformers.HammingDiversityLogitsProcessor
-    :members: __call__
-
 BeamSearch
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/model_doc/bart.rst

@@ -105,6 +105,8 @@ BartModel
 .. autoclass:: transformers.BartModel
     :members: forward
 
+.. autofunction:: transformers.models.bart.modeling_bart._prepare_bart_decoder_inputs
+
 
 BartForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

examples/language-modeling/run_mlm_flax.py

@@ -1,663 +0,0 @@
-# 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
-"""
-import logging
-import os
-import sys
-from dataclasses import dataclass, field
-
-# You can also adapt this script on your own masked language modeling task. Pointers for this are left as comments.
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-from datasets import load_dataset
-from tqdm import tqdm
-
-import jax
-import jax.numpy as jnp
-from flax import jax_utils
-from flax.optim import Adam
-from flax.training import common_utils
-from flax.training.common_utils import get_metrics
-from jax.nn import log_softmax
-from transformers import (
-    CONFIG_MAPPING,
-    MODEL_FOR_MASKED_LM_MAPPING,
-    AutoConfig,
-    AutoTokenizer,
-    FlaxBertForMaskedLM,
-    HfArgumentParser,
-    PreTrainedTokenizerBase,
-    TensorType,
-    TrainingArguments,
-    is_tensorboard_available,
-    set_seed,
-)
-
-# Cache the result
-has_tensorboard = is_tensorboard_available()
-if has_tensorboard:
-    try:
-        from flax.metrics.tensorboard import SummaryWriter
-    except ImportError as ie:
-        has_tensorboard = False
-        print(f"Unable to display metrics through TensorBoard because some package are not installed: {ie}")
-
-else:
-    print(
-        "Unable to display metrics through TensorBoard because the package is not installed: "
-        "Please run pip install tensorboard to enable."
-    )
-
-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.
-    """
-
-    dataset_name: Optional[str] = field(
-        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
-    )
-    dataset_config_name: Optional[str] = field(
-        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
-    )
-    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.dataset_name is None and self.train_file is None and self.validation_file is None:
-            raise ValueError("Need either a dataset name or a training/validation file.")
-        else:
-            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."
-
-
-# Adapted from transformers/data/data_collator.py
-# Letting here for now, let's discuss where it should live
-@dataclass
-class FlaxDataCollatorForLanguageModeling:
-    """
-    Data collator used for language modeling. Inputs are dynamically padded to the maximum length of a batch if they
-    are not all of the same length.
-
-    Args:
-        tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
-            The tokenizer used for encoding the data.
-        mlm (:obj:`bool`, `optional`, defaults to :obj:`True`):
-            Whether or not to use masked language modeling. If set to :obj:`False`, the labels are the same as the
-            inputs with the padding tokens ignored (by setting them to -100). Otherwise, the labels are -100 for
-            non-masked tokens and the value to predict for the masked token.
-        mlm_probability (:obj:`float`, `optional`, defaults to 0.15):
-            The probability with which to (randomly) mask tokens in the input, when :obj:`mlm` is set to :obj:`True`.
-
-    .. note::
-
-        For best performance, this data collator should be used with a dataset having items that are dictionaries or
-        BatchEncoding, with the :obj:`"special_tokens_mask"` key, as returned by a
-        :class:`~transformers.PreTrainedTokenizer` or a :class:`~transformers.PreTrainedTokenizerFast` with the
-        argument :obj:`return_special_tokens_mask=True`.
-    """
-
-    tokenizer: PreTrainedTokenizerBase
-    mlm: bool = True
-    mlm_probability: float = 0.15
-
-    def __post_init__(self):
-        if self.mlm and self.tokenizer.mask_token is None:
-            raise ValueError(
-                "This tokenizer does not have a mask token which is necessary for masked language modeling. "
-                "You should pass `mlm=False` to train on causal language modeling instead."
-            )
-
-    def __call__(self, examples: List[Dict[str, np.ndarray]], pad_to_multiple_of: int) -> Dict[str, np.ndarray]:
-        # Handle dict or lists with proper padding and conversion to tensor.
-        batch = self.tokenizer.pad(examples, pad_to_multiple_of=pad_to_multiple_of, return_tensors=TensorType.NUMPY)
-
-        # If special token mask has been preprocessed, pop it from the dict.
-        special_tokens_mask = batch.pop("special_tokens_mask", None)
-        if self.mlm:
-            batch["input_ids"], batch["labels"] = self.mask_tokens(
-                batch["input_ids"], special_tokens_mask=special_tokens_mask
-            )
-        else:
-            labels = batch["input_ids"].copy()
-            if self.tokenizer.pad_token_id is not None:
-                labels[labels == self.tokenizer.pad_token_id] = -100
-            batch["labels"] = labels
-        return batch
-
-    def mask_tokens(
-            self, inputs: np.ndarray, special_tokens_mask: Optional[np.ndarray]
-    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
-        """
-        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
-        """
-        labels = inputs.copy()
-        # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
-        probability_matrix = np.full(labels.shape, self.mlm_probability)
-        special_tokens_mask = special_tokens_mask.astype("bool")
-
-        probability_matrix[special_tokens_mask] = 0.0
-        masked_indices = np.random.binomial(1, probability_matrix).astype("bool")
-        labels[~masked_indices] = -100  # We only compute loss on masked tokens
-
-        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
-        indices_replaced = np.random.binomial(1, np.full(labels.shape, 0.8)).astype("bool") & masked_indices
-        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
-
-        # 10% of the time, we replace masked input tokens with random word
-        indices_random = np.random.binomial(1, np.full(labels.shape, 0.5)).astype("bool")
-        indices_random &= masked_indices & ~indices_replaced
-
-        random_words = np.random.randint(self.tokenizer.vocab_size, size=labels.shape, dtype="i4")
-        inputs[indices_random] = random_words[indices_random]
-
-        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
-        return inputs, labels
-
-
-def create_learning_rate_scheduler(
-        factors="constant * linear_warmup * rsqrt_decay",
-        base_learning_rate=0.5,
-        warmup_steps=1000,
-        decay_factor=0.5,
-        steps_per_decay=20000,
-        steps_per_cycle=100000,
-):
-    """Creates learning rate schedule.
-    Interprets factors in the factors string which can consist of:
-    * constant: interpreted as the constant value,
-    * linear_warmup: interpreted as linear warmup until warmup_steps,
-    * rsqrt_decay: divide by square root of max(step, warmup_steps)
-    * rsqrt_normalized_decay: divide by square root of max(step/warmup_steps, 1)
-    * decay_every: Every k steps decay the learning rate by decay_factor.
-    * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.
-    Args:
-      factors: string, factors separated by "*" that defines the schedule.
-      base_learning_rate: float, the starting constant for the lr schedule.
-      warmup_steps: int, how many steps to warm up for in the warmup schedule.
-      decay_factor: float, the amount to decay the learning rate by.
-      steps_per_decay: int, how often to decay the learning rate.
-      steps_per_cycle: int, steps per cycle when using cosine decay.
-    Returns:
-      a function learning_rate(step): float -> {"learning_rate": float}, the
-      step-dependent lr.
-    """
-    factors = [n.strip() for n in factors.split("*")]
-
-    def step_fn(step):
-        """Step to learning rate function."""
-        ret = 1.0
-        for name in factors:
-            if name == "constant":
-                ret *= base_learning_rate
-            elif name == "linear_warmup":
-                ret *= jnp.minimum(1.0, step / warmup_steps)
-            elif name == "rsqrt_decay":
-                ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
-            elif name == "rsqrt_normalized_decay":
-                ret *= jnp.sqrt(warmup_steps)
-                ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
-            elif name == "decay_every":
-                ret *= decay_factor ** (step // steps_per_decay)
-            elif name == "cosine_decay":
-                progress = jnp.maximum(0.0, (step - warmup_steps) / float(steps_per_cycle))
-                ret *= jnp.maximum(0.0, 0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
-            else:
-                raise ValueError("Unknown factor %s." % name)
-        return jnp.asarray(ret, dtype=jnp.float32)
-
-    return step_fn
-
-
-def compute_metrics(logits, labels, weights, label_smoothing=0.0):
-    """Compute summary metrics."""
-    loss, normalizer = cross_entropy(logits, labels, weights, label_smoothing)
-    acc, _ = accuracy(logits, labels, weights)
-    metrics = {"loss": loss, "accuracy": acc, "normalizer": normalizer}
-    metrics = jax.lax.psum(metrics, axis_name="batch")
-    return metrics
-
-
-def accuracy(logits, targets, weights=None):
-    """Compute weighted accuracy for log probs and targets.
-    Args:
-     logits: [batch, length, num_classes] float array.
-     targets: categorical targets [batch, length] int array.
-     weights: None or array of shape [batch, length]
-    Returns:
-      Tuple of scalar loss and batch normalizing factor.
-    """
-    if logits.ndim != targets.ndim + 1:
-        raise ValueError(
-            "Incorrect shapes. Got shape %s logits and %s targets" % (str(logits.shape), str(targets.shape))
-        )
-
-    loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)
-    loss *= weights
-
-    return loss.sum(), weights.sum()
-
-
-def cross_entropy(logits, targets, weights=None, label_smoothing=0.0):
-    """Compute cross entropy and entropy for log probs and targets.
-    Args:
-     logits: [batch, length, num_classes] float array.
-     targets: categorical targets [batch, length] int array.
-     weights: None or array of shape [batch, length]
-     label_smoothing: label smoothing constant, used to determine the on and off values.
-    Returns:
-      Tuple of scalar loss and batch normalizing factor.
-    """
-    if logits.ndim != targets.ndim + 1:
-        raise ValueError(
-            "Incorrect shapes. Got shape %s logits and %s targets" % (str(logits.shape), str(targets.shape))
-        )
-
-    vocab_size = logits.shape[-1]
-    confidence = 1.0 - label_smoothing
-    low_confidence = (1.0 - confidence) / (vocab_size - 1)
-    normalizing_constant = -(
-            confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
-    )
-    soft_targets = common_utils.onehot(targets, vocab_size, on_value=confidence, off_value=low_confidence)
-
-    loss = -jnp.sum(soft_targets * log_softmax(logits), axis=-1)
-    loss = loss - normalizing_constant
-
-    if weights is not None:
-        loss = loss * weights
-        normalizing_factor = weights.sum()
-    else:
-        normalizing_factor = np.prod(targets.shape)
-
-    return loss.sum(), normalizing_factor
-
-
-def training_step(optimizer, batch, dropout_rng):
-    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
-
-    def loss_fn(params):
-        targets = batch.pop("labels")
-
-        # Hide away tokens which doesn't participate in the optimization
-        token_mask = jnp.where(targets > 0, 1.0, 0.0)
-
-        pooled, logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)
-        loss, weight_sum = cross_entropy(logits, targets, token_mask)
-        return loss / weight_sum
-
-    step = optimizer.state.step
-    lr = lr_scheduler_fn(step)
-    grad_fn = jax.value_and_grad(loss_fn)
-    loss, grad = grad_fn(optimizer.target)
-    grad = jax.lax.pmean(grad, "batch")
-    optimizer = optimizer.apply_gradient(grad, learning_rate=lr)
-
-    return loss, optimizer, new_dropout_rng
-
-
-def eval_step(params, batch):
-    """
-    Calculate evaluation metrics on a batch.
-    """
-    targets = batch.pop("labels")
-
-    # Hide away tokens which doesn't participate in the optimization
-    token_mask = jnp.where(targets > 0, 1.0, 0.0)
-    _, logits = model(**batch, params=params, train=False)
-
-    return compute_metrics(logits, targets, token_mask)
-
-
-def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
-    nb_samples = len(samples_idx)
-    samples_to_remove = nb_samples % batch_size
-
-    if samples_to_remove != 0:
-        samples_idx = samples_idx[:-samples_to_remove]
-    sections_split = nb_samples // batch_size
-    batch_idx = np.split(samples_idx, sections_split)
-    return batch_idx
-
-
-if __name__ == "__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",
-        level="NOTSET",
-        datefmt="[%X]",
-    )
-
-    # Log on each process the small summary:
-    logger = logging.getLogger(__name__)
-    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):
-    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 guarantees that only one local process can concurrently
-    # download the dataset.
-    if data_args.dataset_name is not None:
-        # Downloading and loading a dataset from the hub.
-        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
-        datasets["validation"] = load_dataset(data_args.dataset_name, data_args.dataset_config_name,
-                                              split="train[90%:95%]")
-        datasets["test"] = load_dataset(data_args.dataset_name, data_args.dataset_config_name, split="train[95%:100%]")
-    else:
-        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.validation_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.
-
-    rng = jax.random.PRNGKey(training_args.seed)
-    dropout_rngs = jax.random.split(rng, jax.local_device_count())
-
-    if model_args.model_name_or_path:
-        model = FlaxBertForMaskedLM.from_pretrained(model_args.model_name_or_path, dtype=jnp.float32, dropout_rate=0.1)
-    elif model_args.config_name:
-        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
-        model = FlaxBertForMaskedLM(config)
-        model.init(jax.random.PRNGKey(training_args.seed), (training_args.train_batch_size, model.config.max_length))
-    else:
-        config = CONFIG_MAPPING[model_args.model_type]()
-        logger.warning("You are instantiating a new config instance from scratch.")
-        model = FlaxBertForMaskedLM(config)
-        model.init(jax.random.PRNGKey(training_args.seed), (training_args.train_batch_size, model.config.max_length))
-
-    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."
-        )
-
-    # 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]
-    columns_to_remove = [column_name for column_name in column_names if column_name != text_column_name]
-
-    padding = "max_length" if data_args.pad_to_max_length else False
-
-
-    def tokenize_function(examples):
-        # Remove empty lines
-        examples = [line for line in examples if len(line) > 0 and not line.isspace()]
-        return tokenizer(
-            examples,
-            return_special_tokens_mask=True,
-            padding=padding,
-            truncation=True,
-            max_length=data_args.max_seq_length,
-        )
-
-
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        input_columns=[text_column_name],
-        batched=True,
-        num_proc=data_args.preprocessing_num_workers,
-        remove_columns=column_names,
-        load_from_cache_file=not data_args.overwrite_cache,
-    )
-
-    # Enable tensorboard only on the master node
-    if has_tensorboard and jax.host_id() == 0:
-        summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir).joinpath("logs").as_posix())
-
-    # Data collator
-    # This one will take care of randomly masking the tokens.
-    if training_args.tokenizers_parallelism:
-        os.environ["TOKENIZERS_PARALLELISM"] = "true"
-    else:
-        os.environ["TOKENIZERS_PARALLELISM"] = "false"
-    data_collator = FlaxDataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=data_args.mlm_probability)
-
-    # Setup optimizer
-    optimizer = Adam(
-        learning_rate=training_args.learning_rate,
-        weight_decay=training_args.weight_decay,
-        beta1=training_args.adam_beta1,
-        beta2=training_args.adam_beta2,
-    ).create(model.params)
-
-    # Create learning rate scheduler
-    if training_args.warmup_steps == 0:
-        training_args.warmup_steps = 1
-    lr_scheduler_fn = create_learning_rate_scheduler(
-        base_learning_rate=training_args.learning_rate, warmup_steps=training_args.warmup_steps
-    )
-
-    # Create parallel version of the training and evaluation steps
-    p_training_step = jax.pmap(training_step, "batch", donate_argnums=(0,))
-    p_eval_step = jax.pmap(eval_step, "batch", donate_argnums=(0,))
-
-    # Replicate the optimizer on each device
-    optimizer = jax_utils.replicate(optimizer)
-
-    # Store some constant
-    nb_epochs = int(training_args.num_train_epochs)
-    batch_size = int(training_args.train_batch_size)
-    eval_batch_size = int(training_args.eval_batch_size)
-
-
-    def evaluation_routine(epoch=-1):
-
-        nb_eval_samples = len(tokenized_datasets["test"])
-        eval_samples_idx = jnp.arange(nb_eval_samples)
-        eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)
-
-        eval_metrics = []
-        for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
-            samples = [tokenized_datasets["test"][int(idx)] for idx in batch_idx]
-            model_inputs = data_collator(samples, pad_to_multiple_of=16)
-
-            # Model forward
-            model_inputs = common_utils.shard(model_inputs.data)
-            metrics = p_eval_step(optimizer.target, model_inputs)
-            eval_metrics.append(metrics)
-
-        eval_metrics_np = get_metrics(eval_metrics)
-        eval_metrics_np = jax.tree_map(jnp.sum, eval_metrics_np)
-        eval_normalizer = eval_metrics_np.pop("normalizer")
-        eval_summary = jax.tree_map(lambda x: x / eval_normalizer, eval_metrics_np)
-
-        # Update progress bar
-        epochs.desc = (
-            f"Epoch... ({epoch + 1}/{nb_epochs} | Loss: {eval_summary['loss']}, Acc: {eval_summary['accuracy']})"
-        )
-        print(f"Epoch... ({epoch + 1}/{nb_epochs} | Loss: {eval_summary['loss']}, Acc: {eval_summary['accuracy']})")
-
-        # Save metrics
-        if has_tensorboard and jax.host_id() == 0:
-            for name, value in eval_summary.items():
-                summary_writer.scalar(name, value, epoch)
-
-
-    epochs = tqdm(range(nb_epochs), desc=f"Epoch ... (1/{nb_epochs})", position=0)
-
-    if training_args.initial_evaluation:
-        evaluation_routine()
-
-    for epoch in epochs:
-
-        # ======================== Training ================================
-        # Create sampling rng
-        rng, training_rng, eval_rng = jax.random.split(rng, 3)
-
-        # Generate an epoch by shuffling sampling indices from the train dataset
-        nb_training_samples = len(tokenized_datasets["train"])
-        training_samples_idx = jax.random.permutation(training_rng, jnp.arange(nb_training_samples))
-        training_batch_idx = generate_batch_splits(training_samples_idx, batch_size)
-
-        # Gather the indexes for creating the batch and do a training step
-        batches = tqdm(training_batch_idx, desc="Training...", position=1)
-        for batch_idx in batches:
-            samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
-            model_inputs = data_collator(samples, pad_to_multiple_of=16)
-
-            # Model forward
-            model_inputs = common_utils.shard(model_inputs.data)
-            loss, optimizer, dropout_rngs = p_training_step(optimizer, model_inputs, dropout_rngs)
-            batches.desc = (
-                f"Loss: {loss})"
-            )
-
-        epochs.write(f"Loss: {loss}")
-
-        # ======================== Evaluating ==============================
-        evaluation_routine(epoch)

model_cards/huggingface/CodeBERTa-language-id/README.md

@@ -1,8 +1,6 @@
 ---
 language: code
 thumbnail: https://cdn-media.huggingface.co/CodeBERTa/CodeBERTa.png
-datasets:
-- code_search_net
 ---
 
 # CodeBERTa-language-id: The World’s fanciest programming language identification algo 🤯

model_cards/huggingface/CodeBERTa-small-v1/README.md

@@ -1,8 +1,6 @@
 ---
 language: code
 thumbnail: https://cdn-media.huggingface.co/CodeBERTa/CodeBERTa.png
-datasets:
-- code_search_net
 ---
 
 # CodeBERTa

notebooks/02-transformers.ipynb

@@ -3,7 +3,6 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "YKdSeUmVSXah",
     "pycharm": {
      "is_executing": false,
      "name": "#%% md\n"
@@ -60,7 +59,6 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "TFHTP6CFSXai",
     "pycharm": {
      "name": "#%% md\n"
     }
@@ -80,7 +78,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "KnT3Jn6fSXai",
     "pycharm": {
      "is_executing": false,
      "name": "#%% code\n"
@@ -95,13 +92,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "UIQGDTIDSXai",
-    "outputId": "9851454a-c898-4fba-a389-9b16462a27c1",
     "pycharm": {
      "is_executing": false,
      "name": "#%% code\n"
@@ -111,13 +103,11 @@
     {
      "data": {
       "text/plain": [
-       "<torch.autograd.grad_mode.set_grad_enabled at 0x7f9c03e5b3c8>"
+       "<torch.autograd.grad_mode.set_grad_enabled at 0x7f10b441e890>"
       ]
      },
      "execution_count": 2,
-     "metadata": {
-      "tags": []
-     },
+     "metadata": {},
      "output_type": "execute_result"
     }
    ],
@@ -130,9 +120,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {
-    "id": "1xMDTHQXSXai",
     "pycharm": {
      "is_executing": false,
      "name": "#%% code\n"
@@ -151,7 +140,6 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "l6EcynhYSXai",
     "pycharm": {
      "name": "#%% md\n"
     }
@@ -164,13 +152,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "XgkFg52fSXai",
-    "outputId": "94b569d4-5415-4327-f39e-c9541b0a53e0",
     "pycharm": {
      "is_executing": false,
      "name": "#%% code\n"
@@ -206,17 +189,13 @@
     "print(\"Tokens PyTorch: {}\".format(tokens_pt))\n",
     "\n",
     "# Now we're ready to go through BERT with out input\n",
-    "outputs = model(tokens_pt)\n",
-    "last_hidden_state = outputs.last_hidden_state\n",
-    "pooler_output = outputs.pooler_output\n",
-    "\n",
-    "print(\"Token wise output: {}, Pooled output: {}\".format(last_hidden_state.shape, pooler_output.shape))"
+    "outputs, pooled = model(tokens_pt)\n",
+    "print(\"Token wise output: {}, Pooled output: {}\".format(outputs.shape, pooled.shape))"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "lBbvwNKXSXaj",
     "pycharm": {
      "name": "#%% md\n"
     }
@@ -236,7 +215,6 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "DCxuDWH2SXaj",
     "pycharm": {
      "name": "#%% md\n"
     }
@@ -249,13 +227,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "sgcNCdXUSXaj",
-    "outputId": "af2fb928-7c17-475b-cf81-89cfc4b1d9e5",
     "pycharm": {
      "is_executing": false,
      "name": "#%% code\n"
@@ -287,18 +260,13 @@
     "for key, value in tokens_pt2.items():\n",
     "    print(\"{}:\\n\\t{}\".format(key, value))\n",
     "\n",
-    "outputs2 = model(**tokens_pt2)\n",
-    "last_hidden_state2 = outputs2.last_hidden_state\n",
-    "pooler_output2 = outputs2.pooler_output\n",
-    "\n",
-    "print(\"Difference with previous code: ({}, {})\".format((last_hidden_state2 - last_hidden_state).sum(), (pooler_output2 - pooler_output).sum()))"
+    "outputs2, pooled2 = model(**tokens_pt2)\n",
+    "print(\"Difference with previous code: ({}, {})\".format((outputs2 - outputs).sum(), (pooled2 - pooled).sum()))"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "gC-7xGYPSXal"
-   },
+   "metadata": {},
    "source": [
     "As you can see above, calling the tokenizer provides a convenient way to generate all the required parameters\n",
     "that will go through the model. \n",
@@ -311,13 +279,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "Pl2HIcwDSXal",
-    "outputId": "22e5d010-47a9-4a12-a67d-208e5016157e",
     "pycharm": {
      "is_executing": false
     }
@@ -357,13 +320,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "1NtvWOgzSXam",
-    "outputId": "e66c47d0-e106-408d-d01c-9ac194ca3ec6",
     "pycharm": {
      "is_executing": false
     }
@@ -400,9 +358,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "vkRYm2HESXan"
-   },
+   "metadata": {},
    "source": [
     "## Frameworks interoperability\n",
     "\n",
@@ -414,9 +370,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {
-    "id": "Kubwm-wJSXan",
     "pycharm": {
      "is_executing": false
     }
@@ -432,13 +387,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 9,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "lJ13tlzOSXan",
-    "outputId": "1e4ac151-a8fc-4b34-946a-da0bc44ed0e6",
     "pycharm": {
      "is_executing": false
     }
@@ -448,8 +398,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "last_hidden_state differences: 1.0094e-05\n",
-      "pooler_output differences: 7.2969e-07\n"
+      "output differences: 1.6236e-05\n",
+      "pooled differences: -1.3039e-08\n"
      ]
     }
    ],
@@ -463,14 +413,13 @@
     "\n",
     "# Models outputs 2 values (The value for each tokens, the pooled representation of the input sentence)\n",
     "# Here we compare the output differences between PyTorch and TensorFlow.\n",
-    "for name in [\"last_hidden_state\", \"pooler_output\"]:\n",
-    "    print(\"{} differences: {:.5}\".format(name, (output_tf[name].numpy() - output_pt[name].numpy()).sum()))"
+    "for name, o_tf, o_pt in zip([\"output\", \"pooled\"], output_tf, output_pt):\n",
+    "    print(\"{} differences: {:.5}\".format(name, (o_tf.numpy() - o_pt.numpy()).sum()))"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "CQf_fpApSXao",
     "pycharm": {
      "name": "#%% md\n"
     }
@@ -493,85 +442,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 10,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/",
-     "height": 185,
-     "referenced_widgets": [
-      "fcffccb427714665bec7d621d00d4ce3",
-      "7aa02ef05fe64489ad6c969dd92d1b07",
-      "817f53e3fa5c43e29ad1b4410c3df7db",
-      "50a441d7a43c4a809a09505b36e83375",
-      "497ba6a585a147459f1346c0661d5c94",
-      "a18c5319739141af9a255bccf25f6884",
-      "cf319210c8134cdba487525c49e4813b",
-      "f1f0272d9bea4e9fad3be8646b45d629",
-      "530b39d56f6b4e0caae3317855c4bcf4",
-      "c5e735694f2c4813a1d6f0d867119f67",
-      "8d53b8dc213f405d8187f3c1f005826d",
-      "d492afe626804d95a5cfac0550913190",
-      "a657a312068b43529afed2050bce572f",
-      "fe230ff13a82400f97cf6f292e8851ba",
-      "be97cf2269d748f3b1a916b5376f7736",
-      "74bd90a09da74db5bcbbe86f044bd664"
-     ]
-    },
-    "id": "wfxMOXb-SXao",
-    "outputId": "fa667556-fbf2-4c86-fc7e-9e3d3ec9da88",
     "pycharm": {
      "is_executing": false
     }
    },
    "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fcffccb427714665bec7d621d00d4ce3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=411.0, style=ProgressStyle(description_…"
-      ]
-     },
-     "metadata": {
-      "tags": []
-     },
-     "output_type": "display_data"
-    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "530b39d56f6b4e0caae3317855c4bcf4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=263273408.0, style=ProgressStyle(descri…"
-      ]
-     },
-     "metadata": {
-      "tags": []
-     },
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n",
-      "CPU times: user 64.4 ms, sys: 0 ns, total: 64.4 ms\n",
-      "Wall time: 72.9 ms\n",
-      "CPU times: user 130 ms, sys: 124 µs, total: 130 ms\n",
-      "Wall time: 131 ms\n"
+      "CPU times: user 232 ms, sys: 0 ns, total: 232 ms\n",
+      "Wall time: 21.1 ms\n",
+      "CPU times: user 511 ms, sys: 0 ns, total: 511 ms\n",
+      "Wall time: 43.9 ms\n"
      ]
     }
    ],
@@ -591,9 +476,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "7lSIc7FbSXao"
-   },
+   "metadata": {},
    "source": [
     "## Community provided models\n",
     "\n",
@@ -607,13 +490,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 11,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "cxLYnadGSXao",
-    "outputId": "70ab584a-e795-490a-8c6a-06e034b3df3d",
     "pycharm": {
      "is_executing": false
     }
@@ -627,7 +505,7 @@
       "Tokens (str)      : ['[CLS]', 'Hug', '##ging', 'Fac', '##e', 'ist', 'eine', 'französische', 'Firma', 'mit', 'Sitz', 'in', 'New', '-', 'York', '.', '[SEP]']\n",
       "Tokens (attn_mask): [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
       "\n",
-      "Token wise output: torch.Size([1, 17, 768]), Pooled output: torch.Size([1, 768])\n"
+      "Token wise output: torch.Size([1, 7, 768]), Pooled output: torch.Size([1, 768])\n"
      ]
     }
    ],
@@ -645,19 +523,13 @@
     "print(\"Tokens (attn_mask): {}\".format(de_input['attention_mask'].tolist()[0]))\n",
     "print()\n",
     "\n",
-    "outputs_de = de_bert(**de_input)\n",
-    "last_hidden_state_de = outputs_de.last_hidden_state\n",
-    "pooler_output_de = outputs_de.pooler_output\n",
+    "output_de, pooled_de = de_bert(**de_input)\n",
     "\n",
-    "print(\"Token wise output: {}, Pooled output: {}\".format(last_hidden_state_de.shape, pooler_output_de.shape))"
+    "print(\"Token wise output: {}, Pooled output: {}\".format(outputs.shape, pooled.shape))"
    ]
   }
  ],
  "metadata": {
-  "colab": {
-   "name": "02-transformers.ipynb",
-   "provenance": []
-  },
   "kernelspec": {
    "display_name": "Python 3",
    "language": "python",
@@ -673,7 +545,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.3"
+   "version": "3.7.4"
   },
   "pycharm": {
    "stem_cell": {
@@ -683,502 +555,8 @@
     },
     "source": []
    }
-  },
-  "widgets": {
-   "application/vnd.jupyter.widget-state+json": {
-    "497ba6a585a147459f1346c0661d5c94": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": "initial"
-     }
-    },
-    "50a441d7a43c4a809a09505b36e83375": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_f1f0272d9bea4e9fad3be8646b45d629",
-      "placeholder": "​",
-      "style": "IPY_MODEL_cf319210c8134cdba487525c49e4813b",
-      "value": " 411/411 [00:19&lt;00:00, 21.5B/s]"
-     }
-    },
-    "530b39d56f6b4e0caae3317855c4bcf4": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_8d53b8dc213f405d8187f3c1f005826d",
-       "IPY_MODEL_d492afe626804d95a5cfac0550913190"
-      ],
-      "layout": "IPY_MODEL_c5e735694f2c4813a1d6f0d867119f67"
-     }
-    },
-    "74bd90a09da74db5bcbbe86f044bd664": {
-     "model_module": "@jupyter-widgets/base",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "7aa02ef05fe64489ad6c969dd92d1b07": {
-     "model_module": "@jupyter-widgets/base",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "817f53e3fa5c43e29ad1b4410c3df7db": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "success",
-      "description": "Downloading: 100%",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_a18c5319739141af9a255bccf25f6884",
-      "max": 411,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_497ba6a585a147459f1346c0661d5c94",
-      "value": 411
-     }
-    },
-    "8d53b8dc213f405d8187f3c1f005826d": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "FloatProgressModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "FloatProgressModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "ProgressView",
-      "bar_style": "success",
-      "description": "Downloading: 100%",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_fe230ff13a82400f97cf6f292e8851ba",
-      "max": 263273408,
-      "min": 0,
-      "orientation": "horizontal",
-      "style": "IPY_MODEL_a657a312068b43529afed2050bce572f",
-      "value": 263273408
-     }
-    },
-    "a18c5319739141af9a255bccf25f6884": {
-     "model_module": "@jupyter-widgets/base",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "a657a312068b43529afed2050bce572f": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "ProgressStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "ProgressStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "bar_color": null,
-      "description_width": "initial"
-     }
-    },
-    "be97cf2269d748f3b1a916b5376f7736": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "c5e735694f2c4813a1d6f0d867119f67": {
-     "model_module": "@jupyter-widgets/base",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "cf319210c8134cdba487525c49e4813b": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "DescriptionStyleModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "DescriptionStyleModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "StyleView",
-      "description_width": ""
-     }
-    },
-    "d492afe626804d95a5cfac0550913190": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "HTMLModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HTMLModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HTMLView",
-      "description": "",
-      "description_tooltip": null,
-      "layout": "IPY_MODEL_74bd90a09da74db5bcbbe86f044bd664",
-      "placeholder": "​",
-      "style": "IPY_MODEL_be97cf2269d748f3b1a916b5376f7736",
-      "value": " 263M/263M [00:06&lt;00:00, 43.5MB/s]"
-     }
-    },
-    "f1f0272d9bea4e9fad3be8646b45d629": {
-     "model_module": "@jupyter-widgets/base",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    },
-    "fcffccb427714665bec7d621d00d4ce3": {
-     "model_module": "@jupyter-widgets/controls",
-     "model_name": "HBoxModel",
-     "state": {
-      "_dom_classes": [],
-      "_model_module": "@jupyter-widgets/controls",
-      "_model_module_version": "1.5.0",
-      "_model_name": "HBoxModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/controls",
-      "_view_module_version": "1.5.0",
-      "_view_name": "HBoxView",
-      "box_style": "",
-      "children": [
-       "IPY_MODEL_817f53e3fa5c43e29ad1b4410c3df7db",
-       "IPY_MODEL_50a441d7a43c4a809a09505b36e83375"
-      ],
-      "layout": "IPY_MODEL_7aa02ef05fe64489ad6c969dd92d1b07"
-     }
-    },
-    "fe230ff13a82400f97cf6f292e8851ba": {
-     "model_module": "@jupyter-widgets/base",
-     "model_name": "LayoutModel",
-     "state": {
-      "_model_module": "@jupyter-widgets/base",
-      "_model_module_version": "1.2.0",
-      "_model_name": "LayoutModel",
-      "_view_count": null,
-      "_view_module": "@jupyter-widgets/base",
-      "_view_module_version": "1.2.0",
-      "_view_name": "LayoutView",
-      "align_content": null,
-      "align_items": null,
-      "align_self": null,
-      "border": null,
-      "bottom": null,
-      "display": null,
-      "flex": null,
-      "flex_flow": null,
-      "grid_area": null,
-      "grid_auto_columns": null,
-      "grid_auto_flow": null,
-      "grid_auto_rows": null,
-      "grid_column": null,
-      "grid_gap": null,
-      "grid_row": null,
-      "grid_template_areas": null,
-      "grid_template_columns": null,
-      "grid_template_rows": null,
-      "height": null,
-      "justify_content": null,
-      "justify_items": null,
-      "left": null,
-      "margin": null,
-      "max_height": null,
-      "max_width": null,
-      "min_height": null,
-      "min_width": null,
-      "object_fit": null,
-      "object_position": null,
-      "order": null,
-      "overflow": null,
-      "overflow_x": null,
-      "overflow_y": null,
-      "padding": null,
-      "right": null,
-      "top": null,
-      "visibility": null,
-      "width": null
-     }
-    }
-   }
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }

src/transformers/__init__.py

@@ -321,14 +321,12 @@ if is_torch_available():
     )
     from .generation_beam_search import BeamScorer, BeamSearchScorer
     from .generation_logits_process import (
-        HammingDiversityLogitsProcessor,
         LogitsProcessor,
         LogitsProcessorList,
         LogitsWarper,
         MinLengthLogitsProcessor,
         NoBadWordsLogitsProcessor,
         NoRepeatNGramLogitsProcessor,
-        PrefixConstrainedLogitsProcessor,
         RepetitionPenaltyLogitsProcessor,
         TemperatureLogitsWarper,
         TopKLogitsWarper,
@@ -378,7 +376,6 @@ if is_torch_available():
         BartForQuestionAnswering,
         BartForSequenceClassification,
         BartModel,
-        BartPretrainedModel,
         PretrainedBartModel,
     )
     from .models.bert import (
@@ -937,7 +934,7 @@ else:
 
 if is_flax_available():
     from .models.auto import FLAX_MODEL_MAPPING, FlaxAutoModel
-    from .models.bert import FlaxBertForMaskedLM, FlaxBertModel
+    from .models.bert import FlaxBertModel
     from .models.roberta import FlaxRobertaModel
 else:
     # Import the same objects as dummies to get them in the namespace.

src/transformers/configuration_utils.py

@@ -95,12 +95,6 @@ class PretrainedConfig(object):
           sentences are finished per batch or not.
         - **num_beams** (:obj:`int`, `optional`, defaults to 1) -- Number of beams for beam search that will be used by
           default in the :obj:`generate` method of the model. 1 means no beam search.
-        - **num_beam_groups** (:obj:`int`, `optional`, defaults to 1) -- Number of groups to divide :obj:`num_beams`
-          into in order to ensure diversity among different groups of beams that will be used by default in the
-          :obj:`generate` method of the model. 1 means no group beam search.
-        - **diversity_penalty** (:obj:`float`, `optional`, defaults to 0.0) -- Value to control diversity for group
-          beam search. that will be used by default in the :obj:`generate` method of the model. 0 means no diversity
-          penalty. The higher the penalty, the more diverse are the outputs.
         - **temperature** (:obj:`float`, `optional`, defaults to 1) -- The value used to module the next token
           probabilities that will be used by default in the :obj:`generate` method of the model. Must be strictly
           positive.
@@ -191,8 +185,6 @@ class PretrainedConfig(object):
         self.do_sample = kwargs.pop("do_sample", False)
         self.early_stopping = kwargs.pop("early_stopping", False)
         self.num_beams = kwargs.pop("num_beams", 1)
-        self.num_beam_groups = kwargs.pop("num_beam_groups", 1)
-        self.diversity_penalty = kwargs.pop("diversity_penalty", 0.0)
         self.temperature = kwargs.pop("temperature", 1.0)
         self.top_k = kwargs.pop("top_k", 50)
         self.top_p = kwargs.pop("top_p", 1.0)

src/transformers/generation_beam_search.py

@@ -122,12 +122,9 @@ class BeamSearchScorer(BeamScorer):
     Adapted in part from `Facebook's XLM beam search code
     <https://github.com/facebookresearch/XLM/blob/9e6f6814d17be4fe5b15f2e6c43eb2b2d76daeb4/src/model/transformer.py#L529>`__.
 
-    Reference for the diverse beam search algorithm and implementation `Ashwin Kalyan's DBS implementation
-    <https://github.com/ashwinkalyan/dbs/blob/master/dbs/beam_utils.lua>`__
-
     Args:
         batch_size (:obj:`int`):
-            Batch Size of :obj:`input_ids` for which standard beam search decoding is run in parallel.
+            Batch Size of :obj:`input_ids` for which beam search decoding is run in parallel.
         max_length (:obj:`int`):
             The maximum length of the sequence to be generated.
         num_beams (:obj:`int`):
@@ -144,9 +141,6 @@ class BeamSearchScorer(BeamScorer):
         num_beam_hyps_to_keep (:obj:`int`, `optional`, defaults to 1):
             The number of beam hypotheses that shall be returned upon calling
             :meth:`~transformer.BeamSearchScorer.finalize`.
-        num_beam_groups (:obj:`int`):
-            Number of groups to divide :obj:`num_beams` into in order to ensure diversity among different groups of
-            beams. See `this paper <https://arxiv.org/pdf/1610.02424.pdf>`__ for more details.
     """
 
     def __init__(
@@ -158,7 +152,6 @@ class BeamSearchScorer(BeamScorer):
         length_penalty: Optional[float] = 1.0,
         do_early_stopping: Optional[bool] = False,
         num_beam_hyps_to_keep: Optional[int] = 1,
-        num_beam_groups: Optional[int] = 1,
     ):
         self.max_length = max_length
         self.num_beams = num_beams
@@ -166,8 +159,6 @@ class BeamSearchScorer(BeamScorer):
         self.length_penalty = length_penalty
         self.do_early_stopping = do_early_stopping
         self.num_beam_hyps_to_keep = num_beam_hyps_to_keep
-        self.num_beam_groups = num_beam_groups
-        self.group_size = self.num_beams // self.num_beam_groups
 
         self._is_init = False
         self._beam_hyps = [
@@ -186,12 +177,6 @@ class BeamSearchScorer(BeamScorer):
                 f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1, one should make use of `greedy_search` instead."
             )
 
-        if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0):
-            raise ValueError(
-                f"`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` "
-                f"has to be divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}."
-            )
-
     @property
     def is_done(self) -> bool:
         return self._done.all()
@@ -207,12 +192,12 @@ class BeamSearchScorer(BeamScorer):
     ) -> Tuple[torch.Tensor]:
         cur_len = input_ids.shape[-1]
         batch_size = len(self._beam_hyps)
-        assert batch_size == (input_ids.shape[0] // self.group_size)
+        assert batch_size == (input_ids.shape[0] // self.num_beams)
 
         device = input_ids.device
-        next_beam_scores = torch.zeros((batch_size, self.group_size), dtype=next_scores.dtype, device=device)
-        next_beam_tokens = torch.zeros((batch_size, self.group_size), dtype=next_tokens.dtype, device=device)
-        next_beam_indices = torch.zeros((batch_size, self.group_size), dtype=next_indices.dtype, device=device)
+        next_beam_scores = torch.zeros((batch_size, self.num_beams), dtype=next_scores.dtype, device=device)
+        next_beam_tokens = torch.zeros((batch_size, self.num_beams), dtype=next_tokens.dtype, device=device)
+        next_beam_indices = torch.zeros((batch_size, self.num_beams), dtype=next_indices.dtype, device=device)
 
         for batch_idx, beam_hyp in enumerate(self._beam_hyps):
             if self._done[batch_idx]:
@@ -233,11 +218,11 @@ class BeamSearchScorer(BeamScorer):
             for beam_token_rank, (next_token, next_score, next_index) in enumerate(
                 zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx])
             ):
-                batch_beam_idx = batch_idx * self.group_size + next_index
+                batch_beam_idx = batch_idx * self.num_beams + next_index
                 # add to generated hypotheses if end of sentence
                 if (eos_token_id is not None) and (next_token.item() == eos_token_id):
                     # if beam_token does not belong to top num_beams tokens, it should not be added
-                    is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size
+                    is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.num_beams
                     if is_beam_token_worse_than_top_num_beams:
                         continue
                     beam_hyp.add(
@@ -252,12 +237,12 @@ class BeamSearchScorer(BeamScorer):
                     beam_idx += 1
 
                 # once the beam for next step is full, don't add more tokens to it.
-                if beam_idx == self.group_size:
+                if beam_idx == self.num_beams:
                     break
 
-            if beam_idx < self.group_size:
+            if beam_idx < self.num_beams:
                 raise ValueError(
-                    f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id: {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected."
+                    f"At most {self.num_beams} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id: {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected."
                 )
 
             # Check if we are done so that we can save a pad step if all(done)
@@ -289,8 +274,7 @@ class BeamSearchScorer(BeamScorer):
             if self._done[batch_idx]:
                 continue
 
-            # all open beam hypotheses are added to the beam hypothesis
-            # beam hypothesis class automatically keeps the best beams
+            # need to add best num_beams hypotheses to generated hyps
             for beam_id in range(self.num_beams):
                 batch_beam_idx = batch_idx * self.num_beams + beam_id
                 final_score = final_beam_scores[batch_beam_idx].item()

src/transformers/generation_logits_process.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import inspect
 import math
 from abc import ABC
 from typing import Callable, Iterable, List
@@ -38,8 +37,6 @@ LOGITS_PROCESSOR_INPUTS_DOCSTRING = r"""
         scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.vocab_size)`):
             Prediction scores of a language modeling head. These can be scores for each vocabulary token before SoftMax
             or scores for each vocabulary token after SoftMax.
-        kwargs:
-            Additional logits processor specific kwargs.
 
     Return:
         :obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.vocab_size)`: The processed prediction scores.
@@ -78,16 +75,9 @@ class LogitsProcessorList(list):
     """
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
         for processor in self:
-            function_args = inspect.signature(processor.__call__).parameters
-            if len(function_args) > 2:
-                assert all(
-                    arg in kwargs for arg in list(function_args.keys())[2:]
-                ), f"Make sure that all the required parameters: {list(function_args.keys())} for {processor.__class__} are passed to the logits processor."
-                scores = processor(input_ids, scores, **kwargs)
-            else:
-                scores = processor(input_ids, scores)
+            scores = processor(input_ids, scores)
         return scores
 
 
@@ -410,65 +400,3 @@ class PrefixConstrainedLogitsProcessor(LogitsProcessor):
                 mask[batch_id * self._num_beams + beam_id, self._prefix_allowed_tokens_fn(batch_id, sent)] = 0
 
         return scores + mask
-
-
-class HammingDiversityLogitsProcessor(LogitsProcessor):
-    r"""
-    :class:`transformers.LogitsProcessor` that enforces diverse beam search. Note that this logits processor is only
-    effective for :meth:`transformers.PretrainedModel.group_beam_search`. See `Diverse Beam Search: Decoding Diverse
-    Solutions from Neural Sequence Models <https://arxiv.org/pdf/1610.02424.pdf>`__ for more details.
-
-    Args:
-        diversity_penalty (:obj:`float`):
-            This value is subtracted from a beam's score if it generates a token same as any beam from other group at a
-            particular time. Note that :obj:`diversity_penalty` is only effective if ``group beam search`` is enabled.
-        num_beams (:obj:`int`):
-            Number of beams used for group beam search. See `this paper <https://arxiv.org/pdf/1610.02424.pdf>`__ for
-            more details.
-        num_beam_groups (:obj:`int`):
-            Number of groups to divide :obj:`num_beams` into in order to ensure diversity among different groups of
-            beams. See `this paper <https://arxiv.org/pdf/1610.02424.pdf>`__ for more details.
-    """
-
-    def __init__(self, diversity_penalty: float, num_beams: int, num_beam_groups: int):
-        if not isinstance(diversity_penalty, float) or (not diversity_penalty > 0.0):
-            raise ValueError("`diversity_penalty` should be a float strictly larger than 0.")
-        self._diversity_penalty = diversity_penalty
-        if not isinstance(num_beams, int) or num_beams < 2:
-            raise ValueError("`num_beams` should be an integer strictly larger than 1.")
-        self._num_beams = num_beams
-        if not isinstance(num_beam_groups, int) or num_beam_groups < 2:
-            raise ValueError("`num_beam_groups` should be an integer strictly larger than 1.")
-        if num_beam_groups > num_beams:
-            raise ValueError("`beam_groups` has to be smaller or equal to `num_beams`.")
-        if num_beam_groups > num_beams:
-            raise ValueError("`beam_groups` has to be smaller or equal to `num_beams`")
-        self._num_sub_beams = num_beams // num_beam_groups
-
-    def __call__(
-        self,
-        input_ids: torch.LongTensor,
-        scores: torch.FloatTensor,
-        current_tokens: torch.LongTensor,
-        beam_group_idx: int,
-    ) -> torch.FloatTensor:
-        # hamming diversity: penalise using same token in current group which was used in previous groups at
-        # the same time step
-        batch_size = current_tokens.shape[0] // self._num_beams
-        group_start_idx = beam_group_idx * self._num_sub_beams
-        group_end_idx = min(group_start_idx + self._num_sub_beams, self._num_beams)
-        group_size = group_end_idx - group_start_idx
-        vocab_size = scores.shape[-1]
-
-        if group_start_idx == 0:
-            return scores
-
-        for batch_idx in range(batch_size):
-            # predicted tokens of last time step of previous groups
-            previous_group_tokens = current_tokens[
-                batch_idx * self._num_beams : batch_idx * self._num_beams + group_start_idx
-            ]
-            token_frequency = torch.bincount(previous_group_tokens, minlength=vocab_size).to(scores.device)
-            scores[batch_idx * group_size : (batch_idx + 1) * group_size] -= self._diversity_penalty * token_frequency
-
-        return scores

src/transformers/generation_utils.py

@@ -22,7 +22,6 @@ from torch.nn import functional as F
 from .file_utils import ModelOutput
 from .generation_beam_search import BeamScorer, BeamSearchScorer
 from .generation_logits_process import (
-    HammingDiversityLogitsProcessor,
     LogitsProcessorList,
     MinLengthLogitsProcessor,
     NoBadWordsLogitsProcessor,
@@ -262,8 +261,6 @@ class GenerationMixin:
         eos_token_id: int,
         prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]],
         num_beams: int,
-        num_beam_groups: int,
-        diversity_penalty: float,
     ) -> LogitsProcessorList:
         """
         This class returns a :obj:`~transformers.LogitsProcessorList` list object that contains all relevant
@@ -278,18 +275,11 @@ class GenerationMixin:
         bad_words_ids = bad_words_ids if bad_words_ids is not None else self.config.bad_words_ids
         min_length = min_length if min_length is not None else self.config.min_length
         eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
-        diversity_penalty = diversity_penalty if diversity_penalty is not None else self.config.diversity_penalty
         # instantiate processors list
         processors = LogitsProcessorList()
 
         # the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files
         # all samplers can be found in `generation_utils_samplers.py`
-        if diversity_penalty is not None and diversity_penalty > 0.0:
-            processors.append(
-                HammingDiversityLogitsProcessor(
-                    diversity_penalty=diversity_penalty, num_beams=num_beams, num_beam_groups=num_beam_groups
-                )
-            )
         if repetition_penalty is not None and repetition_penalty != 1.0:
             processors.append(RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty))
         if no_repeat_ngram_size is not None and no_repeat_ngram_size > 0:
@@ -324,8 +314,6 @@ class GenerationMixin:
         num_return_sequences: Optional[int] = None,
         decoder_start_token_id: Optional[int] = None,
         use_cache: Optional[bool] = None,
-        num_beam_groups: Optional[int] = None,
-        diversity_penalty: Optional[float] = None,
         prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
         **model_kwargs
     ) -> torch.LongTensor:
@@ -393,13 +381,6 @@ class GenerationMixin:
             use_cache: (:obj:`bool`, `optional`, defaults to :obj:`True`):
                 Whether or not the model should use the past last key/values attentions (if applicable to the model) to
                 speed up decoding.
-            num_beam_groups (:obj:`int`, `optional`, defaults to 1):
-                Number of groups to divide :obj:`num_beams` into in order to ensure diversity among different groups of
-                beams. `this paper <https://arxiv.org/pdf/1610.02424.pdf>`__ for more details.
-            diversity_penalty (:obj:`float`, `optional`, defaults to 0.0):
-                This value is subtracted from a beam's score if it generates a token same as any beam from other group
-                at a particular time. Note that :obj:`diversity_penalty` is only effective if ``group beam search`` is
-                enabled.
             prefix_allowed_tokens_fn: (:obj:`Callable[[int, torch.Tensor], List[int]]`, `optional`):
                 If provided, this function constraints the beam search to allowed tokens only at each step. If not
                 provided no constraint is applied. This function takes 2 arguments :obj:`inputs_ids` and the batch ID
@@ -472,7 +453,6 @@ class GenerationMixin:
 
         # set init values
         num_beams = num_beams if num_beams is not None else self.config.num_beams
-        num_beam_groups = num_beam_groups if num_beam_groups is not None else self.config.num_beam_groups
         max_length = max_length if max_length is not None else self.config.max_length
         do_sample = do_sample if do_sample is not None else self.config.do_sample
         num_return_sequences = (
@@ -511,17 +491,10 @@ class GenerationMixin:
                 raise ValueError("Make sure that `model_kwargs` include `encoder_outputs` of type `ModelOutput`.")
 
         # determine generation mode
-        is_greedy_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is False
-        is_sample_gen_mode = (num_beams == 1) and (num_beam_groups == 1) and do_sample is True
-        is_beam_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is False
-        is_beam_sample_gen_mode = (num_beams > 1) and (num_beam_groups == 1) and do_sample is True
-        is_group_beam_gen_mode = (num_beams > 1) and (num_beam_groups > 1)
-        if num_beam_groups > num_beams:
-            raise ValueError("`num_beam_groups` has to be smaller or equal to `num_beams`")
-        if is_group_beam_gen_mode and do_sample is True:
-            raise ValueError(
-                "Diverse beam search cannot be used in sampling mode. Make sure that `do_sample` is set to `False`."
-            )
+        is_greedy_gen_mode = (num_beams == 1) and do_sample is False
+        is_sample_gen_mode = (num_beams == 1) and do_sample is True
+        is_beam_gen_mode = (num_beams > 1) and do_sample is False
+        is_beam_sample_gen_mode = (num_beams > 1) and do_sample is True
 
         # set model_kwargs
         model_kwargs["use_cache"] = use_cache
@@ -535,8 +508,6 @@ class GenerationMixin:
             eos_token_id=eos_token_id,
             prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
             num_beams=num_beams,
-            num_beam_groups=num_beam_groups,
-            diversity_penalty=diversity_penalty,
         )
 
         if is_greedy_gen_mode:
@@ -648,42 +619,6 @@ class GenerationMixin:
                 **model_kwargs,
             )
 
-        elif is_group_beam_gen_mode:
-            batch_size = input_ids.shape[0]
-
-            length_penalty = length_penalty if length_penalty is not None else self.config.length_penalty
-            early_stopping = early_stopping if early_stopping is not None else self.config.early_stopping
-
-            if num_return_sequences > num_beams:
-                raise ValueError("`num_return_sequences` has to be smaller or equal to `num_beams`.")
-
-            if num_beams % num_beam_groups != 0:
-                raise ValueError("`num_beams` should be divisible by `num_beam_groups` for group beam search.")
-
-            diverse_beam_scorer = BeamSearchScorer(
-                batch_size=batch_size,
-                max_length=max_length,
-                num_beams=num_beams,
-                device=self.device,
-                length_penalty=length_penalty,
-                do_early_stopping=early_stopping,
-                num_beam_hyps_to_keep=num_return_sequences,
-                num_beam_groups=num_beam_groups,
-            )
-            # interleave with `num_beams`
-            input_ids, model_kwargs = self._expand_inputs_for_generation(
-                input_ids, expand_size=num_beams, is_encoder_decoder=self.config.is_encoder_decoder, **model_kwargs
-            )
-            return self.group_beam_search(
-                input_ids,
-                diverse_beam_scorer,
-                logits_processor=logits_processor,
-                max_length=max_length,
-                pad_token_id=pad_token_id,
-                eos_token_id=eos_token_id,
-                **model_kwargs,
-            )
-
     def greedy_search(
         self,
         input_ids: torch.LongTensor,
@@ -1273,213 +1208,6 @@ class GenerationMixin:
 
         return decoded
 
-    def group_beam_search(
-        self,
-        input_ids: torch.LongTensor,
-        beam_scorer: BeamScorer,
-        logits_processor: Optional[LogitsProcessorList] = None,
-        max_length: Optional[int] = None,
-        pad_token_id: Optional[int] = None,
-        eos_token_id: Optional[int] = None,
-        **model_kwargs
-    ):
-        r"""
-        Generates sequences for models with a language modeling head using beam search decoding.
-
-        Parameters:
-
-            input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
-                The sequence used as a prompt for the generation. If :obj:`None` the method initializes it as an empty
-                :obj:`torch.LongTensor` of shape :obj:`(1,)`.
-            beam_scorer (:obj:`BeamScorer`):
-                An derived instance of :class:`~transformers.BeamScorer` that defines how beam hypotheses are
-                constructed, stored and sorted during generation. For more information, the documentation of
-                :class:`~transformers.BeamScorer` should be read.
-            logits_processor (:obj:`LogitsProcessorList`, `optional`):
-                An instance of :class:`~transformers.LogitsProcessorList`. List of instances of class derived from
-                :class:`~transformers.LogitsProcessor` used to modify the prediction scores of the language modeling
-                head applied at each generation step.
-            max_length (:obj:`int`, `optional`, defaults to 20):
-                The maximum length of the sequence to be generated.
-            pad_token_id (:obj:`int`, `optional`):
-                The id of the `padding` token.
-            eos_token_id (:obj:`int`, `optional`):
-                The id of the `end-of-sequence` token.
-            model_kwargs:
-                Additional model specific kwargs that will be forwarded to the :obj:`forward` function of the model. If
-                model is an encoder-decoder model the kwargs should include :obj:`encoder_outputs`.
-
-        Return:
-            :obj:`torch.LongTensor` of shape :obj:`(batch_size * num_return_sequences, sequence_length)`: The generated
-            sequences. The second dimension (sequence_length) is either equal to :obj:`max_length` or shorter if all
-            batches finished early due to the :obj:`eos_token_id`.
-
-        Examples::
-
-            >>> from transformers import (
-            ...    AutoTokenizer,
-            ...    AutoModelForSeq2SeqLM,
-            ...    LogitsProcessorList,
-            ...    MinLengthLogitsProcessor,
-            ...    HammingDiversityLogitsProcessor,
-            ...    BeamSearchScorer,
-            ... )
-            >>> import torch
-
-            >>> tokenizer = AutoTokenizer.from_pretrained("t5-base")
-            >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
-
-            >>> encoder_input_str = "translate English to German: How old are you?"
-            >>> encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids
-
-
-            >>> # lets run diverse beam search using 6 beams
-            >>> num_beams = 6
-            >>> # define decoder start token ids
-            >>> input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long)
-            >>> input_ids = input_ids * model.config.decoder_start_token_id
-
-            >>> # add encoder_outputs to model keyword arguments
-            >>> model_kwargs = {
-            ...     "encoder_outputs": model.get_encoder()(encoder_input_ids.repeat_interleave(num_beams, dim=0), return_dict=True)
-            ... }
-
-            >>> # instantiate beam scorer
-            >>> beam_scorer = BeamSearchScorer(
-            ...     batch_size=1,
-            ...     max_length=model.config.max_length,
-            ...     num_beams=num_beams,
-            ...     device=model.device,
-            ...     num_beam_groups=3
-            ... )
-
-            >>> # instantiate logits processors
-            >>> logits_processor = LogitsProcessorList([
-            ...     HammingDiversityLogitsProcessor(5.5, num_beams=6, num_beam_groups=3),
-            ...     MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id),
-            ... ])
-
-            >>> outputs = model.group_beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs)
-
-            >>> print("Generated:", tokenizer.batch_decode(outputs, skip_special_tokens=True))
-        """
-
-        # init values
-        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
-        max_length = max_length if max_length is not None else self.config.max_length
-        pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
-        eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
-
-        batch_size = len(beam_scorer._beam_hyps)
-        num_beams = beam_scorer.num_beams
-        num_beam_groups = beam_scorer.num_beam_groups
-        num_sub_beams = num_beams // num_beam_groups
-        device = input_ids.device
-
-        batch_beam_size, cur_len = input_ids.shape
-
-        assert (
-            num_beams * batch_size == batch_beam_size
-        ), f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
-
-        beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
-        # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in
-        # the same group don't produce same tokens everytime.
-        beam_scores[:, ::num_sub_beams] = 0
-        beam_scores = beam_scores.view((batch_size * num_beams,))
-
-        while cur_len < max_length:
-            # predicted tokens in cur_len step
-            current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device)
-
-            # indices which will form the beams in the next time step
-            reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device)
-
-            # do one decoder step on all beams of all sentences in batch
-            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
-            outputs = self(**model_inputs, return_dict=True)
-
-            for beam_group_idx in range(num_beam_groups):
-                group_start_idx = beam_group_idx * num_sub_beams
-                group_end_idx = min(group_start_idx + num_sub_beams, num_beams)
-                group_size = group_end_idx - group_start_idx
-
-                # indices of beams of current group among all sentences in batch
-                batch_group_indices = []
-                for batch_idx in range(batch_size):
-                    batch_group_indices.extend(
-                        [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)]
-                    )
-                group_input_ids = input_ids[batch_group_indices]
-
-                # select outputs of beams of current group only
-                next_token_logits = outputs.logits[batch_group_indices, -1, :]
-
-                # adjust tokens for Bart, *e.g.*
-                next_token_logits = self.adjust_logits_during_generation(
-                    next_token_logits, cur_len=cur_len, max_length=max_length
-                )
-
-                next_token_scores = F.log_softmax(next_token_logits, dim=-1)  # (batch_size * group_size, vocab_size)
-                vocab_size = next_token_scores.shape[-1]
-
-                next_token_scores = logits_processor(
-                    group_input_ids, next_token_scores, current_tokens=current_tokens, beam_group_idx=beam_group_idx
-                )
-                next_token_scores = next_token_scores + beam_scores[batch_group_indices].unsqueeze(-1).expand_as(
-                    next_token_scores
-                )
-                # reshape for beam search
-
-                next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
-
-                next_token_scores, next_tokens = torch.topk(
-                    next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True
-                )
-
-                next_indices = next_tokens // vocab_size
-                next_tokens = next_tokens % vocab_size
-
-                # stateless
-                beam_outputs = beam_scorer.process(
-                    group_input_ids,
-                    next_token_scores,
-                    next_tokens,
-                    next_indices,
-                    pad_token_id=pad_token_id,
-                    eos_token_id=eos_token_id,
-                )
-                beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"]
-                beam_next_tokens = beam_outputs["next_beam_tokens"]
-                beam_idx = beam_outputs["next_beam_indices"]
-
-                input_ids[batch_group_indices] = group_input_ids[beam_idx]
-                group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
-                current_tokens[batch_group_indices] = group_input_ids[:, -1]
-
-                # (beam_idx // group_size) -> batch_idx
-                # (beam_idx % group_size) -> offset of idx inside the group
-                reordering_indices[batch_group_indices] = (
-                    num_beams * (beam_idx // group_size) + group_start_idx + (beam_idx % group_size)
-                )
-
-            model_kwargs = self._update_model_kwargs_for_generation(
-                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
-            )
-            if model_kwargs["past"] is not None:
-                model_kwargs["past"] = self._reorder_cache(model_kwargs["past"], reordering_indices)
-
-            input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1)
-            cur_len = cur_len + 1
-            if beam_scorer.is_done:
-                break
-
-        decoded = beam_scorer.finalize(
-            input_ids, beam_scores, next_tokens, next_indices, pad_token_id=pad_token_id, eos_token_id=eos_token_id
-        )
-
-        return decoded
-
 
 def top_k_top_p_filtering(
     logits: torch.FloatTensor,

src/transformers/modeling_flax_utils.py

@@ -65,12 +65,13 @@ class FlaxPreTrainedModel(ABC):
     base_model_prefix = ""
     model_class = None
 
-    def __init__(
-        self, config: PretrainedConfig, module: nn.Module, params: Dict, seed: int = 0, dtype: jnp.dtype = jnp.float32
-    ):
+    def __init__(self, config: PretrainedConfig, module: nn.Module, params: Dict, seed: int = 0):
         if config is None:
             raise ValueError("config cannot be None")
 
+        if module is None:
+            raise ValueError("module cannot be None")
+
         if params is None:
             raise ValueError("state cannot be None")
 
@@ -81,23 +82,19 @@ class FlaxPreTrainedModel(ABC):
         # Those are public as their type is generic to every derived classes.
         self.key = PRNGKey(seed)
         self.params = params
-        self.dtype = dtype
+        self.model = module
 
     @property
     def config(self) -> PretrainedConfig:
         return self._config
 
-    @property
-    def module(self) -> nn.Module:
-        return self._module
-
     @staticmethod
     @abstractmethod
     def convert_from_pytorch(pt_state: Dict, config: PretrainedConfig) -> Dict:
         raise NotImplementedError()
 
     @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, dtype: jnp.dtype = jnp.float32, *model_args, **kwargs):
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
         r"""
         Instantiate a pretrained Flax model from a pre-trained model configuration.
         """
@@ -130,9 +127,6 @@ class FlaxPreTrainedModel(ABC):
         else:
             model_kwargs = kwargs
 
-        # Add the dtype to model_kwargs
-        model_kwargs["dtype"] = dtype
-
         # Load model
         if pretrained_model_name_or_path is not None:
             if os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):

src/transformers/models/bart/__init__.py

@@ -31,7 +31,6 @@ if is_torch_available():
         BartForQuestionAnswering,
         BartForSequenceClassification,
         BartModel,
-        BartPretrainedModel,
         PretrainedBartModel,
     )
 

src/transformers/models/bart/modeling_tf_bart.py

@@ -577,9 +577,9 @@ class TFBartDecoder(tf.keras.layers.Layer):
             encoder_padding_mask = invert_mask(encoder_padding_mask)
 
         # embed positions
-        positions = self.embed_positions(input_ids, use_cache=(use_cache and decoder_cached_states is not None))
+        positions = self.embed_positions(input_ids, use_cache=use_cache)
 
-        if use_cache and decoder_cached_states is not None:
+        if use_cache:
             input_ids = input_ids[:, -1:]
             positions = positions[:, -1:]
 
@@ -964,7 +964,7 @@ class TFBartModel(TFPretrainedBartModel):
             else self.config.output_hidden_states
         )
 
-        if not use_cache or past_key_values is None:
+        if not inputs["use_cache"]:
             inputs["decoder_input_ids"], decoder_padding_mask, causal_mask = self._prepare_bart_decoder_inputs(
                 inputs["input_ids"],
                 decoder_input_ids=inputs["decoder_input_ids"],
@@ -1154,7 +1154,6 @@ class TFBartForConditionalGeneration(TFPretrainedBartModel):
             assert (
                 decoder_cached_states
             ), f"decoder cached states must be truthy. got {decoder_cached_states} from the 2nd element of past"
-
         assert isinstance(
             encoder_outputs, TFBaseModelOutput
         ), f"encoder_outputs should be a TFBaseModelOutput, Instead got {type(encoder_outputs)}."

src/transformers/models/bert/__init__.py

@@ -59,4 +59,4 @@ if is_tf_available():
     )
 
 if is_flax_available():
-    from .modeling_flax_bert import FlaxBertForMaskedLM, FlaxBertModel
+    from .modeling_flax_bert import FlaxBertModel

src/transformers/models/bert/modeling_flax_bert.py

@@ -13,14 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Dict, Tuple
+from typing import Callable, Dict
 
 import numpy as np
 
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from jax.random import PRNGKey
 
 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
 from ...modeling_flax_utils import FlaxPreTrainedModel, gelu
@@ -102,8 +101,8 @@ class FlaxBertLayerNorm(nn.Module):
     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.
-    scale_init: Callable[..., np.ndarray] = jax.nn.initializers.ones
-    bias_init: Callable[..., np.ndarray] = jax.nn.initializers.zeros
+    bias_init: jnp.ndarray = nn.initializers.zeros
+    scale_init: jnp.ndarray = nn.initializers.ones
 
     @nn.compact
     def __call__(self, x):
@@ -123,13 +122,11 @@ class FlaxBertLayerNorm(nn.Module):
         mean2 = jnp.mean(jax.lax.square(x), axis=-1, keepdims=True)
         var = mean2 - jax.lax.square(mean)
         mul = jax.lax.rsqrt(var + self.epsilon)
-
         if self.scale:
-            mul = mul * jnp.asarray(self.param("gamma", self.scale_init, (features,)))
+            mul = mul * jnp.asarray(self.param("gamma", self.scale_init, (features,)), self.dtype)
         y = (x - mean) * mul
-
         if self.bias:
-            y = y + jnp.asarray(self.param("beta", self.bias_init, (features,)))
+            y = y + jnp.asarray(self.param("beta", self.bias_init, (features,)), self.dtype)
         return y
 
 
@@ -141,9 +138,7 @@ class FlaxBertEmbedding(nn.Module):
 
     vocab_size: int
     hidden_size: int
-    kernel_init_scale: float = 0.2
-    emb_init: Callable[..., np.ndarray] = jax.nn.initializers.normal(stddev=kernel_init_scale)
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1)
 
     @nn.compact
     def __call__(self, inputs):
@@ -158,105 +153,63 @@ class FlaxBertEmbeddings(nn.Module):
     hidden_size: int
     type_vocab_size: int
     max_length: int
-    kernel_init_scale: float = 0.2
-    dropout_rate: float = 0.0
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask, deterministic: bool = True):
+    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
 
         # Embed
-        w_emb = FlaxBertEmbedding(
-            self.vocab_size,
-            self.hidden_size,
-            kernel_init_scale=self.kernel_init_scale,
-            name="word_embeddings",
-            dtype=self.dtype,
-        )(jnp.atleast_2d(input_ids.astype("i4")))
-        p_emb = FlaxBertEmbedding(
-            self.max_length,
-            self.hidden_size,
-            kernel_init_scale=self.kernel_init_scale,
-            name="position_embeddings",
-            dtype=self.dtype,
-        )(jnp.atleast_2d(position_ids.astype("i4")))
-        t_emb = FlaxBertEmbedding(
-            self.type_vocab_size,
-            self.hidden_size,
-            kernel_init_scale=self.kernel_init_scale,
-            name="token_type_embeddings",
-            dtype=self.dtype,
-        )(jnp.atleast_2d(token_type_ids.astype("i4")))
+        w_emb = FlaxBertEmbedding(self.vocab_size, self.hidden_size, name="word_embeddings")(
+            jnp.atleast_2d(input_ids.astype("i4"))
+        )
+        p_emb = FlaxBertEmbedding(self.max_length, self.hidden_size, name="position_embeddings")(
+            jnp.atleast_2d(position_ids.astype("i4"))
+        )
+        t_emb = FlaxBertEmbedding(self.type_vocab_size, self.hidden_size, name="token_type_embeddings")(
+            jnp.atleast_2d(token_type_ids.astype("i4"))
+        )
 
         # Sum all embeddings
         summed_emb = w_emb + jnp.broadcast_to(p_emb, w_emb.shape) + t_emb
 
         # Layer Norm
-        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(summed_emb)
-        embeddings = nn.Dropout(rate=self.dropout_rate)(layer_norm, deterministic=deterministic)
-        return embeddings
+        layer_norm = FlaxBertLayerNorm(name="layer_norm")(summed_emb)
+
+        return layer_norm
 
 
 class FlaxBertAttention(nn.Module):
     num_heads: int
     head_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, attention_mask):
         # Attention mask comes in as attention_mask.shape == (*batch_sizes, kv_length)
         # FLAX expects: attention_mask.shape == (*batch_sizes, 1, 1, kv_length) such that it is broadcastable
         # with attn_weights.shape == (*batch_sizes, num_heads, q_length, kv_length)
         attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))
-        self_att = nn.attention.SelfAttention(
-            num_heads=self.num_heads,
-            qkv_features=self.head_size,
-            dropout_rate=self.dropout_rate,
-            deterministic=deterministic,
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            bias_init=jax.nn.initializers.zeros,
-            name="self",
-            dtype=self.dtype,
-        )(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
+        )
 
-        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_state)
+        layer_norm = FlaxBertLayerNorm(name="layer_norm")(self_att + hidden_state)
         return layer_norm
 
 
 class FlaxBertIntermediate(nn.Module):
     output_size: int
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
     def __call__(self, hidden_state):
         # TODO: Add ACT2FN reference to change activation function
-        dense = nn.Dense(
-            features=self.output_size,
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            name="dense",
-            dtype=self.dtype,
-        )(hidden_state)
+        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)
         return gelu(dense)
 
 
 class FlaxBertOutput(nn.Module):
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
-
     @nn.compact
-    def __call__(self, intermediate_output, attention_output, deterministic: bool = True):
-        hidden_state = nn.Dense(
-            attention_output.shape[-1],
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            name="dense",
-            dtype=self.dtype,
-        )(intermediate_output)
-        hidden_state = nn.Dropout(rate=self.dropout_rate)(hidden_state, deterministic=deterministic)
-        hidden_state = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_state + attention_output)
+    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)
         return hidden_state
 
 
@@ -264,26 +217,12 @@ class FlaxBertLayer(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
-        attention = FlaxBertAttention(
-            self.num_heads,
-            self.head_size,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="attention",
-            dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
-        intermediate = FlaxBertIntermediate(
-            self.intermediate_size, kernel_init_scale=self.kernel_init_scale, name="intermediate", dtype=self.dtype
-        )(attention)
-        output = FlaxBertOutput(
-            kernel_init_scale=self.kernel_init_scale, dropout_rate=self.dropout_rate, name="output", dtype=self.dtype
-        )(intermediate, attention, deterministic=deterministic)
+    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)
+        output = FlaxBertOutput(name="output")(intermediate, attention)
 
         return output
 
@@ -297,12 +236,9 @@ class FlaxBertLayerCollection(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, inputs, attention_mask, deterministic: bool = True):
+    def __call__(self, inputs, attention_mask):
         assert self.num_layers > 0, f"num_layers should be >= 1, got ({self.num_layers})"
 
         # Initialize input / output
@@ -310,16 +246,8 @@ class FlaxBertLayerCollection(nn.Module):
 
         # Forward over all encoders
         for i in range(self.num_layers):
-            layer = FlaxBertLayer(
-                self.num_heads,
-                self.head_size,
-                self.intermediate_size,
-                kernel_init_scale=self.kernel_init_scale,
-                dropout_rate=self.dropout_rate,
-                name=f"{i}",
-                dtype=self.dtype,
-            )
-            input_i = layer(input_i, attention_mask, deterministic=deterministic)
+            layer = FlaxBertLayer(self.num_heads, self.head_size, self.intermediate_size, name=f"{i}")
+            input_i = layer(input_i, attention_mask)
         return input_i
 
 
@@ -328,39 +256,21 @@ class FlaxBertEncoder(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, attention_mask):
         layer = FlaxBertLayerCollection(
-            self.num_layers,
-            self.num_heads,
-            self.head_size,
-            self.intermediate_size,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="layer",
-            dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+            self.num_layers, self.num_heads, self.head_size, self.intermediate_size, name="layer"
+        )(hidden_state, attention_mask)
         return layer
 
 
 class FlaxBertPooler(nn.Module):
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
-
     @nn.compact
     def __call__(self, hidden_state):
         cls_token = hidden_state[:, 0]
-        out = nn.Dense(
-            hidden_state.shape[-1],
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            name="dense",
-            dtype=self.dtype,
-        )(cls_token)
-        return nn.tanh(out)
+        out = nn.Dense(hidden_state.shape[-1], name="dense")(cls_token)
+        return jax.lax.tanh(out)
 
 
 class FlaxBertModule(nn.Module):
@@ -372,104 +282,24 @@ class FlaxBertModule(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
+    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids):
 
         # Embedding
         embeddings = FlaxBertEmbeddings(
-            self.vocab_size,
-            self.hidden_size,
-            self.type_vocab_size,
-            self.max_length,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="embeddings",
-            dtype=self.dtype,
-        )(input_ids, token_type_ids, position_ids, attention_mask, deterministic=deterministic)
+            self.vocab_size, self.hidden_size, self.type_vocab_size, self.max_length, name="embeddings"
+        )(input_ids, token_type_ids, position_ids, attention_mask)
 
         # N stacked encoding layers
         encoder = FlaxBertEncoder(
-            self.num_encoder_layers,
-            self.num_heads,
-            self.head_size,
-            self.intermediate_size,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="encoder",
-            dtype=self.dtype,
-        )(embeddings, attention_mask, deterministic=deterministic)
+            self.num_encoder_layers, self.num_heads, self.head_size, self.intermediate_size, name="encoder"
+        )(embeddings, attention_mask)
 
-        pooled = FlaxBertPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
+        pooled = FlaxBertPooler(name="pooler")(encoder)
         return encoder, pooled
 
 
-class FlaxBertPredictionHeadTransform(nn.Module):
-    dtype: jnp.dtype = jnp.float32
-
-    @nn.compact
-    def __call__(self, hidden_states):
-        hidden_states = nn.Dense(hidden_states.shape[-1], name="dense", dtype=self.dtype)(hidden_states)
-        hidden_states = nn.elu(hidden_states)  # TODO: ACT2FN[config.hidden_act]
-        return FlaxBertLayerNorm(name="LayerNorm", dtype=self.dtype)(hidden_states)
-
-
-class FlaxBertLMPredictionHead(nn.Module):
-    vocab_size: int
-    dtype: jnp.dtype = jnp.float32
-
-    @nn.compact
-    def __call__(self, hidden_states):
-        # TODO: The output weights are the same as the input embeddings, but there is
-        #   an output-only bias for each token.
-        #   Need a link between the two variables so that the bias is correctly
-        #   resized with `resize_token_embeddings`
-
-        hidden_states = FlaxBertPredictionHeadTransform(name="transform", dtype=self.dtype)(hidden_states)
-        hidden_states = nn.Dense(self.vocab_size, name="decoder", dtype=self.dtype)(hidden_states)
-        return hidden_states
-
-
-class FlaxBertOnlyMLMHead(nn.Module):
-    vocab_size: int
-    hidden_size: int
-    intermediate_size: int
-    head_size: int
-    num_heads: int
-    num_encoder_layers: int
-    type_vocab_size: int
-    max_length: int
-    dropout_rate: float = 0.0
-    dtype: jnp.dtype = jnp.float32
-
-    @nn.compact
-    def __call__(
-        self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, deterministic: bool = True
-    ):
-        # Model
-        encoder, pooled = FlaxBertModule(
-            vocab_size=self.vocab_size,
-            type_vocab_size=self.type_vocab_size,
-            hidden_size=self.hidden_size,
-            intermediate_size=self.intermediate_size,
-            head_size=self.hidden_size,
-            num_heads=self.num_heads,
-            num_encoder_layers=self.num_encoder_layers,
-            max_length=self.max_length,
-            dropout_rate=self.dropout_rate,
-            dtype=self.dtype,
-        )(input_ids, attention_mask, token_type_ids, position_ids, deterministic=deterministic)
-
-        # Compute the prediction scores
-        encoder = nn.Dropout(rate=self.dropout_rate)(encoder, deterministic=deterministic)
-        logits = FlaxBertLMPredictionHead(vocab_size=self.vocab_size, name="predictions", dtype=self.dtype)(encoder)
-
-        return logits, pooled
-
-
 @add_start_docstrings(
     "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
     BERT_START_DOCSTRING,
@@ -555,8 +385,8 @@ class FlaxBertModel(FlaxPreTrainedModel):
 
         return jax_state
 
-    def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
-        module = FlaxBertModule(
+    def __init__(self, config: BertConfig, state: dict, seed: int = 0, **kwargs):
+        model = FlaxBertModule(
             vocab_size=config.vocab_size,
             hidden_size=config.hidden_size,
             type_vocab_size=config.type_vocab_size,
@@ -565,43 +395,16 @@ class FlaxBertModel(FlaxPreTrainedModel):
             num_heads=config.num_attention_heads,
             head_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
-            dropout_rate=config.hidden_dropout_prob,
-            dtype=dtype,
         )
 
-        super().__init__(config, module, state, seed)
+        super().__init__(config, model, state, seed)
+
+    @property
+    def module(self) -> nn.Module:
+        return self._module
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
-    def __call__(
-        self,
-        input_ids,
-        attention_mask=None,
-        token_type_ids=None,
-        position_ids=None,
-        params: dict = None,
-        dropout_rng: PRNGKey = None,
-        train: bool = False,
-    ):
-        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
-            input_ids, attention_mask, token_type_ids, position_ids
-        )
-
-        # Handle any PRNG if needed
-        rngs = {}
-        if dropout_rng is not None:
-            rngs["dropout"] = dropout_rng
-
-        return self.module.apply(
-            {"params": params or self.params},
-            jnp.array(input_ids, dtype="i4"),
-            jnp.array(attention_mask, dtype="i4"),
-            jnp.array(token_type_ids, dtype="i4"),
-            jnp.array(position_ids, dtype="i4"),
-            not train,
-            rngs=rngs,
-        )
-
-    def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
+    def __call__(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None):
         if token_type_ids is None:
             token_type_ids = jnp.ones_like(input_ids)
 
@@ -611,62 +414,10 @@ class FlaxBertModel(FlaxPreTrainedModel):
         if attention_mask is None:
             attention_mask = jnp.ones_like(input_ids)
 
-        return input_ids, attention_mask, token_type_ids, position_ids
-
-    def init(self, rng: jax.random.PRNGKey, input_shape: Tuple):
-        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
-            jnp.zeros(input_shape, dtype="i4"), None, None, None
-        )
-
-        params_rng, dropout_rng = jax.random.split(rng)
-        rngs = {"params": params_rng, "dropout": dropout_rng}
-
-        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
-
-
-class FlaxBertForMaskedLM(FlaxBertModel):
-    def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs):
-        super().__init__(config, state, seed, dtype)
-
-        self._module = FlaxBertOnlyMLMHead(
-            vocab_size=config.vocab_size,
-            type_vocab_size=config.type_vocab_size,
-            hidden_size=config.hidden_size,
-            intermediate_size=config.intermediate_size,
-            head_size=config.hidden_size,
-            num_heads=config.num_attention_heads,
-            num_encoder_layers=config.num_hidden_layers,
-            max_length=config.max_length,
-            **kwargs,
-        )
-
-    def __call__(
-        self,
-        input_ids,
-        attention_mask=None,
-        token_type_ids=None,
-        position_ids=None,
-        params: dict = None,
-        dropout_rng: PRNGKey = None,
-        train: bool = False,
-    ):
-        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
-            input_ids, attention_mask, token_type_ids, position_ids
-        )
-
-        # Handle any PRNG if needed
-        rngs = {}
-        if dropout_rng is not None:
-            rngs["dropout"] = dropout_rng
-
-        pooled, logits = self.module.apply(
-            {"params": params or self.params},
+        return self.model.apply(
+            {"params": self.params},
             jnp.array(input_ids, dtype="i4"),
             jnp.array(attention_mask, dtype="i4"),
             jnp.array(token_type_ids, dtype="i4"),
             jnp.array(position_ids, dtype="i4"),
-            not train,
-            rngs=rngs,
         )
-
-        return logits, pooled

src/transformers/models/rag/modeling_rag.py

@@ -1226,8 +1226,6 @@ class RagTokenForGeneration(RagPreTrainedModel):
         early_stopping=None,
         use_cache=None,
         num_beams=None,
-        num_beam_groups=None,
-        diversity_penalty=None,
         bos_token_id=None,
         pad_token_id=None,
         eos_token_id=None,
@@ -1304,13 +1302,6 @@ class RagTokenForGeneration(RagPreTrainedModel):
                 should not appear in the generated text, use :obj:`tokenizer.encode(bad_word, add_prefix_space=True)`.
             num_beams (:obj:`int`, `optional`, defaults to 1):
                 Number of beams for beam search. 1 means no beam search.
-            num_beam_groups (:obj:`int`, `optional`, defaults to 1):
-                Number of groups to divide :obj:`num_beams` into in order to ensure diversity among different groups of
-                beams. `this paper <https://arxiv.org/pdf/1610.02424.pdf>`__ for more details.
-            diversity_penalty (:obj:`float`, `optional`, defaults to 0.0):
-                This value is subtracted from a beam's score if it generates a token same as any beam from other group
-                at a particular time. Note that :obj:`diversity_penalty` is only effective if ``group beam search`` is
-                enabled.
             num_return_sequences(:obj:`int`, `optional`, defaults to 1):
                 The number of independently computed returned sequences for each element in the batch. Note that this
                 is not the value we pass to the ``generator``'s `:func:`~transformers.PreTrainedModel.generate`
@@ -1335,7 +1326,6 @@ class RagTokenForGeneration(RagPreTrainedModel):
         # set default parameters
         n_docs = n_docs if n_docs is not None else self.config.n_docs
         num_beams = num_beams if num_beams is not None else self.config.num_beams
-        num_beam_groups = num_beam_groups if num_beam_groups is not None else self.config.num_beam_groups
         max_length = max_length if max_length is not None else self.config.max_length
         num_return_sequences = (
             num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
@@ -1422,8 +1412,6 @@ class RagTokenForGeneration(RagPreTrainedModel):
             eos_token_id=eos_token_id,
             prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
             num_beams=num_beams,
-            num_beam_groups=num_beam_groups,
-            diversity_penalty=diversity_penalty,
         )
 
         if num_beams == 1:

src/transformers/models/roberta/modeling_flax_roberta.py

@@ -12,14 +12,13 @@
 # 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.
-from typing import Callable, Dict, Tuple
+from typing import Callable, Dict
 
 import numpy as np
 
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from jax.random import PRNGKey
 
 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
 from ...modeling_flax_utils import FlaxPreTrainedModel, gelu
@@ -102,8 +101,8 @@ class FlaxRobertaLayerNorm(nn.Module):
     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.
-    scale_init: Callable[..., np.ndarray] = jax.nn.initializers.ones
-    bias_init: Callable[..., np.ndarray] = jax.nn.initializers.zeros
+    bias_init: jnp.ndarray = nn.initializers.zeros
+    scale_init: jnp.ndarray = nn.initializers.ones
 
     @nn.compact
     def __call__(self, x):
@@ -123,13 +122,11 @@ class FlaxRobertaLayerNorm(nn.Module):
         mean2 = jnp.mean(jax.lax.square(x), axis=-1, keepdims=True)
         var = mean2 - jax.lax.square(mean)
         mul = jax.lax.rsqrt(var + self.epsilon)
-
         if self.scale:
-            mul = mul * jnp.asarray(self.param("gamma", self.scale_init, (features,)))
+            mul = mul * jnp.asarray(self.param("gamma", self.scale_init, (features,)), self.dtype)
         y = (x - mean) * mul
-
         if self.bias:
-            y = y + jnp.asarray(self.param("beta", self.bias_init, (features,)))
+            y = y + jnp.asarray(self.param("beta", self.bias_init, (features,)), self.dtype)
         return y
 
 
@@ -142,9 +139,7 @@ class FlaxRobertaEmbedding(nn.Module):
 
     vocab_size: int
     hidden_size: int
-    kernel_init_scale: float = 0.2
-    emb_init: Callable[..., np.ndarray] = jax.nn.initializers.normal(stddev=kernel_init_scale)
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1)
 
     @nn.compact
     def __call__(self, inputs):
@@ -160,108 +155,66 @@ class FlaxRobertaEmbeddings(nn.Module):
     hidden_size: int
     type_vocab_size: int
     max_length: int
-    kernel_init_scale: float = 0.2
-    dropout_rate: float = 0.0
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask, deterministic: bool = True):
+    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
 
         # Embed
-        w_emb = FlaxRobertaEmbedding(
-            self.vocab_size,
-            self.hidden_size,
-            kernel_init_scale=self.kernel_init_scale,
-            name="word_embeddings",
-            dtype=self.dtype,
-        )(jnp.atleast_2d(input_ids.astype("i4")))
-        p_emb = FlaxRobertaEmbedding(
-            self.max_length,
-            self.hidden_size,
-            kernel_init_scale=self.kernel_init_scale,
-            name="position_embeddings",
-            dtype=self.dtype,
-        )(jnp.atleast_2d(position_ids.astype("i4")))
-        t_emb = FlaxRobertaEmbedding(
-            self.type_vocab_size,
-            self.hidden_size,
-            kernel_init_scale=self.kernel_init_scale,
-            name="token_type_embeddings",
-            dtype=self.dtype,
-        )(jnp.atleast_2d(token_type_ids.astype("i4")))
+        w_emb = FlaxRobertaEmbedding(self.vocab_size, self.hidden_size, name="word_embeddings")(
+            jnp.atleast_2d(input_ids.astype("i4"))
+        )
+        p_emb = FlaxRobertaEmbedding(self.max_length, self.hidden_size, name="position_embeddings")(
+            jnp.atleast_2d(position_ids.astype("i4"))
+        )
+        t_emb = FlaxRobertaEmbedding(self.type_vocab_size, self.hidden_size, name="token_type_embeddings")(
+            jnp.atleast_2d(token_type_ids.astype("i4"))
+        )
 
         # Sum all embeddings
         summed_emb = w_emb + jnp.broadcast_to(p_emb, w_emb.shape) + t_emb
 
         # Layer Norm
-        layer_norm = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(summed_emb)
-        embeddings = nn.Dropout(rate=self.dropout_rate)(layer_norm, deterministic=deterministic)
-        return embeddings
+        layer_norm = FlaxRobertaLayerNorm(name="layer_norm")(summed_emb)
+
+        return layer_norm
 
 
 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertAttention with Bert->Roberta
 class FlaxRobertaAttention(nn.Module):
     num_heads: int
     head_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, attention_mask):
         # Attention mask comes in as attention_mask.shape == (*batch_sizes, kv_length)
         # FLAX expects: attention_mask.shape == (*batch_sizes, 1, 1, kv_length) such that it is broadcastable
         # with attn_weights.shape == (*batch_sizes, num_heads, q_length, kv_length)
         attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))
-        self_att = nn.attention.SelfAttention(
-            num_heads=self.num_heads,
-            qkv_features=self.head_size,
-            dropout_rate=self.dropout_rate,
-            deterministic=deterministic,
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            bias_init=jax.nn.initializers.zeros,
-            name="self",
-            dtype=self.dtype,
-        )(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
+        )
 
-        layer_norm = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_state)
+        layer_norm = FlaxRobertaLayerNorm(name="layer_norm")(self_att + hidden_state)
         return layer_norm
 
 
 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertIntermediate with Bert->Roberta
 class FlaxRobertaIntermediate(nn.Module):
     output_size: int
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
     def __call__(self, hidden_state):
         # TODO: Add ACT2FN reference to change activation function
-        dense = nn.Dense(
-            features=self.output_size,
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            name="dense",
-            dtype=self.dtype,
-        )(hidden_state)
+        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)
         return gelu(dense)
 
 
 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertOutput with Bert->Roberta
 class FlaxRobertaOutput(nn.Module):
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
-
     @nn.compact
-    def __call__(self, intermediate_output, attention_output, deterministic: bool = True):
-        hidden_state = nn.Dense(
-            attention_output.shape[-1],
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            name="dense",
-            dtype=self.dtype,
-        )(intermediate_output)
-        hidden_state = nn.Dropout(rate=self.dropout_rate)(hidden_state, deterministic=deterministic)
-        hidden_state = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_state + attention_output)
+    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)
         return hidden_state
 
 
@@ -269,29 +222,14 @@ class FlaxRobertaLayer(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
-        attention = FlaxRobertaAttention(
-            self.num_heads,
-            self.head_size,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="attention",
-            dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
-        intermediate = FlaxRobertaIntermediate(
-            self.intermediate_size,
-            kernel_init_scale=self.kernel_init_scale,
-            name="intermediate",
-            dtype=self.dtype,
-        )(attention)
-        output = FlaxRobertaOutput(
-            kernel_init_scale=self.kernel_init_scale, dropout_rate=self.dropout_rate, name="output", dtype=self.dtype
-        )(intermediate, attention, deterministic=deterministic)
+    def __call__(self, hidden_state, attention_mask):
+        attention = FlaxRobertaAttention(self.num_heads, self.head_size, name="attention")(
+            hidden_state, attention_mask
+        )
+        intermediate = FlaxRobertaIntermediate(self.intermediate_size, name="intermediate")(attention)
+        output = FlaxRobertaOutput(name="output")(intermediate, attention)
 
         return output
 
@@ -306,12 +244,9 @@ class FlaxRobertaLayerCollection(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, inputs, attention_mask, deterministic: bool = True):
+    def __call__(self, inputs, attention_mask):
         assert self.num_layers > 0, f"num_layers should be >= 1, got ({self.num_layers})"
 
         # Initialize input / output
@@ -319,16 +254,8 @@ class FlaxRobertaLayerCollection(nn.Module):
 
         # Forward over all encoders
         for i in range(self.num_layers):
-            layer = FlaxRobertaLayer(
-                self.num_heads,
-                self.head_size,
-                self.intermediate_size,
-                kernel_init_scale=self.kernel_init_scale,
-                dropout_rate=self.dropout_rate,
-                name=f"{i}",
-                dtype=self.dtype,
-            )
-            input_i = layer(input_i, attention_mask, deterministic=deterministic)
+            layer = FlaxRobertaLayer(self.num_heads, self.head_size, self.intermediate_size, name=f"{i}")
+            input_i = layer(input_i, attention_mask)
         return input_i
 
 
@@ -338,40 +265,22 @@ class FlaxRobertaEncoder(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, attention_mask):
         layer = FlaxRobertaLayerCollection(
-            self.num_layers,
-            self.num_heads,
-            self.head_size,
-            self.intermediate_size,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="layer",
-            dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+            self.num_layers, self.num_heads, self.head_size, self.intermediate_size, name="layer"
+        )(hidden_state, attention_mask)
         return layer
 
 
 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertPooler with Bert->Roberta
 class FlaxRobertaPooler(nn.Module):
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
-
     @nn.compact
     def __call__(self, hidden_state):
         cls_token = hidden_state[:, 0]
-        out = nn.Dense(
-            hidden_state.shape[-1],
-            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
-            name="dense",
-            dtype=self.dtype,
-        )(cls_token)
-        return nn.tanh(out)
+        out = nn.Dense(hidden_state.shape[-1], name="dense")(cls_token)
+        return jax.lax.tanh(out)
 
 
 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertModule with Bert->Roberta
@@ -384,38 +293,21 @@ class FlaxRobertaModule(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
+    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids):
 
         # Embedding
         embeddings = FlaxRobertaEmbeddings(
-            self.vocab_size,
-            self.hidden_size,
-            self.type_vocab_size,
-            self.max_length,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="embeddings",
-            dtype=self.dtype,
-        )(input_ids, token_type_ids, position_ids, attention_mask, deterministic=deterministic)
+            self.vocab_size, self.hidden_size, self.type_vocab_size, self.max_length, name="embeddings"
+        )(input_ids, token_type_ids, position_ids, attention_mask)
 
         # N stacked encoding layers
         encoder = FlaxRobertaEncoder(
-            self.num_encoder_layers,
-            self.num_heads,
-            self.head_size,
-            self.intermediate_size,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="encoder",
-            dtype=self.dtype,
-        )(embeddings, attention_mask, deterministic=deterministic)
+            self.num_encoder_layers, self.num_heads, self.head_size, self.intermediate_size, name="encoder"
+        )(embeddings, attention_mask)
 
-        pooled = FlaxRobertaPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
+        pooled = FlaxRobertaPooler(name="pooler")(encoder)
         return encoder, pooled
 
 
@@ -504,8 +396,8 @@ class FlaxRobertaModel(FlaxPreTrainedModel):
 
         return jax_state
 
-    def __init__(self, config: RobertaConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
-        module = FlaxRobertaModule(
+    def __init__(self, config: RobertaConfig, state: dict, seed: int = 0, **kwargs):
+        model = FlaxRobertaModule(
             vocab_size=config.vocab_size,
             hidden_size=config.hidden_size,
             type_vocab_size=config.type_vocab_size,
@@ -514,78 +406,31 @@ class FlaxRobertaModel(FlaxPreTrainedModel):
             num_heads=config.num_attention_heads,
             head_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
-            dropout_rate=config.hidden_dropout_prob,
-            dtype=dtype,
         )
 
-        super().__init__(config, module, state, seed)
+        super().__init__(config, model, state, seed)
+
+    @property
+    def module(self) -> nn.Module:
+        return self._module
 
     @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
-    def __call__(
-        self,
-        input_ids,
-        token_type_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        params: dict = None,
-        dropout_rng: PRNGKey = None,
-        train: bool = False,
-    ):
-        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
-            input_ids, attention_mask, token_type_ids, position_ids
-        )
-
-        # Handle any PRNG if needed
-        rngs = {}
-        if dropout_rng is not None:
-            rngs["dropout"] = dropout_rng
-
-        return self.module.apply(
-            {"params": params or self.params},
-            jnp.array(input_ids, dtype="i4"),
-            jnp.array(attention_mask, dtype="i4"),
-            jnp.array(token_type_ids, dtype="i4"),
-            jnp.array(position_ids, dtype="i4"),
-            not train,
-            rngs=rngs,
-        )
-
-    def init(self, rng: jax.random.PRNGKey, input_shape: Tuple):
-        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
-            jnp.zeros(input_shape, dtype="i4"), None, None, None
-        )
-
-        params_rng, dropout_rng = jax.random.split(rng)
-        rngs = {"params": params_rng, "dropout": dropout_rng}
-
-        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
-
-    def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
-
+    def __call__(self, input_ids, token_type_ids=None, attention_mask=None, position_ids=None):
         if token_type_ids is None:
             token_type_ids = jnp.ones_like(input_ids)
 
         if position_ids is None:
-            position_ids = create_position_ids_from_input_ids(input_ids, self.config.pad_token_id)
+            position_ids = jnp.arange(
+                self.config.pad_token_id + 1, jnp.atleast_2d(input_ids).shape[-1] + self.config.pad_token_id + 1
+            )
 
         if attention_mask is None:
             attention_mask = jnp.ones_like(input_ids)
 
-        return input_ids, attention_mask, token_type_ids, position_ids
-
-
-def create_position_ids_from_input_ids(input_ids, padding_idx):
-    """
-    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
-    are ignored. This is modified from fairseq's `utils.make_positions`.
-
-    Args:
-        input_ids: jnp.ndarray
-        padding_idx: int
-
-    Returns: jnp.ndarray
-    """
-    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
-    mask = (input_ids != padding_idx).astype("i4")
-    incremental_indices = jnp.cumsum(mask, axis=1).astype("i4") * mask
-    return incremental_indices.astype("i4") + padding_idx
+        return self.model.apply(
+            {"params": self.params},
+            jnp.array(input_ids, dtype="i4"),
+            jnp.array(attention_mask, dtype="i4"),
+            jnp.array(token_type_ids, dtype="i4"),
+            jnp.array(position_ids, dtype="i4"),
+        )

src/transformers/models/t5/modeling_t5.py

@@ -813,6 +813,9 @@ class T5Stack(T5PreTrainedModel):
     def get_input_embeddings(self):
         return self.embed_tokens
 
+    def get_output_embeddings(self):
+        return self.embed_tokens
+
     def set_input_embeddings(self, new_embeddings):
         self.embed_tokens = new_embeddings
 

src/transformers/trainer.py

@@ -543,7 +543,7 @@ class Trainer:
             self.args.output_dir = checkpoint_dir
             output_dir = os.path.join(self.args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}")
             self.save_model(output_dir)
-            if self.is_world_process_zero():
+            if self.is_world_master():
                 self.state.save_to_json(os.path.join(output_dir, "trainer_state.json"))
                 torch.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                 torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))

src/transformers/training_args.py

@@ -23,12 +23,14 @@ from .file_utils import cached_property, is_torch_available, is_torch_tpu_availa
 from .trainer_utils import EvaluationStrategy
 from .utils import logging
 
+
 if is_torch_available():
     import torch
 
 if is_torch_tpu_available():
     import torch_xla.core.xla_model as xm
 
+
 logger = logging.get_logger(__name__)
 
 
@@ -82,8 +84,6 @@ class TrainingArguments:
                 * :obj:`"steps"`: Evaluation is done (and logged) every :obj:`eval_steps`.
                 * :obj:`"epoch"`: Evaluation is done at the end of each epoch.
 
-        initial_evaluation (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to perform evaluation before doing any training, for example for fine-tuning.
         prediction_loss_only (:obj:`bool`, `optional`, defaults to `False`):
             When performing evaluation and predictions, only returns the loss.
         per_device_train_batch_size (:obj:`int`, `optional`, defaults to 8):
@@ -238,8 +238,6 @@ class TrainingArguments:
         default="no",
         metadata={"help": "Run evaluation during training at each logging step."},
     )
-    initial_evaluation: bool = field(default=False, metadata={
-        "help": "Whether to perform evaluation before doing any training, for example for fine-tuning."})
     prediction_loss_only: bool = field(
         default=False,
         metadata={"help": "When performing evaluation and predictions, only returns the loss."},
@@ -256,14 +254,14 @@ class TrainingArguments:
         default=None,
         metadata={
             "help": "Deprecated, the use of `--per_device_train_batch_size` is preferred. "
-                    "Batch size per GPU/TPU core/CPU for training."
+            "Batch size per GPU/TPU core/CPU for training."
         },
     )
     per_gpu_eval_batch_size: Optional[int] = field(
         default=None,
         metadata={
             "help": "Deprecated, the use of `--per_device_eval_batch_size` is preferred."
-                    "Batch size per GPU/TPU core/CPU for evaluation."
+            "Batch size per GPU/TPU core/CPU for evaluation."
         },
     )
 
@@ -376,12 +374,6 @@ class TrainingArguments:
             "help": "When resuming training, whether or not to skip the first epochs and batches to get to the same training data."
         },
     )
-    tokenizers_parallelism: bool = field(
-        default=True,
-        metadata={
-            "help": "Sets the TOKENIZERS_PARALLELISM variable to True or False to prevent warnings"
-        },
-    )
 
     def __post_init__(self):
         if self.disable_tqdm is None:

src/transformers/utils/dummy_flax_objects.py

@@ -14,15 +14,6 @@ class FlaxAutoModel:
         requires_flax(self)
 
 
-class FlaxBertForMaskedLM:
-    def __init__(self, *args, **kwargs):
-        requires_flax(self)
-
-    @classmethod
-    def from_pretrained(self, *args, **kwargs):
-        requires_flax(self)
-
-
 class FlaxBertModel:
     def __init__(self, *args, **kwargs):
         requires_flax(self)

src/transformers/utils/dummy_pt_objects.py

@@ -118,11 +118,6 @@ class BeamSearchScorer:
         requires_pytorch(self)
 
 
-class HammingDiversityLogitsProcessor:
-    def __init__(self, *args, **kwargs):
-        requires_pytorch(self)
-
-
 class LogitsProcessor:
     def __init__(self, *args, **kwargs):
         requires_pytorch(self)
@@ -153,11 +148,6 @@ class NoRepeatNGramLogitsProcessor:
         requires_pytorch(self)
 
 
-class PrefixConstrainedLogitsProcessor:
-    def __init__(self, *args, **kwargs):
-        requires_pytorch(self)
-
-
 class RepetitionPenaltyLogitsProcessor:
     def __init__(self, *args, **kwargs):
         requires_pytorch(self)
@@ -450,15 +440,6 @@ class BartModel:
         requires_pytorch(self)
 
 
-class BartPretrainedModel:
-    def __init__(self, *args, **kwargs):
-        requires_pytorch(self)
-
-    @classmethod
-    def from_pretrained(self, *args, **kwargs):
-        requires_pytorch(self)
-
-
 class PretrainedBartModel:
     def __init__(self, *args, **kwargs):
         requires_pytorch(self)

tests/test_generation_logits_process.py

@@ -27,7 +27,6 @@ if is_torch_available():
     import torch.nn.functional as F
 
     from transformers.generation_logits_process import (
-        HammingDiversityLogitsProcessor,
         LogitsProcessorList,
         MinLengthLogitsProcessor,
         NoBadWordsLogitsProcessor,
@@ -303,30 +302,3 @@ class LogitsProcessorTest(unittest.TestCase):
         self.assertListEqual(
             torch.isinf(filtered_scores).tolist(), [[False, False, True, True, True], [True, True, False, False, True]]
         )
-
-    def test_hamming_diversity(self):
-        vocab_size = 4
-        num_beams = 2
-        num_beam_groups = 2
-
-        scores = self._get_uniform_logits(num_beams, vocab_size)
-        # batch_idx = 0 -> index batch_idx * num_beam_groups -> idx = 0 * 2 = 0 -> penalises tokens 1
-        # batch_idx = 1 -> index batch_idx * num_beam_groups -> idx = 1 * 2 = 2 -> penalises tokens 1
-        current_tokens = torch.tensor([0, 3, 1, 2], device=torch_device, dtype=torch.long)
-
-        diversity_logits_processor = HammingDiversityLogitsProcessor(
-            diversity_penalty=1.0, num_beams=num_beams, num_beam_groups=num_beam_groups
-        )
-
-        processed_scores = diversity_logits_processor(None, scores, current_tokens, 1)
-
-        self.assertTrue(
-            torch.allclose(
-                processed_scores[0], torch.tensor([-0.7500, 0.2500, 0.2500, 0.2500], device=torch_device), atol=1e-3
-            )
-        )
-        self.assertTrue(
-            torch.allclose(
-                processed_scores[1], torch.tensor([0.2500, -0.7500, 0.2500, 0.2500], device=torch_device), atol=1e-3
-            )
-        )

tests/test_generation_utils.py

@@ -17,16 +17,15 @@
 import unittest
 
 from transformers import is_torch_available
-from transformers.testing_utils import require_torch, slow, torch_device
+from transformers.testing_utils import require_torch, torch_device
 
 
 if is_torch_available():
     import torch
 
-    from transformers import BartForConditionalGeneration, BartTokenizer, top_k_top_p_filtering
+    from transformers import top_k_top_p_filtering
     from transformers.generation_beam_search import BeamSearchScorer
     from transformers.generation_logits_process import (
-        HammingDiversityLogitsProcessor,
         LogitsProcessorList,
         MinLengthLogitsProcessor,
         NoBadWordsLogitsProcessor,
@@ -62,7 +61,7 @@ class GenerationTesterMixin:
         return config, input_ids, attention_mask, max_length
 
     @staticmethod
-    def _get_logits_processor_and_kwargs(input_length, eos_token_id, diversity_penalty=None):
+    def _get_logits_processor_and_kwargs(input_length, eos_token_id):
         process_kwargs = {
             "min_length": input_length + 1,
             "bad_words_ids": [[1, 0]],
@@ -71,13 +70,6 @@ class GenerationTesterMixin:
         }
         logits_processor = LogitsProcessorList(
             (
-                [
-                    HammingDiversityLogitsProcessor(diversity_penalty, num_beams=2, num_beam_groups=2),
-                ]
-                if diversity_penalty is not None
-                else []
-            )
-            + (
                 [
                     MinLengthLogitsProcessor(process_kwargs["min_length"], eos_token_id),
                 ]
@@ -123,28 +115,6 @@ class GenerationTesterMixin:
         )
         return beam_kwargs, beam_scorer
 
-    @staticmethod
-    def _get_diverse_beam_scorer_and_kwargs(batch_size, max_length, num_return_sequences=1):
-        beam_kwargs = {
-            "early_stopping": False,
-            "length_penalty": 2.0,
-            "num_beams": 2,
-            "num_return_sequences": num_return_sequences,
-            "num_beam_groups": 2,  # one beam per group
-            "diversity_penalty": 2.0,
-        }
-        beam_scorer = BeamSearchScorer(
-            batch_size=batch_size,
-            max_length=max_length,
-            num_beams=beam_kwargs["num_beams"],
-            device=torch_device,
-            length_penalty=beam_kwargs["length_penalty"],
-            do_early_stopping=beam_kwargs["early_stopping"],
-            num_beam_hyps_to_keep=num_return_sequences,
-            num_beam_groups=beam_kwargs["num_beam_groups"],
-        )
-        return beam_kwargs, beam_scorer
-
     @staticmethod
     def _get_encoder_outputs(model, input_ids, attention_mask, num_interleave=1):
         encoder = model.get_encoder()
@@ -438,92 +408,6 @@ class GenerationTesterMixin:
 
                 self.assertIsNotNone(output_ids_generate)
 
-    def test_group_beam_search_generate(self):
-        for model_class in self.all_generative_model_classes:
-            config, input_ids, attention_mask, max_length = self._get_input_ids_and_config()
-
-            logits_process_kwargs, logits_processor = self._get_logits_processor_and_kwargs(
-                input_ids.shape[-1], config.eos_token_id, diversity_penalty=2.0
-            )
-
-            model = model_class(config).to(torch_device)
-            model.eval()
-
-            # check `generate()` and `group_beam_search()` are equal
-            if model.config.is_encoder_decoder:
-                max_length = 4
-            beam_kwargs, beam_scorer = self._get_diverse_beam_scorer_and_kwargs(input_ids.shape[0], max_length)
-            output_ids_generate = model.generate(
-                input_ids,
-                attention_mask=attention_mask,
-                do_sample=False,
-                max_length=max_length,
-                **beam_kwargs,
-                **logits_process_kwargs,
-            )
-
-            # group_beam_search does not automatically interleave `batch_size` dim for `num_beams`
-            kwargs = {}
-            if model.config.is_encoder_decoder:
-                encoder_outputs, input_ids_clone, attention_mask_clone = self._get_encoder_outputs(
-                    model, input_ids, attention_mask, num_interleave=beam_scorer.num_beams
-                )
-                kwargs["encoder_outputs"] = encoder_outputs
-                input_ids_clone = input_ids_clone.repeat_interleave(beam_scorer.num_beams, dim=0)
-            else:
-                attention_mask_clone = attention_mask.repeat_interleave(beam_scorer.num_beams, dim=0)
-                input_ids_clone = input_ids.repeat_interleave(beam_scorer.num_beams, dim=0)
-
-            with torch.no_grad():
-                output_ids_group_beam_search = model.group_beam_search(
-                    input_ids_clone,
-                    beam_scorer,
-                    max_length=max_length,
-                    attention_mask=attention_mask_clone,
-                    logits_processor=logits_processor,
-                    **kwargs,
-                )
-            self.assertListEqual(output_ids_generate.tolist(), output_ids_group_beam_search.tolist())
-
-            # check `generate()` and `group_beam_search()` are equal for `num_return_sequences`
-            num_return_sequences = 2
-            if model.config.is_encoder_decoder:
-                max_length = 4
-            beam_kwargs, beam_scorer = self._get_diverse_beam_scorer_and_kwargs(
-                input_ids.shape[0], max_length, num_return_sequences=num_return_sequences
-            )
-
-            output_ids_generate = model.generate(
-                input_ids,
-                attention_mask=attention_mask,
-                do_sample=False,
-                max_length=max_length,
-                **beam_kwargs,
-                **logits_process_kwargs,
-            )
-            # group_beam_search does not automatically interleave `batch_size` dim for `num_beams`
-            kwargs = {}
-            if model.config.is_encoder_decoder:
-                encoder_outputs, input_ids_clone, attention_mask_clone = self._get_encoder_outputs(
-                    model, input_ids, attention_mask, num_interleave=beam_scorer.num_beams
-                )
-                kwargs["encoder_outputs"] = encoder_outputs
-                input_ids_clone = input_ids_clone.repeat_interleave(beam_scorer.num_beams, dim=0)
-            else:
-                attention_mask_clone = attention_mask.repeat_interleave(beam_scorer.num_beams, dim=0)
-                input_ids_clone = input_ids.repeat_interleave(beam_scorer.num_beams, dim=0)
-
-            with torch.no_grad():
-                output_ids_beam_search = model.group_beam_search(
-                    input_ids_clone,
-                    beam_scorer,
-                    max_length=max_length,
-                    attention_mask=attention_mask_clone,
-                    logits_processor=logits_processor,
-                    **kwargs,
-                )
-            self.assertListEqual(output_ids_generate.tolist(), output_ids_beam_search.tolist())
-
 
 @require_torch
 class UtilsFunctionsTest(unittest.TestCase):
@@ -628,31 +512,3 @@ class UtilsFunctionsTest(unittest.TestCase):
 
         self.assertTrue(torch.allclose(non_inf_expected_output, non_inf_output, atol=1e-12))
         self.assertTrue(torch.all(torch.eq(non_inf_expected_idx, non_inf_idx)))
-
-
-@require_torch
-class GenerationIntegrationTests(unittest.TestCase):
-    @slow
-    def test_diverse_beam_search(self):
-        article = """Justin Timberlake and Jessica Biel, welcome to parenthood.
-        The celebrity couple announced the arrival of their son, Silas Randall Timberlake, in statements to People.
-        "Silas was the middle name of Timberlake's maternal grandfather Bill Bomar, who died in 2012, while Randall is the musician's own middle name, as well as his father's first," People reports.
-        The couple announced the pregnancy in January, with an Instagram post. It is the first baby for both."""
-
-        bart_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
-        bart_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn").to(torch_device)
-        input_ids = bart_tokenizer(article, return_tensors="pt").input_ids.to(torch_device)
-
-        outputs = bart_model.generate(
-            input_ids, num_beams=4, num_return_sequences=2, num_beam_groups=4, diversity_penalty=2.0
-        )
-
-        generated_text = bart_tokenizer.batch_decode(outputs, skip_special_tokens=True)
-
-        self.assertListEqual(
-            generated_text,
-            [
-                "The couple announced the birth of their son, Silas Randall Timberlake, in a statement. Silas was the middle name of Timberlake's maternal grandfather Bill Bomar. Randall is the musician's own middle name, as well as his father's first. It is the first baby for both of them.",
-                "Justin Timberlake and Jessica Biel have a son. The baby is named Silas Randall Timberlake. It is the first child for both. The couple announced the pregnancy in January. The name Silas is the middle name of Timberlake's maternal grandfather. It's also his own middle name.",
-            ],
-        )

tests/test_modeling_bart.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 import tempfile
 import unittest
 
@@ -50,68 +49,37 @@ if is_torch_available():
         pipeline,
     )
     from transformers.models.bart.modeling_bart import (
-        BartDecoder,
-        BartEncoder,
-        BartSinusoidalPositionalEmbedding,
+        SinusoidalPositionalEmbedding,
+        _prepare_bart_decoder_inputs,
+        invert_mask,
         shift_tokens_right,
     )
-
-
 PGE_ARTICLE = """ PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow."""
 
 
-def prepare_bart_inputs_dict(
-    config,
-    input_ids,
-    attention_mask=None,
-):
-    if attention_mask is None:
-        attention_mask = input_ids.ne(config.pad_token_id)
-    return {
-        "input_ids": input_ids,
-        "attention_mask": attention_mask,
-    }
-
-
 @require_torch
-class BartModelTester:
+class ModelTester:
     def __init__(
         self,
         parent,
-        batch_size=13,
-        seq_length=7,
-        is_training=True,
-        use_labels=False,
-        vocab_size=99,
-        hidden_size=16,
-        num_hidden_layers=2,
-        num_attention_heads=4,
-        intermediate_size=4,
-        hidden_act="gelu",
-        hidden_dropout_prob=0.1,
-        attention_probs_dropout_prob=0.1,
-        max_position_embeddings=20,
-        eos_token_id=2,
-        pad_token_id=1,
-        bos_token_id=0,
     ):
         self.parent = parent
-        self.batch_size = batch_size
-        self.seq_length = seq_length
-        self.is_training = is_training
-        self.use_labels = use_labels
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.intermediate_size = intermediate_size
-        self.hidden_act = hidden_act
-        self.hidden_dropout_prob = hidden_dropout_prob
-        self.attention_probs_dropout_prob = attention_probs_dropout_prob
-        self.max_position_embeddings = max_position_embeddings
-        self.eos_token_id = eos_token_id
-        self.pad_token_id = pad_token_id
-        self.bos_token_id = bos_token_id
+        self.batch_size = 13
+        self.seq_length = 7
+        self.is_training = True
+        self.use_labels = False
+        self.vocab_size = 99
+        self.hidden_size = 16
+        self.num_hidden_layers = 2
+        self.num_attention_heads = 4
+        self.intermediate_size = 4
+        self.hidden_act = "gelu"
+        self.hidden_dropout_prob = 0.1
+        self.attention_probs_dropout_prob = 0.1
+        self.max_position_embeddings = 20
+        self.eos_token_id = 2
+        self.pad_token_id = 1
+        self.bos_token_id = 0
         torch.manual_seed(0)
 
     def prepare_config_and_inputs(self):
@@ -143,67 +111,21 @@ class BartModelTester:
         config, inputs_dict = self.prepare_config_and_inputs()
         inputs_dict["decoder_input_ids"] = inputs_dict["input_ids"]
         inputs_dict["decoder_attention_mask"] = inputs_dict["attention_mask"]
+        inputs_dict["use_cache"] = False
         return config, inputs_dict
 
-    def create_and_check_decoder_model_past_large_inputs(self, config, inputs_dict):
-        model = BartModel(config=config).get_decoder().to(torch_device).eval()
-        input_ids = inputs_dict["input_ids"]
 
-        # first forward pass
-        outputs = model(input_ids, use_cache=True)
-
-        output, past_key_values = outputs.to_tuple()
-
-        # create hypothetical multiple next token and extent to next_input_ids
-        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
-
-        # append to next input_ids and
-        next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
-
-        output_from_no_past = model(next_input_ids)["last_hidden_state"]
-        output_from_past = model(next_tokens, past_key_values=past_key_values)["last_hidden_state"]
-
-        # select random slice
-        random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
-        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach()
-        output_from_past_slice = output_from_past[:, :, random_slice_idx].detach()
-
-        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])
-
-        # test that outputs are equal for slice
-        self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-2))
-
-    def check_encoder_decoder_model_standalone(self, config, inputs_dict):
-        model = BartModel(config=config).to(torch_device).eval()
-        outputs = model(**inputs_dict)
-
-        encoder_last_hidden_state = outputs.encoder_last_hidden_state
-        last_hidden_state = outputs.last_hidden_state
-
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            encoder = model.get_encoder()
-            encoder.save_pretrained(tmpdirname)
-            encoder = BartEncoder.from_pretrained(tmpdirname).to(torch_device)
-
-        encoder_last_hidden_state_2 = encoder(inputs_dict["input_ids"], attention_mask=inputs_dict["attention_mask"])[
-            0
-        ]
-
-        self.parent.assertTrue((encoder_last_hidden_state_2 - encoder_last_hidden_state).abs().max().item() < 1e-3)
-
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            decoder = model.get_decoder()
-            decoder.save_pretrained(tmpdirname)
-            decoder = BartDecoder.from_pretrained(tmpdirname).to(torch_device)
-
-        last_hidden_state_2 = decoder(
-            input_ids=inputs_dict["decoder_input_ids"],
-            attention_mask=inputs_dict["decoder_attention_mask"],
-            encoder_hidden_states=encoder_last_hidden_state,
-            encoder_attention_mask=inputs_dict["attention_mask"],
-        )[0]
-
-        self.parent.assertTrue((last_hidden_state_2 - last_hidden_state).abs().max().item() < 1e-3)
+def prepare_bart_inputs_dict(
+    config,
+    input_ids,
+    attention_mask=None,
+):
+    if attention_mask is None:
+        attention_mask = input_ids.ne(config.pad_token_id)
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+    }
 
 
 @require_torch
@@ -220,7 +142,7 @@ class BARTModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
     test_missing_keys = False
 
     def setUp(self):
-        self.model_tester = BartModelTester(self)
+        self.model_tester = ModelTester(self)
         self.config_tester = ConfigTester(self, config_class=BartConfig)
 
     def test_config(self):
@@ -247,25 +169,23 @@ class BARTModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs()
         config.use_cache = False
         inputs_dict["input_ids"][:, -2:] = config.pad_token_id
-
+        decoder_input_ids, decoder_attn_mask, causal_mask = _prepare_bart_decoder_inputs(
+            config, inputs_dict["input_ids"]
+        )
         model = BartModel(config).to(torch_device).eval()
+
         decoder_features_with_created_mask = model(**inputs_dict)[0]
-
-        decoder_input_ids = shift_tokens_right(inputs_dict["input_ids"], config.pad_token_id)
-        decoder_attention_mask = decoder_input_ids.ne(config.pad_token_id)
-        decoder_attention_mask[:, 0] = decoder_attention_mask[:, 1]
-
         decoder_features_with_passed_mask = model(
-            decoder_attention_mask=decoder_attention_mask, decoder_input_ids=decoder_input_ids, **inputs_dict
+            decoder_attention_mask=invert_mask(decoder_attn_mask), decoder_input_ids=decoder_input_ids, **inputs_dict
         )[0]
         assert_tensors_close(decoder_features_with_passed_mask, decoder_features_with_created_mask)
-        useless_mask = torch.zeros_like(decoder_attention_mask)
+        useless_mask = torch.zeros_like(decoder_attn_mask)
         decoder_features = model(decoder_attention_mask=useless_mask, **inputs_dict)[0]
         self.assertTrue(isinstance(decoder_features, torch.Tensor))  # no hidden states or attentions
         self.assertEqual(
             decoder_features.size(), (self.model_tester.batch_size, self.model_tester.seq_length, config.d_model)
         )
-        if decoder_attention_mask.min().item() == 0:  # some tokens were masked
+        if decoder_attn_mask.min().item() < -1e3:  # some tokens were masked
             self.assertFalse((decoder_features_with_created_mask == decoder_features).all().item())
 
         # Test different encoder attention masks
@@ -284,43 +204,13 @@ class BARTModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
                 model2, info = model_class.from_pretrained(tmpdirname, output_loading_info=True)
             self.assertEqual(info["missing_keys"], [])
 
-    def test_decoder_model_past_with_large_inputs(self):
-        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)
-
-    def test_encoder_decoder_model_standalone(self):
-        config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common()
-        self.model_tester.check_encoder_decoder_model_standalone(*config_and_inputs)
-
-    # BartForSequenceClassification does not support inputs_embeds
+    @unittest.skip("Passing inputs_embeds not implemented for Bart.")
     def test_inputs_embeds(self):
-        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        pass
 
-        for model_class in (BartModel, BartForConditionalGeneration, BartForQuestionAnswering):
-            model = model_class(config)
-            model.to(torch_device)
-            model.eval()
-
-            inputs = copy.deepcopy(self._prepare_for_class(inputs_dict, model_class))
-
-            if not self.is_encoder_decoder:
-                input_ids = inputs["input_ids"]
-                del inputs["input_ids"]
-            else:
-                encoder_input_ids = inputs["input_ids"]
-                decoder_input_ids = inputs.get("decoder_input_ids", encoder_input_ids)
-                del inputs["input_ids"]
-                inputs.pop("decoder_input_ids", None)
-
-            wte = model.get_input_embeddings()
-            if not self.is_encoder_decoder:
-                inputs["inputs_embeds"] = wte(input_ids)
-            else:
-                inputs["inputs_embeds"] = wte(encoder_input_ids)
-                inputs["decoder_inputs_embeds"] = wte(decoder_input_ids)
-
-            with torch.no_grad():
-                model(**inputs)[0]
+    @unittest.skip("TODO: Decoder embeddings cannot be resized at the moment")
+    def test_resize_embeddings_untied(self):
+        pass
 
     @require_sentencepiece
     @require_tokenizers
@@ -496,6 +386,20 @@ class BartHeadTests(unittest.TestCase):
         model = BartForConditionalGeneration(config).eval().to(torch_device)
         model(**model.dummy_inputs)
 
+    def test_prepare_bart_decoder_inputs(self):
+        config, *_ = self._get_config_and_data()
+        input_ids = _long_tensor(([4, 4, 2]))
+        decoder_input_ids = _long_tensor([[26388, 2, config.pad_token_id]])
+        ignore = float("-inf")
+        decoder_input_ids, decoder_attn_mask, causal_mask = _prepare_bart_decoder_inputs(
+            config, input_ids, decoder_input_ids
+        )
+        expected_causal_mask = torch.tensor(
+            [[0, ignore, ignore], [0, 0, ignore], [0, 0, 0]]  # never attend to the final token, because its pad
+        ).to(input_ids.device)
+        self.assertEqual(decoder_attn_mask.size(), decoder_input_ids.size())
+        self.assertTrue(torch.eq(expected_causal_mask, causal_mask).all())
+
     def test_resize_tokens_embeddings_more(self):
         config, input_ids, _ = self._get_config_and_data()
 
@@ -566,14 +470,14 @@ class BartModelIntegrationTests(unittest.TestCase):
         self.assertTrue(torch.allclose(output[:, :3, :3], expected_slice, atol=TOLERANCE))
 
     @slow
-    def test_base_mask_filling(self):
+    def test_bart_base_mask_filling(self):
         pbase = pipeline(task="fill-mask", model="facebook/bart-base")
         src_text = [" I went to the <mask>."]
         results = [x["token_str"] for x in pbase(src_text)]
         assert "Ġbathroom" in results
 
     @slow
-    def test_large_mask_filling(self):
+    def test_bart_large_mask_filling(self):
         plarge = pipeline(task="fill-mask", model="facebook/bart-large")
         src_text = [" I went to the <mask>."]
         results = [x["token_str"] for x in plarge(src_text)]
@@ -704,7 +608,7 @@ class BartModelIntegrationTests(unittest.TestCase):
 
 
 @require_torch
-class TestBartSinusoidalPositionalEmbeddings(unittest.TestCase):
+class TestSinusoidalPositionalEmbeddings(unittest.TestCase):
     desired_weights = [
         [0, 0, 0, 0, 0],
         [0.84147096, 0.82177866, 0.80180490, 0.78165019, 0.76140374],
@@ -712,30 +616,38 @@ class TestBartSinusoidalPositionalEmbeddings(unittest.TestCase):
     ]
 
     def test_positional_emb_cache_logic(self):
-        emb1 = BartSinusoidalPositionalEmbedding(num_positions=32, embedding_dim=6, padding_idx=1).to(torch_device)
-        no_cache = emb1((4, 10), past_key_values_length=0)
-        yes_cache = emb1((4, 10), past_key_values_length=2)
-
-        self.assertTrue(no_cache.shape == yes_cache.shape == (10, 6))
-        self.assertListEqual(no_cache[2:].tolist(), yes_cache[:-2].tolist())
+        pad = 1
+        input_ids = torch.tensor([[4, 10]], dtype=torch.long, device=torch_device)
+        emb1 = SinusoidalPositionalEmbedding(num_positions=32, embedding_dim=6, padding_idx=pad).to(torch_device)
+        no_cache = emb1(input_ids, use_cache=False)
+        yes_cache = emb1(input_ids, use_cache=True)
+        self.assertEqual((1, 1, 6), yes_cache.shape)  # extra dim to allow broadcasting, feel free to delete!
+        self.assertListEqual(no_cache[-1].tolist(), yes_cache[0][0].tolist())
 
     def test_odd_embed_dim(self):
         # odd embedding_dim is allowed
-        BartSinusoidalPositionalEmbedding(num_positions=4, embedding_dim=5, padding_idx=0).to(torch_device)
+        SinusoidalPositionalEmbedding(num_positions=4, embedding_dim=5, padding_idx=0).to(torch_device)
 
         # odd num_positions is allowed
-        BartSinusoidalPositionalEmbedding(num_positions=5, embedding_dim=4, padding_idx=0).to(torch_device)
+        SinusoidalPositionalEmbedding(num_positions=5, embedding_dim=4, padding_idx=0).to(torch_device)
 
     def test_positional_emb_weights_against_marian(self):
         pad = 1
-        emb1 = BartSinusoidalPositionalEmbedding(num_positions=512, embedding_dim=512, padding_idx=pad).to(
-            torch_device
-        )
+        emb1 = SinusoidalPositionalEmbedding(num_positions=512, embedding_dim=512, padding_idx=pad).to(torch_device)
         weights = emb1.weight.data[:3, :5].tolist()
         for i, (expected_weight, actual_weight) in enumerate(zip(self.desired_weights, weights)):
             for j in range(5):
                 self.assertAlmostEqual(expected_weight[j], actual_weight[j], places=3)
 
+        # test that forward pass is just a lookup, there is no ignore padding logic
+        input_ids = torch.tensor([[4, 10, pad, pad, pad]], dtype=torch.long, device=torch_device)
+        no_cache_pad_zero = emb1(input_ids)
+        self.assertTrue(
+            torch.allclose(
+                torch.tensor(self.desired_weights, device=torch_device), no_cache_pad_zero[:3, :5], atol=1e-3
+            )
+        )
+
     def test_child_config_equivalence(self):
         """Test that configs associated with children of BartForConditionalGeneration are identical."""
         child_classes = [BlenderbotConfig, MBartConfig, MarianConfig, PegasusConfig]

tests/test_modeling_blenderbot.py

@@ -104,6 +104,9 @@ class BlenderbotTesterMixin(ModelTesterMixin, unittest.TestCase):
         self.model_tester = BlenderbotModelTester(self)
         self.config_tester = ConfigTester(self, config_class=BlenderbotConfig)
 
+    def test_inputs_embeds(self):
+        pass
+
     def test_initialization_module(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
         model = BlenderbotForConditionalGeneration(config).model

tests/test_modeling_common.py

@@ -302,8 +302,6 @@ class ModelTesterMixin:
                 # Question Answering model returns start_logits and end_logits
                 if model_class in MODEL_FOR_QUESTION_ANSWERING_MAPPING.values():
                     correct_outlen += 1  # start_logits and end_logits instead of only 1 output
-                if "past_key_values" in outputs:
-                    correct_outlen += 1  # past_key_values have been returned
 
                 self.assertEqual(out_len, correct_outlen)
 
@@ -388,7 +386,7 @@ class ModelTesterMixin:
 
             try:
                 if model.config.is_encoder_decoder:
-                    model.config.use_cache = False  # FSTM still requires this hack -> FSTM should probably be refactored similar to BART afterward
+                    model.config.use_cache = False  # TODO: this should be deleted after bug #7474 is solved
                     input_ids = inputs["input_ids"]
                     attention_mask = inputs["attention_mask"]
                     decoder_input_ids = inputs["decoder_input_ids"]
@@ -1022,6 +1020,7 @@ class ModelTesterMixin:
             )
 
     def test_inputs_embeds(self):
+
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
 
         for model_class in self.all_model_classes:

tests/test_modeling_flax_roberta.py

@@ -57,7 +57,7 @@ class FlaxRobertaModelTest(unittest.TestCase):
                 self.assertEqual(len(fx_outputs), len(pt_outputs), "Output lengths differ between Flax and PyTorch")
 
                 for fx_output, pt_output in zip(fx_outputs, pt_outputs.to_tuple()):
-                    self.assert_almost_equals(fx_output, pt_output.numpy(), 5e-3)
+                    self.assert_almost_equals(fx_output, pt_output.numpy(), 6e-4)
 
     def test_multiple_sequences(self):
         tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")

tests/test_modeling_rag.py

@@ -37,7 +37,7 @@ from transformers.testing_utils import (
     torch_device,
 )
 
-from .test_modeling_bart import BartModelTester
+from .test_modeling_bart import ModelTester as BartModelTester
 from .test_modeling_dpr import DPRModelTester
 from .test_modeling_t5 import T5ModelTester
 

tests/test_modeling_tf_common.py

@@ -344,6 +344,13 @@ class TFModelTesterMixin:
             tf_hidden_states[pt_nans] = 0
 
             max_diff = np.amax(np.abs(tf_hidden_states - pt_hidden_states))
+            # Debug info (remove when fixed)
+            if max_diff >= 4e-2:
+                print("===")
+                print(model_class)
+                print(config)
+                print(inputs_dict)
+                print(pt_inputs_dict)
             self.assertLessEqual(max_diff, 4e-2)
 
             # Check we can load pt model in tf and vice-versa with checkpoint => model functions

utils/check_repo.py

@@ -28,8 +28,6 @@ PATH_TO_DOC = "docs/source/model_doc"
 # Update this list for models that are not tested with a comment explaining the reason it should not be.
 # Being in this list is an exception and should **not** be the rule.
 IGNORE_NON_TESTED = [
-    "BartDecoder",  # Building part of bigger (tested) model.
-    "BartEncoder",  # Building part of bigger (tested) model.
     "BertLMHeadModel",  # Needs to be setup as decoder.
     "DPREncoder",  # Building part of bigger (tested) model.
     "DPRSpanPredictor",  # Building part of bigger (tested) model.
@@ -60,11 +58,9 @@ TEST_FILES_WITH_NO_COMMON_TESTS = [
 # Update this list for models that are not documented with a comment explaining the reason it should not be.
 # Being in this list is an exception and should **not** be the rule.
 IGNORE_NON_DOCUMENTED = [
-    "BartDecoder",  # Building part of bigger (documented) model.
-    "BartEncoder",  # Building part of bigger (documented) model.
     "DPREncoder",  # Building part of bigger (documented) model.
     "DPRSpanPredictor",  # Building part of bigger (documented) model.
-    "T5Stack",  # Building part of bigger (documented) model.
+    "T5Stack",  # Building part of bigger (tested) model.
     "TFDPREncoder",  # Building part of bigger (documented) model.
     "TFDPRSpanPredictor",  # Building part of bigger (documented) model.
 ]
@@ -82,8 +78,6 @@ MODEL_NAME_TO_DOC_FILE = {
 # Update this list for models that are not in any of the auto MODEL_XXX_MAPPING. Being in this list is an exception and
 # should **not** be the rule.
 IGNORE_NON_AUTO_CONFIGURED = [
-    "BartDecoder",
-    "BartEncoder",
     "DPRContextEncoder",
     "DPREncoder",
     "DPRReader",