output-hidden-states (#56)

- add option to output hidden states (681845da60fc92b982a091b193ae1fac09bc61f8)

modeling_xlm_roberta.py

@@ -22,13 +22,13 @@ import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers import AutoTokenizer, PretrainedConfig
-from transformers.modeling_outputs import (MaskedLMOutput,
-                                           SequenceClassifierOutput)
+from transformers.modeling_outputs import MaskedLMOutput, SequenceClassifierOutput
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.bert.modeling_bert import (
-    BaseModelOutputWithPoolingAndCrossAttentions, BertForPreTrainingOutput)
-from transformers.models.xlm_roberta.modeling_xlm_roberta import \
-    XLMRobertaLMHead
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    BertForPreTrainingOutput,
+)
+from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaLMHead
 
 from .rotary import RotaryEmbedding
 from .block import Block
@@ -195,17 +195,30 @@ class XLMRobertaEncoder(nn.Module):
         self._grad_checkpointing = value
 
     def forward(
-        self, hidden_states, key_padding_mask=None, subset_mask=None, adapter_mask=None
+        self,
+        hidden_states,
+        key_padding_mask=None,
+        subset_mask=None,
+        adapter_mask=None,
+        output_hidden_states: Optional[bool] = None,
     ):
         """If subset_mask is not None, we only want output for the subset of the sequence.
         This means that we only compute the last layer output for these tokens.
         subset_mask: (batch, seqlen), dtype=torch.bool
         """
+
+        all_hidden_states = () if output_hidden_states else None
+
+        if output_hidden_states and subset_mask:
+            raise ValueError('output_hidden_states is not supported for subset_masks')
+
         if key_padding_mask is None or not self.use_flash_attn:
             mixer_kwargs = {"adapter_mask": adapter_mask}
             if key_padding_mask is not None:
                 mixer_kwargs["key_padding_mask"] = key_padding_mask.bool()
             for layer in self.layers:
+                if output_hidden_states:
+                    all_hidden_states = all_hidden_states + (hidden_states,)
                 if self._grad_checkpointing:
                     hidden_states = torch.utils.checkpoint.checkpoint(
                         layer,
@@ -215,10 +228,14 @@ class XLMRobertaEncoder(nn.Module):
                     )
                 else:
                     hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
             if subset_mask is not None:
                 hidden_states = hidden_states[subset_mask]
         else:
             batch, seqlen = hidden_states.shape[:2]
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
             hidden_states, indices, cu_seqlens, max_seqlen_in_batch, cu_adapter_mask = (
                 unpad_input(hidden_states, key_padding_mask, adapter_mask)
             )
@@ -239,6 +256,10 @@ class XLMRobertaEncoder(nn.Module):
                         )
                     else:
                         hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
+                    if output_hidden_states:
+                        all_hidden_states = all_hidden_states + (
+                            pad_input(hidden_states, indices, batch, seqlen),
+                        )
                 hidden_states = pad_input(hidden_states, indices, batch, seqlen)
             else:
                 for layer in self.layers[:-1]:
@@ -291,7 +312,7 @@ class XLMRobertaEncoder(nn.Module):
                     hidden_states = self.layers[-1](
                         hidden_states_subset, mixer_kwargs=mixer_kwargs
                     )
-        return hidden_states
+        return all_hidden_states if output_hidden_states else hidden_states
 
 
 class XLMRobertaPooler(nn.Module):
@@ -588,7 +609,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
                     embeddings = self.mean_pooling(
                         token_embs, encoded_input["attention_mask"]
                     )
-            
+
             all_embeddings.extend(embeddings)
 
         all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
@@ -596,9 +617,12 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         truncate_dim = truncate_dim or self.config.truncate_dim
         if truncate_dim:
             all_embeddings = self.truncate_embeddings(all_embeddings, truncate_dim)
-        
+
         if normalize_embeddings:
-            all_embeddings = [torch.nn.functional.normalize(embedding, p=2, dim=0) for embedding in all_embeddings]
+            all_embeddings = [
+                torch.nn.functional.normalize(embedding, p=2, dim=0)
+                for embedding in all_embeddings
+            ]
 
         if convert_to_tensor:
             all_embeddings = torch.stack(all_embeddings)
@@ -659,6 +683,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         attention_mask=None,
         masked_tokens_mask=None,
         return_dict=None,
+        output_hidden_states=None,
         **kwargs,
     ):
         """If masked_tokens_mask is not None (i.e. last_layer_subset == True in XLMForPreTraining),
@@ -711,8 +736,15 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             key_padding_mask=attention_mask,
             subset_mask=subset_mask,
             adapter_mask=adapter_mask,
+            output_hidden_states=output_hidden_states,
         )
 
+        if output_hidden_states:
+            all_hidden_states = sequence_output
+            sequence_output = sequence_output[-1]
+        else:
+            all_hidden_states = None
+
         if masked_tokens_mask is None:
             pooled_output = (
                 self.pooler(sequence_output, adapter_mask=adapter_mask)
@@ -742,6 +774,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         return BaseModelOutputWithPoolingAndCrossAttentions(
             last_hidden_state=sequence_output,
             pooler_output=pooled_output,
+            hidden_states=all_hidden_states,
         )