block_stride + fixes + readme

README.md

@@ -45,6 +45,7 @@ It achieves the following results on the test set:
 | Length | Global tokens | Fine-tuning | Block Size | Sparsity | Connexions | R1    | R2    | RL    | RLsum |
 |:------ |:------------- |:----------- |:---------- |:-------- | :--------- |:----- |:----- |:----- |:----- |
 | 16384  | 64            | Full        | 256        | 0        | 768        | 48.74 | 20.88 | 28.50 | 44.23 |
+| 16384  | 1             | Full        | 256        | 0        | 768        | 48.66 | 20.92 | 28.50 | 44.18 |
 | 16384  | 64            | Global only | 256        | 0        | 768        | 48.08 | 20.42 | 28.00 | 43.65 |
 | 16384  | 1             | None        | 256        | 0        | 768        | 47.03 | 20.19 | 28.26 | 42.69 |
 

modeling_lsg_bart.py

@@ -1167,25 +1167,13 @@ class LSGBartEncoder(LSGBartPretrainedModel):
         pad = t % self.block_size
         
         # Check if t is multiple of block_size and pad
-        if t > b and pad > 0:
+        if self.adaptive and t > b and pad > 0:
             pad_length = self.block_size - pad
             if input_ids is not None:
                 input_ids = torch.nn.functional.pad(input_ids, (0, pad_length), value=self.pad_idx)
             else:
                 inputs_embeds = torch.nn.functional.pad(inputs_embeds.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
             attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=0)
-
-        # else adaptive sequence length
-        elif self.adaptive:
-            # Get last non zero mask index
-            s = int(attention_mask.cumsum(dim=-1).argmax(dim=-1).max()) + 1
-            if s < t and self.block_size is not None:
-                s = max(2, s // self.block_size + 1) * self.block_size if s > b else s
-                if input_ids is not None:
-                    input_ids = input_ids[:, :s]
-                else:
-                    inputs_embeds = inputs_embeds[:, :s]
-                attention_mask = attention_mask[:, :s]
         
         n, t_ = attention_mask.size()
         
@@ -1211,9 +1199,7 @@ class LSGBartEncoder(LSGBartPretrainedModel):
             offset = 0
 
         # Adapt sequence to initial shape
-        if diff > 0:
-            context = torch.nn.functional.pad(context.transpose(-1, -2), pad=(0, diff), value=0).transpose(-1, -2)
-        elif diff < 0:
+        if diff < 0:
             context = context[:, :t + offset]
         
         if return_dict:
@@ -1325,7 +1311,7 @@ class LSGBartEncoder(LSGBartPretrainedModel):
         )
 
 
-class LSGBartDecoder(LSGBartPretrainedModel):
+class LSGBartDecoder(BartDecoder, LSGBartPretrainedModel):
     """
     Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a :class:`LSGBartDecoderLayer`
     Args:
@@ -1334,8 +1320,9 @@ class LSGBartDecoder(LSGBartPretrainedModel):
     """
 
     def __init__(self, config, embed_tokens=None):
-
-        super().__init__(config)
+        
+        LSGBartPretrainedModel.__init__(self, config)
+        
         self.dropout = config.dropout
         self.layerdrop = config.decoder_layerdrop
         self.padding_idx = config.pad_token_id
@@ -1360,202 +1347,6 @@ class LSGBartDecoder(LSGBartPretrainedModel):
         # Initialize weights and apply final processing
         self.post_init()
 
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
-        # create causal mask
-        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-        combined_attention_mask = None
-        if input_shape[-1] > 1:
-            combined_attention_mask = _make_causal_mask(
-                input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length
-            ).to(self.device)
-
-        if attention_mask is not None:
-            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
-            combined_attention_mask = (
-                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
-            )
-
-        return combined_attention_mask
-
-    def resize_inputs(self, inputs_embeds, attention_mask):
-        pad = 0
-        
-        max_len = int(attention_mask.sum(dim=-1).max())
-        pad = attention_mask.size()[-1] - max_len
-        inputs_embeds = inputs_embeds[:, :max_len]
-        attention_mask = attention_mask[..., :max_len]
-        return pad, inputs_embeds, attention_mask
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        encoder_hidden_states=None,
-        encoder_attention_mask=None,
-        head_mask=None,
-        cross_attn_head_mask=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        return_dict=None,
-        ):
-
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        # retrieve input_ids and inputs_embeds
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
-        elif input_ids is not None:
-            input_shape = input_ids.size()
-            input_ids = input_ids.view(-1, input_shape[-1])
-        elif inputs_embeds is not None:
-            input_shape = inputs_embeds.size()[:-1]
-        else:
-            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
-
-        # past_key_values_length
-        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
-
-        # Resize to reduce computation
-        pad = 0
-        if self.adaptive:
-            if attention_mask is not None:
-                pad, inputs_embeds, attention_mask = self.resize_inputs(inputs_embeds, attention_mask)
-                input_shape = inputs_embeds.size()[:-1]
-            if encoder_attention_mask is not None:
-                _, encoder_hidden_states, encoder_attention_mask = self.resize_inputs(encoder_hidden_states, encoder_attention_mask)
-
-        attention_mask = self._prepare_decoder_attention_mask(
-            attention_mask, input_shape, inputs_embeds, past_key_values_length
-        )
-
-        # expand encoder attention mask
-        if encoder_hidden_states is not None and encoder_attention_mask is not None:
-            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            encoder_attention_mask = _expand_mask(encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
-
-        # embed positions
-        positions = self.embed_positions(input_shape, past_key_values_length)
-
-        hidden_states = inputs_embeds + positions
-        hidden_states = self.layernorm_embedding(hidden_states)
-
-        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
-        next_decoder_cache = () if use_cache else None
-
-        # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired
-        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
-            if attn_mask is not None:
-                if attn_mask.size()[0] != (len(self.layers)):
-                    raise ValueError(
-                        "The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
-                    )
-
-        for idx, decoder_layer in enumerate(self.layers):
-            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-            dropout_probability = random.uniform(0, 1)
-            if self.training and (dropout_probability < self.layerdrop):
-                continue
-
-            past_key_value = past_key_values[idx] if past_key_values is not None else None
-
-            if self.gradient_checkpointing and self.training:
-
-                if use_cache:
-                    logger.warning(
-                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                    )
-                    use_cache = False
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs, output_attentions, use_cache)
-
-                    return custom_forward
-
-                layer_outputs = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(decoder_layer),
-                    hidden_states,
-                    attention_mask,
-                    encoder_hidden_states,
-                    encoder_attention_mask,
-                    head_mask[idx] if head_mask is not None else None,
-                    cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
-                    None,
-                )
-            else:
-
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=encoder_attention_mask,
-                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
-                    cross_attn_layer_head_mask=(
-                        cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
-                    ),
-                    past_key_value=past_key_value,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                )
-            hidden_states = layer_outputs[0]
-
-            if use_cache:
-                next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-                if encoder_hidden_states is not None:
-                    all_cross_attentions += (layer_outputs[2],)
-
-        # Resize to original shape
-        hidden_states = torch.nn.functional.pad(hidden_states.transpose(-1, -2), pad=(0, pad), value=0).transpose(-1, -2)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        next_cache = next_decoder_cache if use_cache else None
-        if not return_dict:
-            return tuple(
-                v
-                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]
-                if v is not None
-            )
-        return BaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-            cross_attentions=all_cross_attentions,
-        )
-
 
 class LSGBartModel(LSGBartPretrainedModel):