initial commit

docs/source/main_classes/model.rst

@@ -13,9 +13,10 @@
 Models
 -----------------------------------------------------------------------------------------------------------------------
 
-The base classes :class:`~transformers.PreTrainedModel` and :class:`~transformers.TFPreTrainedModel` implement the
-common methods for loading/saving a model either from a local file or directory, or from a pretrained model
-configuration provided by the library (downloaded from HuggingFace's AWS S3 repository).
+The base classes :class:`~transformers.PreTrainedModel`, :class:`~transformers.TFPreTrainedModel`, and
+:class:`~transformers.FlaxPreTrainedModel` implement the common methods for loading/saving a model either from a local
+file or directory, or from a pretrained model configuration provided by the library (downloaded from HuggingFace's AWS
+S3 repository).
 
 :class:`~transformers.PreTrainedModel` and :class:`~transformers.TFPreTrainedModel` also implement a few methods which
 are common among all the models to:
@@ -57,6 +58,13 @@ TFModelUtilsMixin
     :members:
 
 
+FlaxPreTrainedModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.FlaxPreTrainedModel
+    :members:
+
+
 Generation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

examples/language-modeling/run_mlm_flax.py

@@ -379,7 +379,7 @@ def training_step(optimizer, batch, dropout_rng):
         # 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)
+        logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
         loss, weight_sum = cross_entropy(logits, targets, token_mask)
         return loss / weight_sum
 
@@ -401,7 +401,7 @@ def eval_step(params, batch):
 
     # 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)
+    logits = model(**batch, params=params, train=False)[0]
 
     return compute_metrics(logits, targets, token_mask)
 
@@ -413,7 +413,7 @@ def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndar
     if samples_to_remove != 0:
         samples_idx = samples_idx[:-samples_to_remove]
     sections_split = nb_samples // batch_size
-    batch_idx = jnp.split(samples_idx, sections_split)
+    batch_idx = np.split(samples_idx, sections_split)
     return batch_idx
 
 
@@ -554,8 +554,13 @@ if __name__ == "__main__":
     rng = jax.random.PRNGKey(training_args.seed)
     dropout_rngs = jax.random.split(rng, jax.local_device_count())
 
-    model = FlaxBertForMaskedLM.from_pretrained("bert-base-cased", dtype=jnp.float32, dropout_rate=0.1)
-    model.init(jax.random.PRNGKey(training_args.seed), (training_args.train_batch_size, model.config.max_length))
+    model = FlaxBertForMaskedLM.from_pretrained(
+        "bert-base-cased",
+        dtype=jnp.float32,
+        input_shape=(training_args.train_batch_size, config.max_position_embeddings),
+        seed=training_args.seed,
+        dropout_rate=0.1,
+    )
 
     # Setup optimizer
     optimizer = Adam(

examples/research_projects/performer/run_mlm_performer.py

@@ -0,0 +1,659 @@
+# 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"}
+    )
+    validation_split_percentage: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
+    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)
+
+        logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+        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)[0]
+
+    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)
+        if "validation" not in datasets.keys():
+            datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+            )
+            datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+            )
+    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.
+    if model_args.config_name:
+        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
+    elif model_args.model_name_or_path:
+        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
+    else:
+        config = CONFIG_MAPPING[model_args.model_type]()
+        logger.warning("You are instantiating a new config instance from scratch.")
+
+    if model_args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+        )
+    elif model_args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+        )
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    if training_args.do_train:
+        column_names = datasets["train"].column_names
+    else:
+        column_names = datasets["validation"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    padding = "max_length" if data_args.pad_to_max_length else False
+
+    def tokenize_function(examples):
+        # Remove empty lines
+        examples = [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.
+    data_collator = FlaxDataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=data_args.mlm_probability)
+
+    # Initialize our training
+    rng = jax.random.PRNGKey(training_args.seed)
+    dropout_rngs = jax.random.split(rng, jax.local_device_count())
+
+    model = FlaxBertForMaskedLM.from_pretrained(
+        "bert-base-cased",
+        dtype=jnp.float32,
+        input_shape=(training_args.train_batch_size, config.max_position_embeddings),
+        seed=training_args.seed,
+        dropout_rate=0.1,
+    )
+
+    # 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
+    lr_scheduler_fn = create_learning_rate_scheduler(
+        base_learning_rate=training_args.learning_rate, warmup_steps=max(training_args.warmup_steps, 1)
+    )
+
+    # 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)
+
+    epochs = tqdm(range(nb_epochs), desc=f"Epoch ... (1/{nb_epochs})", position=0)
+    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
+        for batch_idx in tqdm(training_batch_idx, desc="Training...", position=1):
+            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)
+
+        epochs.write(f"Loss: {loss}")
+
+        # ======================== Evaluating ==============================
+        nb_eval_samples = len(tokenized_datasets["validation"])
+        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["validation"][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']})"
+        )
+
+        # Save metrics
+        if has_tensorboard and jax.host_id() == 0:
+            for name, value in eval_summary.items():
+                summary_writer.scalar(name, value, epoch)

setup.py

@@ -97,7 +97,7 @@ _deps = [
     "fastapi",
     "filelock",
     "flake8>=3.8.3",
-    "flax==0.2.2",
+    "flax>=0.2.2",
     "fugashi>=1.0",
     "ipadic>=1.0.0,<2.0",
     "isort>=5.5.4",
@@ -175,7 +175,7 @@ class DepsTableUpdateCommand(Command):
             "deps = {",
             entries,
             "}",
-            ""
+            "",
         ]
         target = "src/transformers/dependency_versions_table.py"
         print(f"updating {target}")
@@ -232,14 +232,14 @@ extras["dev"] = (
 # when modifying the following list, make sure to update src/transformers/dependency_versions_check.py
 install_requires = [
     deps["dataclasses"] + ";python_version<'3.7'",  # dataclasses for Python versions that don't have it
-    deps["filelock"],   # filesystem locks, e.g., to prevent parallel downloads
+    deps["filelock"],  # filesystem locks, e.g., to prevent parallel downloads
     deps["numpy"],
     deps["packaging"],  # utilities from PyPA to e.g., compare versions
-    deps["regex"],      # for OpenAI GPT
-    deps["requests"],   # for downloading models over HTTPS
-    deps["sacremoses"], # for XLM
+    deps["regex"],  # for OpenAI GPT
+    deps["requests"],  # for downloading models over HTTPS
+    deps["sacremoses"],  # for XLM
     deps["tokenizers"],
-    deps["tqdm"],       # progress bars in model download and training scripts
+    deps["tqdm"],  # progress bars in model download and training scripts
 ]
 
 setup(

src/transformers/__init__.py

@@ -936,6 +936,7 @@ else:
 
 
 if is_flax_available():
+    from .modeling_flax_utils import FlaxPreTrainedModel
     from .models.auto import FLAX_MODEL_MAPPING, FlaxAutoModel
     from .models.bert import FlaxBertForMaskedLM, FlaxBertModel
     from .models.roberta import FlaxRobertaModel

src/transformers/dependency_versions_table.py

@@ -10,7 +10,7 @@ deps = {
     "fastapi": "fastapi",
     "filelock": "filelock",
     "flake8": "flake8>=3.8.3",
-    "flax": "flax==0.2.2",
+    "flax": "flax>=0.2.2",
     "fugashi": "fugashi>=1.0",
     "ipadic": "ipadic>=1.0.0,<2.0",
     "isort": "isort>=5.5.4",
@@ -40,8 +40,8 @@ deps = {
     "sphinx-rtd-theme": "sphinx-rtd-theme==0.4.3",
     "sphinx": "sphinx==3.2.1",
     "starlette": "starlette",
-    "tensorflow-cpu": "tensorflow-cpu>=2.0",
-    "tensorflow": "tensorflow>=2.0",
+    "tensorflow-cpu": "tensorflow-cpu>=2.0,<2.4",
+    "tensorflow": "tensorflow>=2.0,<2.4",
     "timeout-decorator": "timeout-decorator",
     "tokenizers": "tokenizers==0.9.4",
     "torch": "torch>=1.0",

src/transformers/file_utils.py

@@ -247,6 +247,7 @@ TRANSFORMERS_CACHE = os.getenv("TRANSFORMERS_CACHE", PYTORCH_TRANSFORMERS_CACHE)
 WEIGHTS_NAME = "pytorch_model.bin"
 TF2_WEIGHTS_NAME = "tf_model.h5"
 TF_WEIGHTS_NAME = "model.ckpt"
+FLAX_WEIGHTS_NAME = "flax_model.msgpack"
 CONFIG_NAME = "config.json"
 MODEL_CARD_NAME = "modelcard.json"
 

src/transformers/modeling_flax_utils.py

@@ -16,17 +16,18 @@
 import os
 from abc import ABC, abstractmethod
 from pickle import UnpicklingError
-from typing import Dict
+from typing import Dict, Set, Tuple, Union
 
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from flax.serialization import to_bytes
-from flax.traverse_util import unflatten_dict
+from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
+from flax.serialization import from_bytes, to_bytes
+from flax.traverse_util import flatten_dict, unflatten_dict
 from jax.random import PRNGKey
 
 from .configuration_utils import PretrainedConfig
-from .file_utils import WEIGHTS_NAME, cached_path, hf_bucket_url, is_remote_url
+from .file_utils import FLAX_WEIGHTS_NAME, WEIGHTS_NAME, cached_path, hf_bucket_url, is_remote_url
 from .utils import logging
 
 
@@ -34,20 +35,8 @@ logger = logging.get_logger(__name__)
 
 
 @jax.jit
-def gelu(x):
-    r"""
-    Gaussian error linear unit activation function.
-
-    Computes the element-wise function:
-
-    .. math::
-      \mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{tanh} \left(
-        \sqrt{\frac{2}{\pi}} \left(x + 0.044715 x^3 \right) \right) \right)
-
-    We explicitly use the approximation rather than the exact formulation for speed. For more information, see
-    `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_, section 2.
-    """
-    return x * 0.5 * (1.0 + jax.lax.erf(x / jnp.sqrt(2.0)))
+def gelu_new(x):
+    return nn.gelu(x, approximate=True)
 
 
 ACT2FN = {
@@ -55,24 +44,40 @@ ACT2FN = {
     "relu": nn.relu,
     "silu": nn.swish,
     "swish": nn.swish,
-    "gelu_new": gelu,
+    "gelu_new": gelu_new,
 }
 
 
 class FlaxPreTrainedModel(ABC):
+    r"""
+    Base class for all models.
+
+    :class:`~transformers.FlaxPreTrainedModel` takes care of storing the configuration of the models and handles
+    methods for loading, downloading and saving models.
+
+    Class attributes (overridden by derived classes):
+
+        - **config_class** (:class:`~transformers.PretrainedConfig`) -- A subclass of
+          :class:`~transformers.PretrainedConfig` to use as configuration class for this model architecture.
+        - **base_model_prefix** (:obj:`str`) -- A string indicating the attribute associated to the base model in
+          derived classes of the same architecture adding modules on top of the base model.
+    """
     config_class = None
-    pretrained_model_archive_map = {}
     base_model_prefix = ""
-    model_class = None
 
     def __init__(
-        self, config: PretrainedConfig, module: nn.Module, params: Dict, seed: int = 0, dtype: jnp.dtype = jnp.float32
+        self,
+        config: PretrainedConfig,
+        module: nn.Module,
+        input_shape: Tuple = (1, 1),
+        seed: int = 0,
+        dtype: jnp.dtype = jnp.float32,
     ):
         if config is None:
             raise ValueError("config cannot be None")
 
-        if params is None:
-            raise ValueError("state cannot be None")
+        if module is None:
+            raise ValueError("module cannot be None")
 
         # Those are private to be exposed as typed property on derived classes.
         self._config = config
@@ -80,9 +85,18 @@ 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
 
+        # randomely initialized parameters
+        random_params = self.init(self.key, input_shape)
+
+        # save required_params as set
+        self._required_params = set(flatten_dict(unfreeze(random_params)).keys())
+        self.params = random_params
+
+    def init(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> Dict:
+        raise NotImplementedError(f"init method has to be implemented for {self}")
+
     @property
     def config(self) -> PretrainedConfig:
         return self._config
@@ -91,24 +105,130 @@ class FlaxPreTrainedModel(ABC):
     def module(self) -> nn.Module:
         return self._module
 
+    @property
+    def params(self) -> Union[Dict, FrozenDict]:
+        return self._params
+
+    @property
+    def required_params(self) -> Set:
+        return self._required_params
+
+    @params.setter
+    def params(self, params: Union[Dict, FrozenDict]):
+        if isinstance(params, FrozenDict):
+            params = unfreeze(params)
+        param_keys = set(flatten_dict(params).keys())
+        if len(self.required_params - param_keys) > 0:
+            raise ValueError(
+                "Some parameters are missing. Make sure that `params` include the following "
+                f"parameters {self.required_params - param_keys}"
+            )
+        self._params = freeze(params)
+
     @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: Union[str, os.PathLike],
+        dtype: jnp.dtype = jnp.float32,
+        *model_args,
+        **kwargs
+    ):
+
         r"""
-        Instantiate a pretrained Flax model from a pre-trained model configuration.
+        Instantiate a pretrained flax model from a pre-trained model configuration.
+
+        The warning `Weights from XXX not initialized from pretrained model` means that the weights of XXX do not come
+        pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
+        task.
+
+        The warning `Weights from XXX not used in YYY` means that the layer XXX is not used by YYY, therefore those
+        weights are discarded.
+
+        Parameters:
+            pretrained_model_name_or_path (:obj:`str` or :obj:`os.PathLike`):
+                Can be either:
+
+                    - A string, the `model id` of a pretrained model hosted inside a model repo on huggingface.co.
+                      Valid model ids can be located at the root-level, like ``bert-base-uncased``, or namespaced under
+                      a user or organization name, like ``dbmdz/bert-base-german-cased``.
+                    - A path to a `directory` containing model weights saved using
+                      :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g., ``./my_model_directory/``.
+                    - A path or url to a `pt index checkpoint file` (e.g, ``./tf_model/model.ckpt.index``). In this
+                      case, ``from_pt`` should be set to :obj:`True`.
+            model_args (sequence of positional arguments, `optional`):
+                All remaning positional arguments will be passed to the underlying model's ``__init__`` method.
+            config (:obj:`Union[PretrainedConfig, str, os.PathLike]`, `optional`):
+                Can be either:
+
+                    - an instance of a class derived from :class:`~transformers.PretrainedConfig`,
+                    - a string or path valid as input to :func:`~transformers.PretrainedConfig.from_pretrained`.
+
+                Configuration for the model to use instead of an automatically loaded configuation. Configuration can
+                be automatically loaded when:
+
+                    - The model is a model provided by the library (loaded with the `model id` string of a pretrained
+                      model).
+                    - The model was saved using :func:`~transformers.PreTrainedModel.save_pretrained` and is reloaded
+                      by supplying the save directory.
+                    - The model is loaded by supplying a local directory as ``pretrained_model_name_or_path`` and a
+                      configuration JSON file named `config.json` is found in the directory.
+            cache_dir (:obj:`Union[str, os.PathLike]`, `optional`):
+                Path to a directory in which a downloaded pretrained model configuration should be cached if the
+                standard cache should not be used.
+            from_pt (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Load the model weights from a PyTorch checkpoint save file (see docstring of
+                ``pretrained_model_name_or_path`` argument).
+            force_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+                cached versions if they exist.
+            resume_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to delete incompletely received files. Will attempt to resume the download if such a
+                file exists.
+            proxies (:obj:`Dict[str, str], `optional`):
+                A dictionary of proxy servers to use by protocol or endpoint, e.g., :obj:`{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+            local_files_only(:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to only look at local files (i.e., do not try to download the model).
+            revision(:obj:`str`, `optional`, defaults to :obj:`"main"`):
+                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+                git-based system for storing models and other artifacts on huggingface.co, so ``revision`` can be any
+                identifier allowed by git.
+            kwargs (remaining dictionary of keyword arguments, `optional`):
+                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
+                :obj:`output_attentions=True`). Behaves differently depending on whether a ``config`` is provided or
+                automatically loaded:
+
+                    - If a configuration is provided with ``config``, ``**kwargs`` will be directly passed to the
+                      underlying model's ``__init__`` method (we assume all relevant updates to the configuration have
+                      already been done)
+                    - If a configuration is not provided, ``kwargs`` will be first passed to the configuration class
+                      initialization function (:func:`~transformers.PretrainedConfig.from_pretrained`). Each key of
+                      ``kwargs`` that corresponds to a configuration attribute will be used to override said attribute
+                      with the supplied ``kwargs`` value. Remaining keys that do not correspond to any configuration
+                      attribute will be passed to the underlying model's ``__init__`` function.
+
+        Examples::
+
+            >>> from transformers import BertConfig, FlaxBertModel
+            >>> # Download model and configuration from huggingface.co and cache.
+            >>> model = FlaxBertModel.from_pretrained('bert-base-cased')
+            >>> # Model was saved using `save_pretrained('./test/saved_model/')` (for example purposes, not runnable).
+            >>> model = FlaxBertModel.from_pretrained('./test/saved_model/')
+            >>> # Loading from a PyTorch checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable).
+            >>> config = BertConfig.from_json_file('./pt_model/config.json')
+            >>> model = FlaxBertModel.from_pretrained('./pt_model/pytorch_model.bin', from_pt=True, config=config)
         """
         config = kwargs.pop("config", None)
-        # state_dict = kwargs.pop("state_dict", None)
         cache_dir = kwargs.pop("cache_dir", None)
-        # from_tf = kwargs.pop("from_tf", False)
+        from_pt = kwargs.pop("from_pt", False)
         force_download = kwargs.pop("force_download", False)
         resume_download = kwargs.pop("resume_download", False)
         proxies = kwargs.pop("proxies", None)
-        # output_loading_info = kwargs.pop("output_loading_info", False)
         local_files_only = kwargs.pop("local_files_only", False)
         revision = kwargs.pop("revision", None)
 
@@ -135,10 +255,28 @@ class FlaxPreTrainedModel(ABC):
 
         # 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):
+            if os.path.isdir(pretrained_model_name_or_path):
+                if from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
+                    # Load from a PyTorch checkpoint
+                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):
+                    # Load from a Flax checkpoint
+                    archive_file = os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)
+                else:
+                    raise EnvironmentError(
+                        "Error no file named {} found in directory {} or `from_pt` set to False".format(
+                            [FLAX_WEIGHTS_NAME, WEIGHTS_NAME],
+                            pretrained_model_name_or_path,
+                        )
+                    )
+            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
                 archive_file = pretrained_model_name_or_path
             else:
-                archive_file = hf_bucket_url(pretrained_model_name_or_path, filename=WEIGHTS_NAME, revision=revision)
+                archive_file = hf_bucket_url(
+                    pretrained_model_name_or_path,
+                    filename=WEIGHTS_NAME if from_pt else FLAX_WEIGHTS_NAME,
+                    revision=revision,
+                )
 
             # redirect to the cache, if necessary
             try:
@@ -169,31 +307,96 @@ class FlaxPreTrainedModel(ABC):
         # Instantiate model.
         with open(resolved_archive_file, "rb") as state_f:
             try:
-                from flax.serialization import from_bytes
-
-                state = from_bytes(cls.model_class, state_f)
-            except TypeError:
-                try:
+                if from_pt:
                     import torch
 
                     state = torch.load(state_f)
-                    state = {k: v.numpy() for k, v in state.items()}
-                    state = cls.convert_from_pytorch(state, config)
-                    state = unflatten_dict({tuple(k.split(".")[1:]): v for k, v in state.items()})
-                except UnpicklingError:
-                    raise EnvironmentError(
-                        f"Unable to convert model {archive_file} to Flax deserializable object. "
-                        "Supported format are PyTorch archive or Flax msgpack"
-                    )
 
-        return cls(config, state, *model_args, **model_kwargs)
+                    state = convert_state_dict_from_pt(cls, state, config)
+                else:
+                    state = from_bytes(cls, state_f.read())
+            except UnpicklingError:
+                raise EnvironmentError(
+                    f"Unable to convert pytorch model {archive_file} to Flax deserializable object. "
+                )
 
-    def save_pretrained(self, folder):
-        folder_abs = os.path.abspath(folder)
+        # init random models
+        model = cls(config, *model_args, **model_kwargs)
 
-        if not os.path.exists(folder_abs):
-            os.mkdir(folder_abs)
+        # if model is base model only use model_prefix key
+        if cls.base_model_prefix not in dict(model.params) and cls.base_model_prefix in state:
+            state = state[cls.base_model_prefix]
 
-        with open(os.path.join(folder_abs, f"{self._config.model_type}.flax", "wb")) as f:
+        # flatten dicts
+        state = flatten_dict(state)
+        random_state = flatten_dict(unfreeze(model.params))
+
+        missing_keys = model.required_params - set(state.keys())
+        unexpected_keys = set(state.keys()) - model.required_params
+
+        # add missing keys as random parameters
+        for missing_key in missing_keys:
+            state[missing_key] = random_state[missing_key]
+
+        if len(unexpected_keys) > 0:
+            logger.warning(
+                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when "
+                f"initializing {model.__class__.__name__}: {unexpected_keys}\n"
+                f"- This IS expected if you are initializing {model.__class__.__name__} from the checkpoint of a model trained on another task "
+                f"or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n"
+                f"- This IS NOT expected if you are initializing {model.__class__.__name__} from the checkpoint of a model that you expect "
+                f"to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
+            )
+        else:
+            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
+
+        if len(missing_keys) > 0:
+            logger.warning(
+                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at {pretrained_model_name_or_path} "
+                f"and are newly initialized: {missing_keys}\n"
+                f"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference."
+            )
+        else:
+            logger.info(
+                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at {pretrained_model_name_or_path}.\n"
+                f"If your task is similar to the task the model of the checkpoint was trained on, "
+                f"you can already use {model.__class__.__name__} for predictions without further training."
+            )
+
+        # set correct parameters
+        model.params = unflatten_dict(state)
+        return model
+
+    def save_pretrained(self, save_directory: Union[str, os.PathLike]):
+        """
+        Save a model and its configuration file to a directory, so that it can be re-loaded using the
+        `:func:`~transformers.FlaxPreTrainedModel.from_pretrained`` class method
+
+        Arguments:
+            save_directory (:obj:`str` or :obj:`os.PathLike`):
+                Directory to which to save. Will be created if it doesn't exist.
+        """
+        if os.path.isfile(save_directory):
+            logger.error("Provided path ({}) should be a directory, not a file".format(save_directory))
+            return
+        os.makedirs(save_directory, exist_ok=True)
+
+        # get abs dir
+        save_directory = os.path.abspath(save_directory)
+        # save config as well
+        self.config.save_pretrained(save_directory)
+
+        # save model
+        with open(os.path.join(save_directory, FLAX_WEIGHTS_NAME), "wb") as f:
             model_bytes = to_bytes(self.params)
             f.write(model_bytes)
+
+
+def convert_state_dict_from_pt(model_class: ABC, state: Dict, config: PretrainedConfig):
+    """
+    Converts a PyTorch parameter state dict to an equivalent Flax parameter state dict
+    """
+    state = {k: v.numpy() for k, v in state.items()}
+    state = model_class.convert_from_pytorch(state, config)
+    state = unflatten_dict({tuple(k.split(".")): v for k, v in state.items()})
+    return state

src/transformers/models/auto/modeling_flax_auto.py

@@ -27,15 +27,6 @@ from .configuration_auto import AutoConfig, BertConfig, RobertaConfig
 logger = logging.get_logger(__name__)
 
 
-ALL_PRETRAINED_MODEL_ARCHIVE_MAP = dict(
-    (key, value)
-    for pretrained_map in [
-        FlaxBertModel.pretrained_model_archive_map,
-        FlaxRobertaModel.pretrained_model_archive_map,
-    ]
-    for key, value, in pretrained_map.items()
-)
-
 FLAX_MODEL_MAPPING = OrderedDict(
     [
         (RobertaConfig, FlaxRobertaModel),
@@ -114,10 +105,9 @@ class FlaxAutoModel(object):
                   organization name, like ``dbmdz/bert-base-german-cased``.
                 - a path to a `directory` containing model weights saved using
                   :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g.: ``./my_model_directory/``.
-                - a path or url to a `tensorflow index checkpoint file` (e.g. `./tf_model/model.ckpt.index`). In this
-                  case, ``from_tf`` should be set to True and a configuration object should be provided as ``config``
-                  argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model
-                  using the provided conversion scripts and loading the PyTorch model afterwards.
+                - a path or url to a `pytorch index checkpoint file` (e.g. `./pt_model/pytorch_model.bin`). In this
+                  case, ``from_pt`` should be set to True and a configuration object should be provided as ``config``
+                  argument.
 
             model_args: (`optional`) Sequence of positional arguments:
                 All remaining positional arguments will be passed to the underlying model's ``__init__`` method
@@ -133,13 +123,6 @@ class FlaxAutoModel(object):
                 - the model is loaded by supplying a local directory as ``pretrained_model_name_or_path`` and a
                   configuration JSON file named `config.json` is found in the directory.
 
-            state_dict: (`optional`) dict:
-                an optional state dictionary for the model to use instead of a state dictionary loaded from saved
-                weights file. This option can be used if you want to create a model from a pretrained configuration but
-                load your own weights. In this case though, you should check if using
-                :func:`~transformers.FlaxPreTrainedModel.save_pretrained` and
-                :func:`~transformers.FlaxPreTrainedModel.from_pretrained` is not a simpler option.
-
             cache_dir: (`optional`) string:
                 Path to a directory in which a downloaded pre-trained model configuration should be cached if the
                 standard cache should not be used.

src/transformers/models/bert/modeling_flax_bert.py

@@ -20,10 +20,11 @@ import numpy as np
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
+from flax.core.frozen_dict import FrozenDict
 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
+from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel
 from ...utils import logging
 from .configuration_bert import BertConfig
 
@@ -205,7 +206,7 @@ class FlaxBertAttention(nn.Module):
     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_states, attention_mask, deterministic: bool = True):
         # 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)
@@ -219,27 +220,28 @@ class FlaxBertAttention(nn.Module):
             bias_init=jax.nn.initializers.zeros,
             name="self",
             dtype=self.dtype,
-        )(hidden_state, attention_mask)
+        )(hidden_states, attention_mask)
 
-        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_state)
+        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_states)
         return layer_norm
 
 
 class FlaxBertIntermediate(nn.Module):
     output_size: int
+    hidden_act: str = "gelu"
     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(
+    def __call__(self, hidden_states):
+        hidden_states = 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)
-        return gelu(dense)
+        )(hidden_states)
+        hidden_states = ACT2FN[self.hidden_act](hidden_states)
+        return hidden_states
 
 
 class FlaxBertOutput(nn.Module):
@@ -249,27 +251,28 @@ class FlaxBertOutput(nn.Module):
 
     @nn.compact
     def __call__(self, intermediate_output, attention_output, deterministic: bool = True):
-        hidden_state = nn.Dense(
+        hidden_states = 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)
-        return hidden_state
+        hidden_states = nn.Dropout(rate=self.dropout_rate)(hidden_states, deterministic=deterministic)
+        hidden_states = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_states + attention_output)
+        return hidden_states
 
 
 class FlaxBertLayer(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     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_states, attention_mask, deterministic: bool = True):
         attention = FlaxBertAttention(
             self.num_heads,
             self.head_size,
@@ -277,9 +280,13 @@ class FlaxBertLayer(nn.Module):
             dropout_rate=self.dropout_rate,
             name="attention",
             dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+        )(hidden_states, attention_mask, deterministic=deterministic)
         intermediate = FlaxBertIntermediate(
-            self.intermediate_size, kernel_init_scale=self.kernel_init_scale, name="intermediate", dtype=self.dtype
+            self.intermediate_size,
+            kernel_init_scale=self.kernel_init_scale,
+            hidden_act=self.hidden_act,
+            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
@@ -297,6 +304,7 @@ class FlaxBertLayerCollection(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     dropout_rate: float = 0.0
     kernel_init_scale: float = 0.2
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
@@ -316,6 +324,7 @@ class FlaxBertLayerCollection(nn.Module):
                 self.intermediate_size,
                 kernel_init_scale=self.kernel_init_scale,
                 dropout_rate=self.dropout_rate,
+                hidden_act=self.hidden_act,
                 name=f"{i}",
                 dtype=self.dtype,
             )
@@ -328,22 +337,24 @@ class FlaxBertEncoder(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     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_states, attention_mask, deterministic: bool = True):
         layer = FlaxBertLayerCollection(
             self.num_layers,
             self.num_heads,
             self.head_size,
             self.intermediate_size,
+            hidden_act=self.hidden_act,
             kernel_init_scale=self.kernel_init_scale,
             dropout_rate=self.dropout_rate,
             name="layer",
             dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+        )(hidden_states, attention_mask, deterministic=deterministic)
         return layer
 
 
@@ -352,10 +363,10 @@ class FlaxBertPooler(nn.Module):
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state):
-        cls_token = hidden_state[:, 0]
+    def __call__(self, hidden_states):
+        cls_token = hidden_states[:, 0]
         out = nn.Dense(
-            hidden_state.shape[-1],
+            hidden_states.shape[-1],
             kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
             name="dense",
             dtype=self.dtype,
@@ -363,62 +374,20 @@ class FlaxBertPooler(nn.Module):
         return nn.tanh(out)
 
 
-class FlaxBertModule(nn.Module):
-    vocab_size: int
-    hidden_size: int
-    type_vocab_size: int
-    max_length: int
-    num_encoder_layers: int
-    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):
-
-        # 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)
-
-        # 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)
-
-        pooled = FlaxBertPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
-        return encoder, pooled
-
-
 class FlaxBertPredictionHeadTransform(nn.Module):
+    hidden_act: str = "gelu"
     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)
+        hidden_states = ACT2FN[self.hidden_act](hidden_states)
+        return FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_states)
 
 
 class FlaxBertLMPredictionHead(nn.Module):
     vocab_size: int
+    hidden_act: str = "gelu"
     dtype: jnp.dtype = jnp.float32
 
     @nn.compact
@@ -428,64 +397,57 @@ class FlaxBertLMPredictionHead(nn.Module):
         #   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 = FlaxBertPredictionHeadTransform(
+            name="transform", hidden_act=self.hidden_act, 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
+    hidden_act: str = "gelu"
     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
+    def __call__(self, hidden_states):
+        hidden_states = FlaxBertLMPredictionHead(
+            vocab_size=self.vocab_size, hidden_act=self.hidden_act, name="predictions", dtype=self.dtype
+        )(hidden_states)
+        return hidden_states
 
 
-@add_start_docstrings(
-    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
-    BERT_START_DOCSTRING,
-)
-class FlaxBertModel(FlaxPreTrainedModel):
+class FlaxBertPreTrainedModel(FlaxPreTrainedModel):
     """
-    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
-    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
-    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
-    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
     """
 
-    model_class = FlaxBertModule
     config_class = BertConfig
     base_model_prefix = "bert"
 
+    def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
+        if token_type_ids is None:
+            token_type_ids = jnp.ones_like(input_ids)
+
+        if position_ids is None:
+            position_ids = jnp.arange(jnp.atleast_2d(input_ids).shape[-1])
+
+        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) -> FrozenDict:
+        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}
+
+        return self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
+
     @staticmethod
     def convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:
         jax_state = dict(pt_state)
@@ -501,6 +463,11 @@ class FlaxBertModel(FlaxPreTrainedModel):
                 key = key.replace("weight", "kernel")
                 jax_state[key] = tensor
 
+            if "decoder.weight" in key:
+                del jax_state[key]
+                key = key.replace("weight", "kernel")
+                jax_state[key] = tensor.T
+
             # SelfAttention needs also to replace "weight" by "kernel"
             if {"query", "key", "value"} & key_parts:
 
@@ -526,7 +493,7 @@ class FlaxBertModel(FlaxPreTrainedModel):
                 jax_state[key] = tensor
 
             # There are some transposed parameters w.r.t their PyTorch counterpart
-            if "intermediate.dense.kernel" in key or "output.dense.kernel" in key:
+            if "intermediate.dense.kernel" in key or "output.dense.kernel" in key or "transform.dense.kernel" in key:
                 jax_state[key] = tensor.T
 
             # Self Attention output projection needs to be transposed
@@ -539,6 +506,11 @@ class FlaxBertModel(FlaxPreTrainedModel):
             if "pooler.dense.kernel" in key:
                 jax_state[key] = tensor.T
 
+            # Hack to correctly load some pytorch models
+            if "predictions.bias" in key:
+                del jax_state[key]
+                jax_state[".".join(key.split(".")[:2]) + ".decoder.bias"] = tensor
+
             # Handle LayerNorm conversion
             if "LayerNorm" in key:
                 del jax_state[key]
@@ -555,7 +527,22 @@ class FlaxBertModel(FlaxPreTrainedModel):
 
         return jax_state
 
-    def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
+
+@add_start_docstrings(
+    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
+    BERT_START_DOCSTRING,
+)
+class FlaxBertModel(FlaxBertPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    """
+
+    def __init__(
+        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
+    ):
         module = FlaxBertModule(
             vocab_size=config.vocab_size,
             hidden_size=config.hidden_size,
@@ -566,10 +553,12 @@ class FlaxBertModel(FlaxPreTrainedModel):
             head_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
             dropout_rate=config.hidden_dropout_prob,
+            hidden_act=config.hidden_act,
             dtype=dtype,
+            **kwargs,
         )
 
-        super().__init__(config, module, state, seed)
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     def __call__(
@@ -601,34 +590,62 @@ class FlaxBertModel(FlaxPreTrainedModel):
             rngs=rngs,
         )
 
-    def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
-        if token_type_ids is None:
-            token_type_ids = jnp.ones_like(input_ids)
 
-        if position_ids is None:
-            position_ids = jnp.arange(jnp.atleast_2d(input_ids).shape[-1])
+class FlaxBertModule(nn.Module):
+    vocab_size: int
+    hidden_size: int
+    type_vocab_size: int
+    max_length: int
+    num_encoder_layers: int
+    num_heads: int
+    head_size: int
+    intermediate_size: int
+    hidden_act: str = "gelu"
+    dropout_rate: float = 0.0
+    kernel_init_scale: float = 0.2
+    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    add_pooling_layer: bool = True
 
-        if attention_mask is None:
-            attention_mask = jnp.ones_like(input_ids)
+    @nn.compact
+    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
 
-        return 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)
 
-    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
-        )
+        # 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,
+            hidden_act=self.hidden_act,
+            name="encoder",
+            dtype=self.dtype,
+        )(embeddings, attention_mask, deterministic=deterministic)
 
-        params_rng, dropout_rng = jax.random.split(rng)
-        rngs = {"params": params_rng, "dropout": dropout_rng}
+        if not self.add_pooling_layer:
+            return encoder
 
-        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
+        pooled = FlaxBertPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
+        return encoder, pooled
 
 
-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(
+class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
+    def __init__(
+        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
+    ):
+        module = FlaxBertForMaskedLMModule(
             vocab_size=config.vocab_size,
             type_vocab_size=config.type_vocab_size,
             hidden_size=config.hidden_size,
@@ -636,10 +653,13 @@ class FlaxBertForMaskedLM(FlaxBertModel):
             head_size=config.hidden_size,
             num_heads=config.num_attention_heads,
             num_encoder_layers=config.num_hidden_layers,
-            max_length=config.max_length,
+            max_length=config.max_position_embeddings,
+            hidden_act=config.hidden_act,
             **kwargs,
         )
 
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
+
     def __call__(
         self,
         input_ids,
@@ -659,7 +679,7 @@ class FlaxBertForMaskedLM(FlaxBertModel):
         if dropout_rng is not None:
             rngs["dropout"] = dropout_rng
 
-        pooled, logits = self.module.apply(
+        return self.module.apply(
             {"params": params or self.params},
             jnp.array(input_ids, dtype="i4"),
             jnp.array(attention_mask, dtype="i4"),
@@ -669,4 +689,45 @@ class FlaxBertForMaskedLM(FlaxBertModel):
             rngs=rngs,
         )
 
-        return logits, pooled
+
+class FlaxBertForMaskedLMModule(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
+    hidden_act: str
+    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 = 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,
+            hidden_act=self.hidden_act,
+            dtype=self.dtype,
+            add_pooling_layer=False,
+            name="bert",
+        )(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 = FlaxBertOnlyMLMHead(
+            vocab_size=self.vocab_size, hidden_act=self.hidden_act, name="cls", dtype=self.dtype
+        )(encoder)
+
+        return (logits,)

src/transformers/models/roberta/modeling_flax_roberta.py

@@ -19,10 +19,11 @@ import numpy as np
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
+from flax.core.frozen_dict import FrozenDict
 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
+from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel
 from ...utils import logging
 from .configuration_roberta import RobertaConfig
 
@@ -33,6 +34,23 @@ _CONFIG_FOR_DOC = "RobertaConfig"
 _TOKENIZER_FOR_DOC = "RobertaTokenizer"
 
 
+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
+
+
 ROBERTA_START_DOCSTRING = r"""
 
     This model inherits from :class:`~transformers.FlaxPreTrainedModel`. Check the superclass documentation for the
@@ -208,7 +226,7 @@ class FlaxRobertaAttention(nn.Module):
     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_states, attention_mask, deterministic: bool = True):
         # 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)
@@ -222,28 +240,29 @@ class FlaxRobertaAttention(nn.Module):
             bias_init=jax.nn.initializers.zeros,
             name="self",
             dtype=self.dtype,
-        )(hidden_state, attention_mask)
+        )(hidden_states, attention_mask)
 
-        layer_norm = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_state)
+        layer_norm = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_states)
         return layer_norm
 
 
 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertIntermediate with Bert->Roberta
 class FlaxRobertaIntermediate(nn.Module):
     output_size: int
+    hidden_act: str = "gelu"
     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(
+    def __call__(self, hidden_states):
+        hidden_states = 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)
-        return gelu(dense)
+        )(hidden_states)
+        hidden_states = ACT2FN[self.hidden_act](hidden_states)
+        return hidden_states
 
 
 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertOutput with Bert->Roberta
@@ -254,27 +273,28 @@ class FlaxRobertaOutput(nn.Module):
 
     @nn.compact
     def __call__(self, intermediate_output, attention_output, deterministic: bool = True):
-        hidden_state = nn.Dense(
+        hidden_states = 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)
-        return hidden_state
+        hidden_states = nn.Dropout(rate=self.dropout_rate)(hidden_states, deterministic=deterministic)
+        hidden_states = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_states + attention_output)
+        return hidden_states
 
 
 class FlaxRobertaLayer(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     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_states, attention_mask, deterministic: bool = True):
         attention = FlaxRobertaAttention(
             self.num_heads,
             self.head_size,
@@ -282,10 +302,11 @@ class FlaxRobertaLayer(nn.Module):
             dropout_rate=self.dropout_rate,
             name="attention",
             dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+        )(hidden_states, attention_mask, deterministic=deterministic)
         intermediate = FlaxRobertaIntermediate(
             self.intermediate_size,
             kernel_init_scale=self.kernel_init_scale,
+            hidden_act=self.hidden_act,
             name="intermediate",
             dtype=self.dtype,
         )(attention)
@@ -306,6 +327,7 @@ class FlaxRobertaLayerCollection(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     dropout_rate: float = 0.0
     kernel_init_scale: float = 0.2
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
@@ -325,6 +347,7 @@ class FlaxRobertaLayerCollection(nn.Module):
                 self.intermediate_size,
                 kernel_init_scale=self.kernel_init_scale,
                 dropout_rate=self.dropout_rate,
+                hidden_act=self.hidden_act,
                 name=f"{i}",
                 dtype=self.dtype,
             )
@@ -338,22 +361,24 @@ class FlaxRobertaEncoder(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     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_states, attention_mask, deterministic: bool = True):
         layer = FlaxRobertaLayerCollection(
             self.num_layers,
             self.num_heads,
             self.head_size,
             self.intermediate_size,
+            hidden_act=self.hidden_act,
             kernel_init_scale=self.kernel_init_scale,
             dropout_rate=self.dropout_rate,
             name="layer",
             dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+        )(hidden_states, attention_mask, deterministic=deterministic)
         return layer
 
 
@@ -363,10 +388,10 @@ class FlaxRobertaPooler(nn.Module):
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state):
-        cls_token = hidden_state[:, 0]
+    def __call__(self, hidden_states):
+        cls_token = hidden_states[:, 0]
         out = nn.Dense(
-            hidden_state.shape[-1],
+            hidden_states.shape[-1],
             kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
             name="dense",
             dtype=self.dtype,
@@ -374,64 +399,12 @@ class FlaxRobertaPooler(nn.Module):
         return nn.tanh(out)
 
 
-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertModule with Bert->Roberta
-class FlaxRobertaModule(nn.Module):
-    vocab_size: int
-    hidden_size: int
-    type_vocab_size: int
-    max_length: int
-    num_encoder_layers: int
-    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):
-
-        # 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)
-
-        # 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)
-
-        pooled = FlaxRobertaPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
-        return encoder, pooled
-
-
-@add_start_docstrings(
-    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
-    ROBERTA_START_DOCSTRING,
-)
-class FlaxRobertaModel(FlaxPreTrainedModel):
+class FlaxRobertaPreTrainedModel(FlaxPreTrainedModel):
     """
-    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
-    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
-    all you need`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
-    Kaiser and Illia Polosukhin.
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
     """
 
-    model_class = FlaxRobertaModule
     config_class = RobertaConfig
     base_model_prefix = "roberta"
 
@@ -504,7 +477,49 @@ class FlaxRobertaModel(FlaxPreTrainedModel):
 
         return jax_state
 
-    def __init__(self, config: RobertaConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
+    def init(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
+        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}
+
+        return 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):
+        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)
+
+        if attention_mask is None:
+            attention_mask = jnp.ones_like(input_ids)
+
+        return input_ids, attention_mask, token_type_ids, position_ids
+
+
+@add_start_docstrings(
+    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
+    ROBERTA_START_DOCSTRING,
+)
+class FlaxRobertaModel(FlaxRobertaPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
+    Kaiser and Illia Polosukhin.
+    """
+
+    def __init__(
+        self,
+        config: RobertaConfig,
+        input_shape: Tuple = (1, 1),
+        seed: int = 0,
+        dtype: jnp.dtype = jnp.float32,
+        **kwargs
+    ):
         module = FlaxRobertaModule(
             vocab_size=config.vocab_size,
             hidden_size=config.hidden_size,
@@ -513,12 +528,14 @@ class FlaxRobertaModel(FlaxPreTrainedModel):
             num_encoder_layers=config.num_hidden_layers,
             num_heads=config.num_attention_heads,
             head_size=config.hidden_size,
+            hidden_act=config.hidden_act,
             intermediate_size=config.intermediate_size,
             dropout_rate=config.hidden_dropout_prob,
             dtype=dtype,
+            **kwargs,
         )
 
-        super().__init__(config, module, state, seed)
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
 
     @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     def __call__(
@@ -550,42 +567,53 @@ class FlaxRobertaModel(FlaxPreTrainedModel):
             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}
+# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertModule with Bert->Roberta
+class FlaxRobertaModule(nn.Module):
+    vocab_size: int
+    hidden_size: int
+    type_vocab_size: int
+    max_length: int
+    num_encoder_layers: int
+    num_heads: int
+    head_size: int
+    intermediate_size: int
+    hidden_act: str = "gelu"
+    dropout_rate: float = 0.0
+    kernel_init_scale: float = 0.2
+    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    add_pooling_layer: bool = True
 
-        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
+    @nn.compact
+    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
 
-    def _check_inputs(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)
 
-        if token_type_ids is None:
-            token_type_ids = jnp.ones_like(input_ids)
+        # 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,
+            hidden_act=self.hidden_act,
+            name="encoder",
+            dtype=self.dtype,
+        )(embeddings, attention_mask, deterministic=deterministic)
 
-        if position_ids is None:
-            position_ids = create_position_ids_from_input_ids(input_ids, self.config.pad_token_id)
+        if not self.add_pooling_layer:
+            return encoder
 
-        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
+        pooled = FlaxRobertaPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
+        return encoder, pooled

src/transformers/utils/dummy_flax_objects.py

@@ -2,6 +2,15 @@
 from ..file_utils import requires_flax
 
 
+class FlaxPreTrainedModel:
+    def __init__(self, *args, **kwargs):
+        requires_flax(self)
+
+    @classmethod
+    def from_pretrained(self, *args, **kwargs):
+        requires_flax(self)
+
+
 FLAX_MODEL_MAPPING = None
 
 

tests/test_modeling_flax_bert.py

@@ -14,14 +14,16 @@
 
 import unittest
 
+import numpy as np
+
 from transformers import BertConfig, is_flax_available
-from transformers.testing_utils import require_flax
+from transformers.testing_utils import require_flax, slow
 
 from .test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask
 
 
 if is_flax_available():
-    from transformers.models.bert.modeling_flax_bert import FlaxBertModel
+    from transformers.models.bert.modeling_flax_bert import FlaxBertForMaskedLM, FlaxBertModel
 
 
 class FlaxBertModelTester(unittest.TestCase):
@@ -105,7 +107,14 @@ class FlaxBertModelTester(unittest.TestCase):
 @require_flax
 class FlaxBertModelTest(FlaxModelTesterMixin, unittest.TestCase):
 
-    all_model_classes = (FlaxBertModel,) if is_flax_available() else ()
+    all_model_classes = (FlaxBertModel, FlaxBertForMaskedLM) if is_flax_available() else ()
 
     def setUp(self):
         self.model_tester = FlaxBertModelTester(self)
+
+    @slow
+    def test_model_from_pretrained(self):
+        for model_class_name in self.all_model_classes:
+            model = model_class_name.from_pretrained("bert-base-cased")
+            outputs = model(np.ones((1, 1)))
+            self.assertIsNotNone(outputs)

tests/test_modeling_flax_common.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import random
+import tempfile
 
 import numpy as np
 
@@ -26,7 +27,7 @@ if is_flax_available():
 
     import jax
     import jax.numpy as jnp
-    from flax.traverse_util import unflatten_dict
+    from transformers.modeling_flax_utils import convert_state_dict_from_pt
 
     os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.12"  # assumed parallelism: 8
 
@@ -59,21 +60,13 @@ def random_attention_mask(shape, rng=None):
     return attn_mask
 
 
-def convert_pt_model_to_flax(pt_model, config, flax_model_cls):
-    state = pt_model.state_dict()
-    state = {k: v.numpy() for k, v in state.items()}
-    state = flax_model_cls.convert_from_pytorch(state, config)
-    state = unflatten_dict({tuple(k.split(".")): v for k, v in state.items()})
-    return flax_model_cls(config, state, dtype=jnp.float32)
-
-
 @require_flax
 class FlaxModelTesterMixin:
     model_tester = None
     all_model_classes = ()
 
     def assert_almost_equals(self, a: np.ndarray, b: np.ndarray, tol: float):
-        diff = np.abs((a - b)).sum()
+        diff = np.abs((a - b)).max()
         self.assertLessEqual(diff, tol, f"Difference between torch and flax is {diff} (>= {tol}).")
 
     @require_torch
@@ -86,30 +79,54 @@ class FlaxModelTesterMixin:
                 pt_model_class = getattr(transformers, pt_model_class_name)
                 pt_model = pt_model_class(config).eval()
 
-                fx_model = convert_pt_model_to_flax(pt_model, config, model_class)
+                fx_state = convert_state_dict_from_pt(model_class, pt_model.state_dict(), config)
+                fx_model = model_class(config, dtype=jnp.float32)
+                fx_model.params = fx_state
 
                 pt_inputs = {k: torch.tensor(v.tolist()) for k, v in inputs_dict.items()}
 
                 with torch.no_grad():
                     pt_outputs = pt_model(**pt_inputs).to_tuple()
+
                 fx_outputs = fx_model(**inputs_dict)
                 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):
-                    self.assert_almost_equals(fx_output, pt_output.numpy(), 5e-3)
+                    self.assert_almost_equals(fx_output, pt_output.numpy(), 1e-3)
+
+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    pt_model.save_pretrained(tmpdirname)
+                    fx_model_loaded = model_class.from_pretrained(tmpdirname, from_pt=True)
+
+                fx_outputs_loaded = fx_model_loaded(**inputs_dict)
+                self.assertEqual(
+                    len(fx_outputs_loaded), len(pt_outputs), "Output lengths differ between Flax and PyTorch"
+                )
+                for fx_output_loaded, pt_output in zip(fx_outputs_loaded, pt_outputs):
+                    self.assert_almost_equals(fx_output_loaded, pt_output.numpy(), 5e-3)
+
+    def test_from_pretrained_save_pretrained(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            with self.subTest(model_class.__name__):
+                model = model_class(config)
+
+                outputs = model(**inputs_dict)
+
+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    model.save_pretrained(tmpdirname)
+                    model_loaded = model_class.from_pretrained(tmpdirname)
+
+                outputs_loaded = model_loaded(**inputs_dict)
+                for output_loaded, output in zip(outputs_loaded, outputs):
+                    self.assert_almost_equals(output_loaded, output, 5e-3)
 
-    @require_torch
     def test_jit_compilation(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
 
         for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__):
-
-                # TODO later: have some way to initialize easily a Flax model from config, for now I go through PT
-                pt_model_class_name = model_class.__name__[4:]  # Skip the "Flax" at the beginning
-                pt_model_class = getattr(transformers, pt_model_class_name)
-                pt_model = pt_model_class(config).eval()
-
-                model = convert_pt_model_to_flax(pt_model, config, model_class)
+                model = model_class(config)
 
                 @jax.jit
                 def model_jitted(input_ids, attention_mask=None, token_type_ids=None):
@@ -125,3 +142,14 @@ class FlaxModelTesterMixin:
                 self.assertEqual(len(outputs), len(jitted_outputs))
                 for jitted_output, output in zip(jitted_outputs, outputs):
                     self.assertEqual(jitted_output.shape, output.shape)
+
+    def test_naming_convention(self):
+        for model_class in self.all_model_classes:
+            model_class_name = model_class.__name__
+            module_class_name = (
+                model_class_name[:-5] + "Module" if model_class_name[-5:] == "Model" else model_class_name + "Module"
+            )
+            bert_modeling_flax_module = __import__(model_class.__module__, fromlist=[module_class_name])
+            module_cls = getattr(bert_modeling_flax_module, module_class_name)
+
+            self.assertIsNotNone(module_cls)

tests/test_modeling_flax_roberta.py

@@ -14,8 +14,10 @@
 
 import unittest
 
+import numpy as np
+
 from transformers import RobertaConfig, is_flax_available
-from transformers.testing_utils import require_flax
+from transformers.testing_utils import require_flax, slow
 
 from .test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask
 
@@ -109,3 +111,10 @@ class FlaxRobertaModelTest(FlaxModelTesterMixin, unittest.TestCase):
 
     def setUp(self):
         self.model_tester = FlaxRobertaModelTester(self)
+
+    @slow
+    def test_model_from_pretrained(self):
+        for model_class_name in self.all_model_classes:
+            model = model_class_name.from_pretrained("roberta-base")
+            outputs = model(np.ones((1, 1)))
+            self.assertIsNotNone(outputs)