yuewang-sf commited on
Commit
8300c1c
1 Parent(s): fe5f115

Upload 2 files

Browse files
Files changed (2) hide show
  1. configuration_codet5p.py +113 -0
  2. modeling_codet5p.py +982 -0
configuration_codet5p.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
3
+
4
+ """ CodeT5+ model configuration"""
5
+ from transformers.configuration_utils import PretrainedConfig
6
+ from transformers.utils import logging
7
+ import copy
8
+
9
+ logger = logging.get_logger(__name__)
10
+
11
+
12
+ # Adapted from transformers.models.codegen.configuration_codegen.CodeGenConfig
13
+ class CodeT5pModuleConfig(PretrainedConfig):
14
+ model_type = "codet5p_module"
15
+ attribute_map = {
16
+ "max_position_embeddings": "n_positions",
17
+ "hidden_size": "n_embd",
18
+ "num_attention_heads": "n_head",
19
+ "num_hidden_layers": "n_layer",
20
+ }
21
+
22
+ def __init__(
23
+ self,
24
+ vocab_size=50400,
25
+ n_positions=2048,
26
+ n_ctx=2048,
27
+ n_embd=4096,
28
+ n_layer=28,
29
+ n_head=16,
30
+ rotary_dim=64,
31
+ n_inner=None,
32
+ activation_function="gelu_new",
33
+ resid_pdrop=0.0,
34
+ embd_pdrop=0.0,
35
+ attn_pdrop=0.0,
36
+ layer_norm_epsilon=1e-5,
37
+ initializer_range=0.02,
38
+ scale_attn_weights=True,
39
+ use_cache=True,
40
+ bos_token_id=50256,
41
+ eos_token_id=50256,
42
+ tie_word_embeddings=False,
43
+ **kwargs
44
+ ):
45
+ self.vocab_size = vocab_size
46
+ self.n_ctx = n_ctx
47
+ self.n_positions = n_positions
48
+ self.n_embd = n_embd
49
+ self.n_layer = n_layer
50
+ self.n_head = n_head
51
+ self.n_inner = n_inner
52
+ self.rotary_dim = rotary_dim
53
+ self.activation_function = activation_function
54
+ self.resid_pdrop = resid_pdrop
55
+ self.embd_pdrop = embd_pdrop
56
+ self.attn_pdrop = attn_pdrop
57
+ self.layer_norm_epsilon = layer_norm_epsilon
58
+ self.initializer_range = initializer_range
59
+ self.scale_attn_weights = scale_attn_weights
60
+ self.use_cache = use_cache
61
+
62
+ self.bos_token_id = bos_token_id
63
+ self.eos_token_id = eos_token_id
64
+
65
+ super().__init__(
66
+ bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
67
+ )
68
+
69
+
70
+ # Adapted from transformers.models.encoder_decoder.configuration_encoder_decoder.EncoderDecoderConfig
71
+ class CodeT5pConfig(PretrainedConfig):
72
+ model_type = "codet5p"
73
+ is_composition = True
74
+
75
+ def __init__(self, **kwargs):
76
+ super().__init__(**kwargs)
77
+ assert (
78
+ "encoder" in kwargs and "decoder" in kwargs
79
+ ), "Config has to be initialized with encoder and decoder config"
80
+ encoder_config = kwargs.pop("encoder")
81
+ decoder_config = kwargs.pop("decoder")
82
+ encoder_model_type = encoder_config.pop("model_type")
83
+ decoder_model_type = decoder_config.pop("model_type")
84
+
85
+ if encoder_model_type != decoder_model_type:
86
+ logger.warning("Encoder and decoder model types are different")
87
+
88
+ self.encoder = CodeT5pModuleConfig(**encoder_config)
89
+ self.decoder = CodeT5pModuleConfig(**decoder_config)
90
+ self.is_encoder_decoder = True
91
+
92
+ @classmethod
93
+ def from_encoder_decoder_configs(
94
+ cls, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs
95
+ ) -> PretrainedConfig:
96
+ logger.info("Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config")
97
+ decoder_config.is_decoder = True
98
+ decoder_config.add_cross_attention = True
99
+
100
+ return cls(encoder=encoder_config.to_dict(), decoder=decoder_config.to_dict(), **kwargs)
101
+
102
+ def to_dict(self):
103
+ """
104
+ Serializes this instance to a Python dictionary. Override the default *to_dict()* from *PretrainedConfig*.
105
+
106
+ Returns:
107
+ `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
108
+ """
109
+ output = copy.deepcopy(self.__dict__)
110
+ output["encoder"] = self.encoder.to_dict()
111
+ output["decoder"] = self.decoder.to_dict()
112
+ output["model_type"] = self.__class__.model_type
113
+ return output
modeling_codet5p.py ADDED
@@ -0,0 +1,982 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
3
+ """ PyTorch CodeT5+ 2B 6B 16B models.
4
+ The implementation is mainly based on transformers.models.codegen.modeling_codegen by adding cross-attention
5
+ and transformers.models.encoder_decoder.modeling_encoder_decoder.EncoderDecoderModel.
6
+ """
7
+ from typing import Optional, Tuple, Union
8
+ import torch
9
+ import torch.utils.checkpoint
10
+ from torch import nn
11
+ from torch.nn import CrossEntropyLoss
12
+
13
+ from transformers.activations import ACT2FN
14
+ from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput, \
15
+ BaseModelOutputWithPast, CausalLMOutputWithPast, \
16
+ BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions
17
+ from transformers.modeling_utils import PreTrainedModel
18
+ from transformers.configuration_utils import PretrainedConfig
19
+ from transformers.utils import add_code_sample_docstrings, add_start_docstrings, logging
20
+ from configuration_codet5p import CodeT5pConfig, CodeT5pModuleConfig
21
+
22
+ logger = logging.get_logger(__name__)
23
+
24
+ CODET5P_PRETRAINED_MODEL_ARCHIVE_LIST = [
25
+ "Salesforce/codet5p-220m",
26
+ "Salesforce/codet5p-770m",
27
+ "Salesforce/codet5p-220m-py",
28
+ "Salesforce/codet5p-770m-py",
29
+ "Salesforce/codet5p-2b",
30
+ "Salesforce/codet5p-6b",
31
+ "Salesforce/codet5p-16b",
32
+ "Salesforce/instructcodet5p-16b",
33
+ # See all CodeT5+ models at https://huggingface.co/models?filter=codet5p
34
+ ]
35
+
36
+
37
+ # Copied from transformers.models.gptj.modeling_gptj.fixed_pos_embedding
38
+ def fixed_pos_embedding(x, seq_dim=1, seq_len=None):
39
+ dim = x.shape[-1]
40
+ if seq_len is None:
41
+ seq_len = x.shape[seq_dim]
42
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))
43
+ sinusoid_inp = (
44
+ torch.einsum("i , j -> i j", torch.arange(seq_len, dtype=torch.float), inv_freq).to(x.device).float()
45
+ )
46
+ return torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)
47
+
48
+
49
+ # Copied from transformers.models.gptj.modeling_gptj.rotate_every_two
50
+ def rotate_every_two(x):
51
+ x1 = x[:, :, :, ::2]
52
+ x2 = x[:, :, :, 1::2]
53
+ x = torch.stack((-x2, x1), dim=-1)
54
+ return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)')
55
+
56
+
57
+ # Copied from transformers.models.gptj.modeling_gptj.duplicate_interleave
58
+ def duplicate_interleave(m):
59
+ """
60
+ A simple version of `torch.repeat_interleave` for duplicating a matrix while interleaving the copy.
61
+ """
62
+ dim0 = m.shape[0]
63
+ m = m.view(-1, 1) # flatten the matrix
64
+ m = m.repeat(1, 2) # repeat all elements into the 2nd dimension
65
+ m = m.view(dim0, -1) # reshape into a matrix, interleaving the copy
66
+ return m
67
+
68
+
69
+ # Copied from transformers.models.gptj.modeling_gptj.apply_rotary_pos_emb
70
+ def apply_rotary_pos_emb(x, sincos, offset=0):
71
+ sin, cos = (duplicate_interleave(t)[None, offset: x.shape[1] + offset, None, :] for t in sincos)
72
+ # einsum notation for lambda t: repeat(t[offset:x.shape[1]+offset,:], "n d -> () n () (d j)", j=2)
73
+ return (x * cos) + (rotate_every_two(x) * sin)
74
+
75
+
76
+ # Adapted from transformers.models.codegen.modeling_codegen.CodeGenAttention
77
+ class CodeT5pAttention(nn.Module):
78
+ def __init__(self, config, is_cross_attention=False, is_decoder=True):
79
+ super().__init__()
80
+
81
+ max_positions = config.max_position_embeddings
82
+ self.register_buffer(
83
+ "causal_mask",
84
+ torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
85
+ 1, 1, max_positions, max_positions
86
+ ),
87
+ )
88
+
89
+ self.attn_dropout = nn.Dropout(config.attn_pdrop)
90
+ self.resid_dropout = nn.Dropout(config.resid_pdrop)
91
+
92
+ self.embed_dim = config.hidden_size
93
+ self.num_attention_heads = config.num_attention_heads
94
+ self.head_dim = self.embed_dim // self.num_attention_heads
95
+ if self.head_dim * self.num_attention_heads != self.embed_dim:
96
+ raise ValueError(
97
+ f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
98
+ f" `num_attention_heads`: {self.num_attention_heads})."
99
+ )
100
+
101
+ self.scale_attn = torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype())
102
+ self.is_decoder = is_decoder
103
+ self.is_cross_attention = is_cross_attention
104
+ if self.is_cross_attention:
105
+ self.qkv_proj = nn.Linear(self.embed_dim, self.embed_dim * 2, bias=False)
106
+ self.q_attn = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
107
+ else:
108
+ self.qkv_proj = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=False)
109
+
110
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
111
+ self.rotary_dim = None
112
+ if config.rotary_dim is not None:
113
+ self.rotary_dim = config.rotary_dim
114
+
115
+ def _split_heads(self, x, n_head, dim_head, mp_num):
116
+ reshaped = x.reshape(x.shape[:-1] + (n_head // mp_num, dim_head))
117
+ reshaped = reshaped.reshape(x.shape[:-2] + (-1,) + reshaped.shape[-1:])
118
+ return reshaped
119
+
120
+ def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
121
+ """
122
+ Merges attn_head_size dim and num_attn_heads dim into n_ctx
123
+ """
124
+ if len(tensor.shape) == 5:
125
+ tensor = tensor.permute(0, 1, 3, 2, 4).contiguous()
126
+ elif len(tensor.shape) == 4:
127
+ tensor = tensor.permute(0, 2, 1, 3).contiguous()
128
+ else:
129
+ raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
130
+ new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
131
+ return tensor.view(new_shape)
132
+
133
+ def _attn(
134
+ self,
135
+ query,
136
+ key,
137
+ value,
138
+ attention_mask=None,
139
+ head_mask=None,
140
+ ):
141
+ # Keep the attention weights computation in fp32 to avoid overflow issues
142
+ query = query.to(torch.float32)
143
+ key = key.to(torch.float32)
144
+
145
+ attn_weights = torch.matmul(query, key.transpose(-1, -2))
146
+ attn_weights = attn_weights / self.scale_attn
147
+
148
+ if not self.is_cross_attention and self.is_decoder:
149
+ # compute causal mask from causal mask buffer
150
+ query_length, key_length = query.size(-2), key.size(-2)
151
+ causal_mask = self.causal_mask[:, :, key_length - query_length: key_length, :key_length]
152
+ mask_value = torch.finfo(attn_weights.dtype).min
153
+ # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
154
+ # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
155
+ mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
156
+ attn_weights = torch.where(causal_mask.bool(), attn_weights, mask_value)
157
+
158
+ if attention_mask is not None:
159
+ # Apply the attention mask
160
+ attn_weights = attn_weights + attention_mask
161
+
162
+ attn_weights = nn.Softmax(dim=-1)(attn_weights)
163
+ attn_weights = attn_weights.to(value.dtype)
164
+ attn_weights = self.attn_dropout(attn_weights)
165
+
166
+ # Mask heads if we want to
167
+ if head_mask is not None:
168
+ attn_weights = attn_weights * head_mask
169
+
170
+ attn_output = torch.matmul(attn_weights, value)
171
+
172
+ return attn_output, attn_weights
173
+
174
+ def forward(
175
+ self,
176
+ hidden_states: Optional[torch.FloatTensor],
177
+ attention_mask: Optional[torch.FloatTensor] = None,
178
+ layer_past: Optional[Tuple[torch.Tensor]] = None,
179
+ head_mask: Optional[torch.FloatTensor] = None,
180
+ encoder_hidden_states: Optional[torch.Tensor] = None,
181
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
182
+ use_cache: Optional[bool] = False,
183
+ output_attentions: Optional[bool] = False,
184
+ ) -> Union[
185
+ Tuple[torch.Tensor, Tuple[torch.Tensor]],
186
+ Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]],
187
+ ]:
188
+
189
+ if encoder_hidden_states is not None:
190
+ if not hasattr(self, "q_attn"):
191
+ raise ValueError(
192
+ "If class is used as cross attention, the weights `q_attn` have to be defined. "
193
+ "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
194
+ )
195
+
196
+ mp_num = 4
197
+ local_dim = self.head_dim * self.num_attention_heads // mp_num
198
+ q = self.q_attn(hidden_states)
199
+ q_split = q.reshape(q.shape[:-1] + (mp_num, -1))
200
+ query = torch.split(q_split, local_dim, dim=-1)[0]
201
+
202
+ qkv = self.qkv_proj(encoder_hidden_states)
203
+ qkv_split = qkv.reshape(qkv.shape[:-1] + (mp_num, -1))
204
+ value, key = torch.split(qkv_split, local_dim, dim=-1)
205
+
206
+ attention_mask = encoder_attention_mask
207
+ else:
208
+ qkv = self.qkv_proj(hidden_states)
209
+ mp_num = 4
210
+ qkv_split = qkv.reshape(qkv.shape[:-1] + (mp_num, -1))
211
+
212
+ local_dim = self.head_dim * self.num_attention_heads // mp_num
213
+ query, value, key = torch.split(qkv_split, local_dim, dim=-1)
214
+
215
+ query = self._split_heads(query, self.num_attention_heads, self.head_dim, mp_num=mp_num)
216
+ key = self._split_heads(key, self.num_attention_heads, self.head_dim, mp_num=mp_num)
217
+
218
+ value = self._split_heads(value, self.num_attention_heads, self.head_dim, mp_num=mp_num)
219
+ value = value.permute(0, 2, 1, 3)
220
+
221
+ seq_len = key.shape[1]
222
+ offset = 0
223
+
224
+ if layer_past is not None:
225
+ offset = layer_past[0].shape[-2]
226
+ seq_len += offset
227
+
228
+ if self.rotary_dim is not None:
229
+ k_rot = key[:, :, :, : self.rotary_dim]
230
+ k_pass = key[:, :, :, self.rotary_dim:]
231
+
232
+ q_rot = query[:, :, :, : self.rotary_dim]
233
+ q_pass = query[:, :, :, self.rotary_dim:]
234
+
235
+ sincos = fixed_pos_embedding(k_rot, 1, seq_len=seq_len)
236
+ k_rot = apply_rotary_pos_emb(k_rot, sincos, offset=offset)
237
+ seq_len_q = query.shape[1]
238
+ sincos_q = fixed_pos_embedding(q_rot, 1, seq_len=seq_len_q)
239
+ q_rot = apply_rotary_pos_emb(q_rot, sincos_q, offset=offset)
240
+
241
+ key = torch.cat([k_rot, k_pass], dim=-1)
242
+ query = torch.cat([q_rot, q_pass], dim=-1)
243
+ else:
244
+ sincos = fixed_pos_embedding(key, 1, seq_len=seq_len)
245
+ key = apply_rotary_pos_emb(key, sincos, offset=offset)
246
+ query = apply_rotary_pos_emb(query, sincos, offset=offset)
247
+
248
+ key = key.permute(0, 2, 1, 3)
249
+ query = query.permute(0, 2, 1, 3)
250
+
251
+ if layer_past is not None:
252
+ past_key = layer_past[0]
253
+ past_value = layer_past[1]
254
+ key = torch.cat((past_key, key), dim=-2)
255
+ value = torch.cat((past_value, value), dim=-2)
256
+
257
+ if use_cache is True:
258
+ present = (key, value)
259
+ else:
260
+ present = None
261
+
262
+ # compute self-attention: V x Softmax(QK^T)
263
+ attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
264
+
265
+ attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
266
+ attn_output = self.out_proj(attn_output)
267
+ attn_output = self.resid_dropout(attn_output)
268
+
269
+ outputs = (attn_output, present)
270
+ if output_attentions:
271
+ outputs += (attn_weights,)
272
+
273
+ return outputs # a, present, (attentions)
274
+
275
+
276
+ # Adapted from transformers.models.codegen.modeling_codegen.CodeGenMLP
277
+ class CodeT5pMLP(nn.Module):
278
+ def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * embed_dim
279
+ super().__init__()
280
+ embed_dim = config.n_embd
281
+
282
+ self.fc_in = nn.Linear(embed_dim, intermediate_size)
283
+ self.fc_out = nn.Linear(intermediate_size, embed_dim)
284
+
285
+ self.act = ACT2FN[config.activation_function]
286
+ self.dropout = nn.Dropout(config.resid_pdrop)
287
+
288
+ def forward(self, hidden_states: Optional[torch.FloatTensor]) -> torch.FloatTensor:
289
+ hidden_states = self.fc_in(hidden_states)
290
+ hidden_states = self.act(hidden_states)
291
+ hidden_states = self.fc_out(hidden_states)
292
+ hidden_states = self.dropout(hidden_states)
293
+ return hidden_states
294
+
295
+
296
+ # Adapted from transformers.models.codegen.modeling_codegen.CodeGenBlock
297
+ class CodeT5pBlock(nn.Module):
298
+ def __init__(self, config, layer_idx=None):
299
+ super().__init__()
300
+ inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
301
+ self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
302
+
303
+ if config.is_decoder is False:
304
+ self.attn = CodeT5pAttention(config, is_cross_attention=False, is_decoder=False)
305
+ else:
306
+ self.attn = CodeT5pAttention(config)
307
+ self.mlp = CodeT5pMLP(inner_dim, config)
308
+
309
+ # Adding 1 cross-attention layer at the final decoder layer
310
+ self.add_cross_attention_by_layer = True \
311
+ if config.add_cross_attention and layer_idx == config.n_layer - 1 else False
312
+
313
+ if config.add_cross_attention and self.add_cross_attention_by_layer:
314
+ self.crossattention = CodeT5pAttention(config, is_cross_attention=True)
315
+
316
+ def forward(
317
+ self,
318
+ hidden_states: Optional[torch.FloatTensor],
319
+ layer_past: Optional[Tuple[torch.Tensor]] = None,
320
+ attention_mask: Optional[torch.FloatTensor] = None,
321
+ head_mask: Optional[torch.FloatTensor] = None,
322
+ encoder_hidden_states: Optional[torch.Tensor] = None,
323
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
324
+ use_cache: Optional[bool] = False,
325
+ output_attentions: Optional[bool] = False,
326
+ ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
327
+ residual = hidden_states
328
+ hidden_states = self.ln_1(hidden_states)
329
+ attn_outputs = self.attn(
330
+ hidden_states,
331
+ layer_past=layer_past,
332
+ attention_mask=attention_mask,
333
+ head_mask=head_mask,
334
+ use_cache=use_cache,
335
+ output_attentions=output_attentions,
336
+ )
337
+ attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
338
+ outputs = attn_outputs[1:]
339
+ feed_forward_hidden_states = self.mlp(hidden_states)
340
+
341
+ if encoder_hidden_states is not None and self.add_cross_attention_by_layer:
342
+ # add one self-attention block for cross-attention
343
+ if not hasattr(self, "crossattention"):
344
+ raise ValueError(
345
+ f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
346
+ "cross-attention layers by setting `config.add_cross_attention=True`"
347
+ )
348
+ # residual = hidden_states
349
+ # hidden_states = self.ln_cross_attn(residual)
350
+ cross_attn_outputs = self.crossattention(
351
+ hidden_states,
352
+ attention_mask=attention_mask,
353
+ head_mask=head_mask,
354
+ encoder_hidden_states=encoder_hidden_states,
355
+ encoder_attention_mask=encoder_attention_mask,
356
+ output_attentions=output_attentions,
357
+ )
358
+ xattn_output = cross_attn_outputs[0]
359
+ attn_output = attn_output + xattn_output
360
+ outputs = outputs + cross_attn_outputs[2:] # add cross attentions if we output attention weights
361
+
362
+ hidden_states = attn_output + feed_forward_hidden_states + residual
363
+
364
+ if use_cache:
365
+ outputs = (hidden_states,) + outputs
366
+ else:
367
+ outputs = (hidden_states,) + outputs[1:]
368
+
369
+ return outputs # hidden_states, present, (attentions)
370
+
371
+
372
+ # Adapted from transformers.models.codegen.modeling_codegen.CodeGenPreTrainedModel
373
+ class CodeT5pPreTrainedModel(PreTrainedModel):
374
+ """
375
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
376
+ models.
377
+ """
378
+ config_class = CodeT5pModuleConfig
379
+ base_model_prefix = "transformer"
380
+ supports_gradient_checkpointing = True
381
+ _no_split_modules = ["CodeT5pBlock"]
382
+
383
+ def __init__(self, *inputs, **kwargs):
384
+ super().__init__(*inputs, **kwargs)
385
+
386
+ def _init_weights(self, module):
387
+ """Initialize the weights."""
388
+ if isinstance(module, (nn.Linear,)):
389
+ # Slightly different from Mesh Transformer JAX which uses truncated_normal for initialization
390
+ # cf https://github.com/pytorch/pytorch/pull/5617
391
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
392
+ if module.bias is not None:
393
+ module.bias.data.zero_()
394
+ elif isinstance(module, nn.Embedding):
395
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
396
+ if module.padding_idx is not None:
397
+ module.weight.data[module.padding_idx].zero_()
398
+ elif isinstance(module, nn.LayerNorm):
399
+ module.bias.data.zero_()
400
+ module.weight.data.fill_(1.0)
401
+
402
+ def _set_gradient_checkpointing(self, module, value=False):
403
+ if isinstance(module, CodeT5pModel):
404
+ module.gradient_checkpointing = value
405
+
406
+
407
+ # Adapted from transformers.models.codegen.modeling_codegen.CodeGenModel
408
+ class CodeT5pModel(CodeT5pPreTrainedModel):
409
+ def __init__(self, config):
410
+ super().__init__(config)
411
+
412
+ self.embed_dim = config.n_embd
413
+ self.vocab_size = config.vocab_size
414
+ self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
415
+ self.drop = nn.Dropout(config.embd_pdrop)
416
+ self.h = nn.ModuleList([CodeT5pBlock(config, idx) for idx in range(config.n_layer)])
417
+ self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
418
+ self.rotary_dim = min(config.rotary_dim, config.n_ctx // config.num_attention_heads)
419
+
420
+ self.gradient_checkpointing = False
421
+
422
+ # Initialize weights and apply final processing
423
+ self.post_init()
424
+
425
+ def get_input_embeddings(self):
426
+ return self.wte
427
+
428
+ def set_input_embeddings(self, new_embeddings):
429
+ self.wte = new_embeddings
430
+
431
+ def forward(
432
+ self,
433
+ input_ids: Optional[torch.LongTensor] = None,
434
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
435
+ attention_mask: Optional[torch.FloatTensor] = None,
436
+ token_type_ids: Optional[torch.LongTensor] = None,
437
+ position_ids: Optional[torch.LongTensor] = None,
438
+ head_mask: Optional[torch.FloatTensor] = None,
439
+ inputs_embeds: Optional[torch.FloatTensor] = None,
440
+ encoder_hidden_states: Optional[torch.Tensor] = None,
441
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
442
+ use_cache: Optional[bool] = None,
443
+ output_attentions: Optional[bool] = None,
444
+ output_hidden_states: Optional[bool] = None,
445
+ return_dict: Optional[bool] = None,
446
+ ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
447
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
448
+ output_hidden_states = (
449
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
450
+ )
451
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
452
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
453
+
454
+ if input_ids is not None and inputs_embeds is not None:
455
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
456
+ elif input_ids is not None:
457
+ input_shape = input_ids.size()
458
+ input_ids = input_ids.view(-1, input_shape[-1])
459
+ batch_size = input_ids.shape[0]
460
+ elif inputs_embeds is not None:
461
+ input_shape = inputs_embeds.size()[:-1]
462
+ batch_size = inputs_embeds.shape[0]
463
+ else:
464
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
465
+
466
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
467
+
468
+ if token_type_ids is not None:
469
+ token_type_ids = token_type_ids.view(-1, input_shape[-1])
470
+
471
+ if position_ids is not None:
472
+ position_ids = position_ids.view(-1, input_shape[-1])
473
+
474
+ if past_key_values is None:
475
+ past_length = 0
476
+ past_key_values = tuple([None] * len(self.h))
477
+ else:
478
+ past_length = past_key_values[0][0].size(-2)
479
+
480
+ if position_ids is None:
481
+ position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
482
+ position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
483
+
484
+ # Attention mask.
485
+ if attention_mask is not None:
486
+ if batch_size <= 0:
487
+ raise ValueError("batch_size has to be defined and > 0")
488
+ attention_mask = attention_mask.view(batch_size, -1)
489
+ # We create a 3D attention mask from a 2D tensor mask.
490
+ # Sizes are [batch_size, 1, 1, to_seq_length]
491
+ # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
492
+ # this attention mask is more simple than the triangular masking of causal attention
493
+ # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
494
+ attention_mask = attention_mask[:, None, None, :]
495
+
496
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
497
+ # masked positions, this operation will create a tensor which is 0.0 for
498
+ # positions we want to attend and the dtype's smallest value for masked positions.
499
+ # Since we are adding it to the raw scores before the softmax, this is
500
+ # effectively the same as removing these entirely.
501
+ attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
502
+ attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
503
+
504
+ # If a 2D or 3D attention mask is provided for the cross-attention
505
+ # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
506
+ if self.config.add_cross_attention and encoder_hidden_states is not None:
507
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
508
+ encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
509
+ if encoder_attention_mask is None:
510
+ encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
511
+ encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
512
+ else:
513
+ encoder_attention_mask = None
514
+
515
+ # Prepare head mask if needed
516
+ # 1.0 in head_mask indicate we keep the head
517
+ # attention_probs has shape bsz x num_attention_heads x N x N
518
+ # head_mask has shape n_layer x batch x num_attention_heads x N x N
519
+ head_mask = self.get_head_mask(head_mask, self.config.n_layer)
520
+
521
+ if inputs_embeds is None:
522
+ inputs_embeds = self.wte(input_ids)
523
+
524
+ hidden_states = inputs_embeds
525
+
526
+ if token_type_ids is not None:
527
+ token_type_embeds = self.wte(token_type_ids)
528
+ hidden_states = hidden_states + token_type_embeds
529
+
530
+ hidden_states = self.drop(hidden_states)
531
+
532
+ output_shape = input_shape + (hidden_states.size(-1),)
533
+
534
+ presents = () if use_cache else None
535
+ all_self_attentions = () if output_attentions else None
536
+ all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
537
+ all_hidden_states = () if output_hidden_states else None
538
+ for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
539
+ if output_hidden_states:
540
+ all_hidden_states = all_hidden_states + (hidden_states,)
541
+
542
+ if self.gradient_checkpointing and self.training:
543
+ if use_cache:
544
+ logger.warning(
545
+ "`use_cache=True` is incompatible with `config.gradient_checkpointing=True`. Setting "
546
+ "`use_cache=False`..."
547
+ )
548
+ use_cache = False
549
+
550
+ def create_custom_forward(module):
551
+ def custom_forward(*inputs):
552
+ # None for past_key_value
553
+ return module(*inputs, use_cache, output_attentions)
554
+
555
+ return custom_forward
556
+
557
+ outputs = torch.utils.checkpoint.checkpoint(
558
+ create_custom_forward(block),
559
+ hidden_states,
560
+ None,
561
+ attention_mask,
562
+ head_mask[i],
563
+ encoder_hidden_states,
564
+ encoder_attention_mask,
565
+ )
566
+ else:
567
+ outputs = block(
568
+ hidden_states,
569
+ layer_past=layer_past,
570
+ attention_mask=attention_mask,
571
+ head_mask=head_mask[i],
572
+ encoder_hidden_states=encoder_hidden_states,
573
+ encoder_attention_mask=encoder_attention_mask,
574
+ use_cache=use_cache,
575
+ output_attentions=output_attentions,
576
+ )
577
+
578
+ hidden_states = outputs[0]
579
+ if use_cache is True:
580
+ presents = presents + (outputs[1],)
581
+
582
+ if output_attentions:
583
+ all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
584
+ if self.config.add_cross_attention and self.add_cross_attention_by_layer:
585
+ all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)
586
+
587
+ hidden_states = self.ln_f(hidden_states)
588
+
589
+ hidden_states = hidden_states.view(output_shape)
590
+ # Add last hidden state
591
+ if output_hidden_states:
592
+ all_hidden_states = all_hidden_states + (hidden_states,)
593
+
594
+ if not return_dict:
595
+ return tuple(
596
+ v for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions] if
597
+ v is not None)
598
+
599
+ return BaseModelOutputWithPastAndCrossAttentions(
600
+ last_hidden_state=hidden_states,
601
+ past_key_values=presents,
602
+ hidden_states=all_hidden_states,
603
+ attentions=all_self_attentions,
604
+ cross_attentions=all_cross_attentions,
605
+ )
606
+
607
+
608
+ # Adapted from transformers.models.codegen.modeling_codegen.CodeGenForCausalLM
609
+ class CodeT5pForCausalLM(CodeT5pPreTrainedModel):
610
+ _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.causal_mask"]
611
+
612
+ def __init__(self, config):
613
+ super().__init__(config)
614
+ self.transformer = CodeT5pModel(config)
615
+ self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
616
+
617
+ # Initialize weights and apply final processing
618
+ self.post_init()
619
+
620
+ def get_output_embeddings(self):
621
+ return self.lm_head
622
+
623
+ def set_output_embeddings(self, new_embeddings):
624
+ self.lm_head = new_embeddings
625
+
626
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
627
+ token_type_ids = kwargs.get("token_type_ids", None)
628
+ # only last token for inputs_ids if past is defined in kwargs
629
+ if past_key_values:
630
+ input_ids = input_ids[:, -1].unsqueeze(-1)
631
+ if token_type_ids is not None:
632
+ token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
633
+
634
+ attention_mask = kwargs.get("attention_mask", None)
635
+ position_ids = kwargs.get("position_ids", None)
636
+
637
+ if attention_mask is not None and position_ids is None:
638
+ # create position_ids on the fly for batch generation
639
+ position_ids = attention_mask.long().cumsum(-1) - 1
640
+ position_ids.masked_fill_(attention_mask == 0, 1)
641
+ if past_key_values:
642
+ position_ids = position_ids[:, -1].unsqueeze(-1)
643
+ else:
644
+ position_ids = None
645
+ return {
646
+ "input_ids": input_ids,
647
+ "past_key_values": past_key_values,
648
+ "use_cache": kwargs.get("use_cache"),
649
+ "position_ids": position_ids,
650
+ "attention_mask": attention_mask,
651
+ "token_type_ids": token_type_ids,
652
+ }
653
+
654
+ def forward(
655
+ self,
656
+ input_ids: Optional[torch.LongTensor] = None,
657
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
658
+ attention_mask: Optional[torch.FloatTensor] = None,
659
+ token_type_ids: Optional[torch.LongTensor] = None,
660
+ position_ids: Optional[torch.LongTensor] = None,
661
+ head_mask: Optional[torch.FloatTensor] = None,
662
+ inputs_embeds: Optional[torch.FloatTensor] = None,
663
+ encoder_hidden_states: Optional[torch.Tensor] = None,
664
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
665
+ labels: Optional[torch.LongTensor] = None,
666
+ use_cache: Optional[bool] = None,
667
+ output_attentions: Optional[bool] = None,
668
+ output_hidden_states: Optional[bool] = None,
669
+ return_dict: Optional[bool] = None,
670
+ ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
671
+ r"""
672
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
673
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
674
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
675
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
676
+ """
677
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
678
+
679
+ transformer_outputs = self.transformer(
680
+ input_ids,
681
+ past_key_values=past_key_values,
682
+ attention_mask=attention_mask,
683
+ token_type_ids=token_type_ids,
684
+ position_ids=position_ids,
685
+ head_mask=head_mask,
686
+ inputs_embeds=inputs_embeds,
687
+ encoder_hidden_states=encoder_hidden_states,
688
+ encoder_attention_mask=encoder_attention_mask,
689
+ use_cache=use_cache,
690
+ output_attentions=output_attentions,
691
+ output_hidden_states=output_hidden_states,
692
+ return_dict=return_dict,
693
+ )
694
+ hidden_states = transformer_outputs[0]
695
+
696
+ # make sure sampling in fp16 works correctly and
697
+ # compute loss in fp32 to match with mesh-tf version
698
+ # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
699
+ lm_logits = self.lm_head(hidden_states).to(torch.float32)
700
+
701
+ loss = None
702
+ if labels is not None:
703
+ # Shift so that tokens < n predict n
704
+ shift_logits = lm_logits[..., :-1, :].contiguous()
705
+ shift_labels = labels[..., 1:].contiguous()
706
+ # Flatten the tokens
707
+ loss_fct = CrossEntropyLoss()
708
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
709
+
710
+ loss = loss.to(hidden_states.dtype)
711
+
712
+ if not return_dict:
713
+ output = (lm_logits,) + transformer_outputs[1:]
714
+ return ((loss,) + output) if loss is not None else output
715
+
716
+ return CausalLMOutputWithCrossAttentions(
717
+ loss=loss,
718
+ logits=lm_logits,
719
+ past_key_values=transformer_outputs.past_key_values,
720
+ hidden_states=transformer_outputs.hidden_states,
721
+ attentions=transformer_outputs.attentions,
722
+ cross_attentions=transformer_outputs.cross_attentions,
723
+ )
724
+
725
+ @staticmethod
726
+ def _reorder_cache(
727
+ past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
728
+ ) -> Tuple[Tuple[torch.Tensor]]:
729
+ """
730
+ This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
731
+ [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
732
+ beam_idx at every generation step.
733
+ """
734
+ return tuple(
735
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
736
+ for layer_past in past_key_values
737
+ )
738
+
739
+
740
+ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
741
+ """
742
+ Shift input ids one token to the right.
743
+ """
744
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
745
+ shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
746
+ if decoder_start_token_id is None:
747
+ raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
748
+ shifted_input_ids[:, 0] = decoder_start_token_id
749
+
750
+ if pad_token_id is None:
751
+ raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
752
+ # replace possible -100 values in labels by `pad_token_id`
753
+ shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
754
+
755
+ return shifted_input_ids
756
+
757
+
758
+ # Adapted from transformers.models.encoder_decoder.modeling_encoder_decoder.EncoderDecoderModel
759
+ class CodeT5pEncoderDecoderModel(PreTrainedModel):
760
+ config_class = CodeT5pConfig
761
+
762
+ def __init__(
763
+ self,
764
+ config: Optional[PretrainedConfig] = None,
765
+ encoder: Optional[PreTrainedModel] = None,
766
+ decoder: Optional[PreTrainedModel] = None,
767
+ ):
768
+ if config is None and (encoder is None or decoder is None):
769
+ raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")
770
+ if config is None:
771
+ config = CodeT5pConfig.from_encoder_decoder_configs(encoder.config, decoder.config)
772
+ else:
773
+ if not isinstance(config, self.config_class):
774
+ raise ValueError(f"Config: {config} has to be of type {self.config_class}")
775
+
776
+ if config.decoder.cross_attention_hidden_size is not None:
777
+ if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
778
+ raise ValueError(
779
+ "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"
780
+ f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
781
+ f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"
782
+ " `config.encoder.hidden_size`."
783
+ )
784
+
785
+ # initialize with config
786
+ super().__init__(config)
787
+
788
+ if encoder is None:
789
+ encoder = CodeT5pModel(config.encoder)
790
+
791
+ if decoder is None:
792
+ decoder = CodeT5pForCausalLM(config.decoder)
793
+
794
+ self.encoder = encoder
795
+ self.decoder = decoder
796
+
797
+ if self.encoder.config.to_dict() != self.config.encoder.to_dict():
798
+ logger.warning(
799
+ f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"
800
+ f" {self.config.encoder}"
801
+ )
802
+ if self.decoder.config.to_dict() != self.config.decoder.to_dict():
803
+ logger.warning(
804
+ f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
805
+ f" {self.config.decoder}"
806
+ )
807
+
808
+ # make sure that the individual model's config refers to the shared config
809
+ # so that the updates to the config will be synced
810
+ self.encoder.config = self.config.encoder
811
+ self.decoder.config = self.config.decoder
812
+
813
+ # encoder outputs might need to be projected to different dimension for decoder
814
+ if (
815
+ self.encoder.config.hidden_size != self.decoder.config.hidden_size
816
+ and self.decoder.config.cross_attention_hidden_size is None
817
+ ):
818
+ self.enc_to_dec_proj = nn.Linear(self.encoder.config.hidden_size, self.decoder.config.hidden_size)
819
+
820
+ if self.encoder.get_output_embeddings() is not None:
821
+ raise ValueError(
822
+ f"The encoder {self.encoder} should not have a LM Head. Please use a model without LM Head"
823
+ )
824
+ # tie encoder, decoder weights if config set accordingly
825
+ self.tie_weights()
826
+
827
+ def tie_weights(self):
828
+ # tie encoder & decoder if needed
829
+ if self.config.tie_encoder_decoder:
830
+ # tie encoder and decoder base model
831
+ decoder_base_model_prefix = self.decoder.base_model_prefix
832
+ self._tie_encoder_decoder_weights(
833
+ self.encoder, self.decoder._modules[decoder_base_model_prefix], self.decoder.base_model_prefix
834
+ )
835
+
836
+ def get_encoder(self):
837
+ return self.encoder
838
+
839
+ def get_decoder(self):
840
+ return self.decoder
841
+
842
+ def get_input_embeddings(self):
843
+ return self.encoder.get_input_embeddings()
844
+
845
+ def get_output_embeddings(self):
846
+ return self.decoder.get_output_embeddings()
847
+
848
+ def set_output_embeddings(self, new_embeddings):
849
+ return self.decoder.set_output_embeddings(new_embeddings)
850
+
851
+ @classmethod
852
+ def from_pretrained(cls, *args, **kwargs):
853
+ # At the moment fast initialization is not supported for composite models
854
+ if kwargs.get("_fast_init", False):
855
+ logger.warning(
856
+ "Fast initialization is currently not supported for EncoderDecoderModel. "
857
+ "Falling back to slow initialization..."
858
+ )
859
+ kwargs["_fast_init"] = False
860
+ return super().from_pretrained(*args, **kwargs)
861
+
862
+ def forward(
863
+ self,
864
+ input_ids: Optional[torch.LongTensor] = None,
865
+ attention_mask: Optional[torch.FloatTensor] = None,
866
+ decoder_input_ids: Optional[torch.LongTensor] = None,
867
+ decoder_attention_mask: Optional[torch.BoolTensor] = None,
868
+ encoder_outputs: Optional[Tuple[torch.FloatTensor]] = None,
869
+ past_key_values: Tuple[Tuple[torch.FloatTensor]] = None,
870
+ inputs_embeds: Optional[torch.FloatTensor] = None,
871
+ decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
872
+ labels: Optional[torch.LongTensor] = None,
873
+ use_cache: Optional[bool] = None,
874
+ output_attentions: Optional[bool] = None,
875
+ output_hidden_states: Optional[bool] = None,
876
+ return_dict: Optional[bool] = None,
877
+ **kwargs,
878
+ ) -> Union[Tuple, Seq2SeqLMOutput]:
879
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
880
+
881
+ kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
882
+
883
+ kwargs_decoder = {
884
+ argument[len("decoder_"):]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
885
+ }
886
+
887
+ if encoder_outputs is None:
888
+ encoder_outputs = self.encoder(
889
+ input_ids=input_ids,
890
+ attention_mask=attention_mask,
891
+ inputs_embeds=inputs_embeds,
892
+ output_attentions=output_attentions,
893
+ output_hidden_states=output_hidden_states,
894
+ return_dict=return_dict,
895
+ **kwargs_encoder,
896
+ )
897
+ elif isinstance(encoder_outputs, tuple):
898
+ encoder_outputs = BaseModelOutput(*encoder_outputs)
899
+
900
+ encoder_hidden_states = encoder_outputs[0]
901
+
902
+ # optionally project encoder_hidden_states
903
+ if (
904
+ self.encoder.config.hidden_size != self.decoder.config.hidden_size
905
+ and self.decoder.config.cross_attention_hidden_size is None
906
+ ):
907
+ encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
908
+
909
+ if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
910
+ decoder_input_ids = shift_tokens_right(
911
+ labels, self.config.pad_token_id, self.config.decoder_start_token_id
912
+ )
913
+
914
+ # Decode
915
+ decoder_outputs = self.decoder(
916
+ input_ids=decoder_input_ids,
917
+ attention_mask=decoder_attention_mask,
918
+ encoder_hidden_states=encoder_hidden_states,
919
+ encoder_attention_mask=attention_mask,
920
+ inputs_embeds=decoder_inputs_embeds,
921
+ output_attentions=output_attentions,
922
+ output_hidden_states=output_hidden_states,
923
+ use_cache=use_cache,
924
+ past_key_values=past_key_values,
925
+ return_dict=return_dict,
926
+ **kwargs_decoder,
927
+ )
928
+
929
+ # Compute loss independent from decoder (as some shift the logits inside them)
930
+ loss = None
931
+ if labels is not None:
932
+ # warnings.warn(DEPRECATION_WARNING, FutureWarning)
933
+ logits = decoder_outputs.logits if return_dict else decoder_outputs[0]
934
+ loss_fct = CrossEntropyLoss()
935
+ loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1))
936
+
937
+ if not return_dict:
938
+ if loss is not None:
939
+ return (loss,) + decoder_outputs + encoder_outputs
940
+ else:
941
+ return decoder_outputs + encoder_outputs
942
+
943
+ return Seq2SeqLMOutput(
944
+ loss=loss,
945
+ logits=decoder_outputs.logits,
946
+ past_key_values=decoder_outputs.past_key_values,
947
+ decoder_hidden_states=decoder_outputs.hidden_states,
948
+ decoder_attentions=decoder_outputs.attentions,
949
+ cross_attentions=decoder_outputs.cross_attentions,
950
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
951
+ encoder_hidden_states=encoder_outputs.hidden_states,
952
+ encoder_attentions=encoder_outputs.attentions,
953
+ )
954
+
955
+ def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
956
+ return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
957
+
958
+ def prepare_inputs_for_generation(
959
+ self, input_ids, past=None, attention_mask=None, use_cache=None, encoder_outputs=None, **kwargs
960
+ ):
961
+ decoder_inputs = self.decoder.prepare_inputs_for_generation(input_ids, past=past)
962
+ decoder_attention_mask = decoder_inputs["attention_mask"] if "attention_mask" in decoder_inputs else None
963
+ input_dict = {
964
+ "attention_mask": attention_mask,
965
+ "decoder_attention_mask": decoder_attention_mask,
966
+ "decoder_input_ids": decoder_inputs["input_ids"],
967
+ "encoder_outputs": encoder_outputs,
968
+ "past_key_values": decoder_inputs["past_key_values"],
969
+ "use_cache": use_cache,
970
+ }
971
+ return input_dict
972
+
973
+ def resize_token_embeddings(self, *args, **kwargs):
974
+ raise NotImplementedError(
975
+ "Resizing the embedding layers via the EncoderDecoderModel directly is not supported. Please use the"
976
+ " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"
977
+ " model.decoder.resize_token_embeddings(...))"
978
+ )
979
+
980
+ def _reorder_cache(self, past, beam_idx):
981
+ # apply decoder cache reordering here
982
+ return self.decoder._reorder_cache(past, beam_idx)