Integrate XLM-R into PyText

Summary: Adding the ability to load and finetune XLM-R models in PyText.

Reviewed By: rutyrinott

Differential Revision: D18382033

fbshipit-source-id: e201b2f50129814950784fc255f4e6bfb4610352

Author: Kartikay Khandelwal

pytext/data/roberta_tensorizer.py

@@ -19,7 +19,7 @@ class Config(BERTTensorizerBase.Config):
         vocab_file: str = (
             "manifold://pytext_training/tree/static/vocabs/bpe/gpt2/dict.txt"
         )
-        tokenizer: GPT2BPETokenizer.Config = GPT2BPETokenizer.Config()
+        tokenizer: Tokenizer.Config = GPT2BPETokenizer.Config()
         max_seq_len: int = 256
 
     @classmethod

pytext/docs/source/xlm_r.rst

@@ -0,0 +1,47 @@
+# Unsupervised Cross-lingual Representation Learning at Scale (XLM-RoBERTa)
+
+## Introduction
+
+XLM-R (XLM-RoBERTa) is scaled cross lingual sentence encoder. It is trained on `2.5T` of data across `100` languages data filtered from Common Crawl. XLM-R achieves state-of-the-arts results on multiple cross lingual benchmarks.
+
+## Pre-trained models
+
+Model | Description | #params | vocab size | Download
+---|---|---|---|---
+`xlmr.base.v0` | XLM-R using the BERT-base architecture | 250M | 250k | [xlm.base.v0.tar.gz](https://dl.fbaipublicfiles.com/fairseq/models/xlmr.base.v0.tar.gz)
+`xlmr.large.v0` | XLM-R using the BERT-large architecture | 560M | 250k | [xlm.large.v0.tar.gz](https://dl.fbaipublicfiles.com/fairseq/models/xlmr.large.v0.tar.gz)
+
+(Note: The above models are still under training, we will update the weights, once fully trained, the results are based on the above checkpoints.)
+
+## Results
+
+**[XNLI (Conneau et al., 2018)](https://arxiv.org/abs/1809.05053)**
+
+Model | average | en | fr | es | de | el | bg | ru | tr | ar | vi | th | zh | hi | sw | ur
+---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---
+`roberta.large.mnli` _(TRANSLATE-TEST)_ | 77.8 | 91.3 | 82.9 | 84.3 | 81.2 | 81.7 | 83.1 | 78.3 | 76.8 | 76.6 | 74.2 | 74.1 | 77.5 | 70.9 | 66.7 | 66.8
+`xlmr.large.v0` _(TRANSLATE-TRAIN-ALL)_ | **82.4** | 88.7 | 85.2 | 85.6 | 84.6 | 83.6 | 85.5 | 82.4 | 81.6 | 80.9 | 83.4 | 80.9 | 83.3 | 79.8 | 75.9 | 74.3
+
+**[MLQA (Lewis et al., 2018)](https://arxiv.org/abs/1910.07475)**
+
+Model | average | en | es | de | ar | hi | vi | zh
+---|---|---|---|---|---|---|---|---
+`BERT-large` | - | 80.2/67.4 | - | - | - | - | - | -
+`mBERT` | 57.7 / 41.6 | 77.7 / 65.2 | 64.3 / 46.6 | 57.9 / 44.3 | 45.7 / 29.8| 43.8 / 29.7 | 57.1 / 38.6 | 57.5 / 37.3
+`xlmr.large.v0` | **70.0 / 52.2** | 80.1 / 67.7 | 73.2 / 55.1 | 68.3 / 53.7 | 62.8 / 43.7 | 68.3 / 51.0 | 70.5 / 50.1 | 67.1 / 44.4
+
+
+## Citation
+
+```bibtex
+@article{,
+    title = {Unsupervised Cross-lingual Representation Learning at Scale},
+    author = {Alexis Conneau and Kartikay Khandelwal
+        and Naman Goyal and Vishrav Chaudhary and Guillaume Wenzek
+        and Francisco Guzm\'an and Edouard Grave and Myle Ott
+        and Luke Zettlemoyer and Veselin Stoyanov
+    },
+    journal={},
+    year = {2019},
+}
+```

pytext/models/roberta.py

@@ -57,17 +57,25 @@ class RoBERTaEncoder(RoBERTaEncoderBase):
     """A PyTorch RoBERTa implementation"""
 
     class Config(RoBERTaEncoderBase.Config):
+        embedding_dim: int = 768
+        vocab_size: int = 50265
         num_encoder_layers: int = 12
         num_attention_heads: int = 12
         model_path: str = (
             "manifold://pytext_training/tree/static/models/roberta_base_torch.pt"
         )
+        # Loading the state dict of the model depends on whether the model was
+        # previously finetuned in PyText or not. If it was finetuned then we
+        # dont need to translate the state dict and can just load it`
+        # directly.
+        is_finetuned: bool = False
 
     def __init__(self, config: Config, output_encoded_layers: bool, **kwarg) -> None:
         super().__init__(config, output_encoded_layers=output_encoded_layers)
         # assert config.pretrained_encoder.load_path, "Load path cannot be empty."
         self.encoder = SentenceEncoder(
             transformer=Transformer(
+                vocab_size=config.vocab_size,
                 embedding_dim=config.embedding_dim,
                 layers=[
                     TransformerLayer(
@@ -84,7 +92,13 @@ def __init__(self, config: Config, output_encoded_layers: bool, **kwarg) -> None
             config.model_path,
             map_location=lambda s, l: default_restore_location(s, "cpu"),
         )
-        self.encoder.load_roberta_state_dict(roberta_state["model"])
+        # In case the model has previously been loaded in PyText and finetuned,
+        # then we dont need to do the special state dict translation. Load
+        # it directly
+        if not config.is_finetuned:
+            self.encoder.load_roberta_state_dict(roberta_state["model"])
+        else:
+            self.load_state_dict(roberta_state)
         self.representation_dim = self.encoder.transformer.token_embedding.weight.size(
             -1
         )