[Longformer] Docs and clean API (#4464)

* add longformer docs

* improve docs

docs/source/index.rst

@@ -109,3 +109,4 @@ The library currently contains PyTorch and Tensorflow implementations, pre-train
     model_doc/dialogpt
     model_doc/reformer
     model_doc/marian
+    model_doc/longformer

docs/source/model_doc/longformer.rst

+Longformer
+----------------------------------------------------
+**DISCLAIMER:** This model is still a work in progress, if you see something strange,
+file a `Github Issue <https://github.com/huggingface/transformers/issues/new?assignees=&labels=&template=bug-report.md&title>`_
+
+Overview
+~~~~~
+The Longformer model was presented in `Longformer: The Long-Document Transformer <https://arxiv.org/pdf/2004.05150.pdf>`_ by Iz Beltagy, Matthew E. Peters, Arman Cohan.
+Here the abstract: 
+
+*Transformer-based models are unable to process long sequences due to their self-attention operation, which scales quadratically with the sequence length. To address this limitation, we introduce the Longformer with an attention mechanism that scales linearly with sequence length, making it easy to process documents of thousands of tokens or longer. Longformer's attention mechanism is a drop-in replacement for the standard self-attention and combines a local windowed attention with a task motivated global attention. Following prior work on long-sequence transformers, we evaluate Longformer on character-level language modeling and achieve state-of-the-art results on text8 and enwik8. In contrast to most prior work, we also pretrain Longformer and finetune it on a variety of downstream tasks. Our pretrained Longformer consistently outperforms RoBERTa on long document tasks and sets new state-of-the-art results on WikiHop and TriviaQA.*
+
+The Authors' code can be found `here <https://github.com/allenai/longformer>`_ .
+
+Longformer Self Attention
+~~~~~~~~~~~~~~~~~~~~
+Longformer self attention employs self attention on both a "local" context and a "global" context.
+Most tokens only attend "locally" to each other meaning that each token attends to its :math:`\frac{1}{2} w` previous tokens and :math:`\frac{1}{2} w` succeding tokens with :math:`w` being the window length as defined in `config.attention_window`. Note that `config.attention_window` can be of type ``list`` to define a different :math:`w` for each layer. 
+A selecetd few tokens attend "globally" to all other tokens, as it is conventionally done for all tokens in *e.g.* `BertSelfAttention`.
+
+Note that "locally" and "globally" attending tokens are projected by different query, key and value matrices.
+Also note that every "locally" attending token not only attends to tokens within its window :math:`w`, but also to all "globally" attending tokens so that global attention is *symmetric*.
+
+The user can define which tokens are masked, which tokens attend "locally" and which tokens attend "globally" by setting the `config.attention_mask` `torch.Tensor` appropriately. In contrast to other models `Longformer` accepts the following values in `config.attention_mask`: `0` - the token is masked and not attended at all (as is done in other models), `1` - the token attends "locally", `2` - token attends "globally". For more information please also refer to :func:`~transformers.LongformerModel.forward` method.
+
+Using Longformer self attention, the memory and time complexity of the query-key matmul operation, which usually represents the memory and time bottleneck, can be reduced from :math:`\mathcal{O}(n_s \times n_s)` to :math:`\mathcal{O}(n_s \times w)`, with :math:`n_s` being the sequence length and :math:`w` being the average window size. It is assumed that the number of "globally" attending tokens is insignificant as compared to the number of "locally" attending tokens.
+
+For more information, please refer to the official `paper <https://arxiv.org/pdf/2004.05150.pdf>`_ .
+
+
+Training
+~~~~~~~~~~~~~~~~~~~~
+``LongformerForMaskedLM`` is trained the exact same way, ``RobertaForMaskedLM`` is trained and 
+should be used as follows:
+
+::
+
+  input_ids = tokenizer.encode('This is a sentence from [MASK] training data', return_tensors='pt')
+  mlm_labels = tokenizer.encode('This is a sentence from the training data', return_tensors='pt')
+
+  loss = model(input_ids, labels=input_ids, masked_lm_labels=mlm_labels)[0]
+
+
+LongformerConfig
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LongformerConfig
+    :members:
+
+
+LongformerTokenizer
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LongformerTokenizer
+    :members: 
+
+
+LongformerModel
+~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LongformerModel
+    :members:
+
+
+LongformerForMaskedLM
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LongformerForMaskedLM
+    :members:

docs/source/pretrained_models.rst

@@ -305,9 +305,9 @@ For a list that includes community-uploaded models, refer to `https://huggingfac
 | MarianMT          | ``Helsinki-NLP/opus-mt-{src}-{tgt}``                       | | 12-layer, 512-hidden, 8-heads, ~74M parameter Machine translation models. Parameter counts vary depending on vocab size.            |
 |                   |                                                            | | (see `model list <https://huggingface.co/Helsinki-NLP>`_)                                                                           |
 +-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Longformer        | ``longformer-base-4096``                                   | | 12-layer, 768-hidden, 12-heads, ~149M parameters                                                                                     |
+| Longformer        | ``longformer-base-4096``                                   | | 12-layer, 768-hidden, 12-heads, ~149M parameters                                                                                    |
 |                   |                                                            | | Starting from RoBERTa-base checkpoint, trained on documents of max length 4,096                                                     |
 |                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                   | ``longformer-large-4096``                                  | | 24-layer, 1024-hidden, 16-heads, ~435M parameters                                                                                    |
+|                   | ``longformer-large-4096``                                  | | 24-layer, 1024-hidden, 16-heads, ~435M parameters                                                                                   |
 |                   |                                                            | | Starting from RoBERTa-large checkpoint, trained on documents of max length 4,096                                                    |
 +-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+

src/transformers/configuration_longformer.py

@@ -38,6 +38,12 @@ class LongformerConfig(RobertaConfig):
         The :class:`~transformers.LongformerConfig` class directly inherits :class:`~transformers.RobertaConfig`.
         It reuses the same defaults. Please check the parent class for more information.
 
+        Args:
+            attention_window (:obj:`int` or :obj:`List[int]`, optional, defaults to 512):
+                Size of an attention window around each token. If :obj:`int`, use the same size for all layers.
+                To specify a different window size for each layer, use a :obj:`List[int]` where
+                ``len(attention_window) == num_hidden_layers``.
+
         Example::
 
             from transformers import LongformerConfig, LongformerModel
@@ -58,18 +64,6 @@ class LongformerConfig(RobertaConfig):
     pretrained_config_archive_map = LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP
     model_type = "longformer"
 
-    def __init__(self, attention_window: Union[List[int], int] = 512, attention_mode: str = "longformer", **kwargs):
-        """
-        Args:
-            attention_window (:obj:`int` or :obj:`List[int]`, optional, defaults to 512):
-                Size of an attention window around each token. If :obj:`int`, use the same size for all layers.
-                To specify a different window size for each layer, use a :obj:`List[int]` where
-                `len(attention_window) == num_hidden_layers`.
-            attention_mode (:obj:`str`, optional, possible values ['longformer', 'bert'], defaults to 'longformer'):
-                Type of selfattention. Use 'longformer' for :obj:`LongformerSelfAttention` or 'bert' for
-                standard BERT full n^2 self attention using :obj:`modeling_bert.BertSelfAttention`. Note that full n^2
-                selfattention is supported just for comparison, but it will OOM for long sequences.
-        """
+    def __init__(self, attention_window: Union[List[int], int] = 512, **kwargs):
         super().__init__(**kwargs)
         self.attention_window = attention_window
-        self.attention_mode = attention_mode

src/transformers/modeling_longformer.py

@@ -476,25 +476,18 @@ class LongformerModel(RobertaModel):
         super().__init__(config)
 
         if isinstance(config.attention_window, int):
-            assert config.attention_window % 2 == 0, "`attention_window` has to be an even value"
-            assert config.attention_window > 0, "`attention_window` has to be positive"
+            assert config.attention_window % 2 == 0, "`config.attention_window` has to be an even value"
+            assert config.attention_window > 0, "`config.attention_window` has to be positive"
             config.attention_window = [config.attention_window] * config.num_hidden_layers  # one value per layer
         else:
             assert len(config.attention_window) == config.num_hidden_layers, (
-                "`len(attention_window)` should equal `num_hidden_layers`. "
+                "`len(config.attention_window)` should equal `config.num_hidden_layers`. "
                 f"Expected {config.num_hidden_layers}, given {len(config.attention_window)}"
             )
 
-        if config.attention_mode == "bert":
-            pass  # do nothing, use the default `modeling_bert.BertSelfAttention` (will OOM for long sequences)
-        elif config.attention_mode == "longformer":
-            for i, layer in enumerate(self.encoder.layer):
-                # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
-                layer.attention.self = LongformerSelfAttention(config, layer_id=i)
-        else:
-            raise ValueError(
-                f'Expected values of `attention_mode` are "longformer" or "bert", given {config.attention_mode}'
-            )
+        for i, layer in enumerate(self.encoder.layer):
+            # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
+            layer.attention.self = LongformerSelfAttention(config, layer_id=i)
 
         self.init_weights()