sunzeyeah commited on
Commit
c4c2f24
1 Parent(s): 77f733c
config.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/root/autodl-tmp//Data/models/pangu-350M",
3
+ "activation_function": "gelu",
4
+ "architectures": [
5
+ "GPTPanguForCausalLM"
6
+ ],
7
+ "attn_pdrop": 0.1,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_gptpangu.GPTPanguConfig",
10
+ "AutoModelForCausalLM": "modeling_gptpangu.GPTPanguForCausalLM",
11
+ "AutoTokenizer": [
12
+ "tokenization_gptpangu.GPTPanguTokenizer",
13
+ null
14
+ ]
15
+ },
16
+ "embd_pdrop": 0.1,
17
+ "hidden_size": 1024,
18
+ "initializer_range": 0.02,
19
+ "intermediate_size": null,
20
+ "layer_norm_epsilon": 1e-05,
21
+ "max_position_embeddings": 1024,
22
+ "model_type": "gpt_pangu",
23
+ "num_heads": 16,
24
+ "num_layers": 24,
25
+ "resid_pdrop": 0.1,
26
+ "scale_attn_weights": true,
27
+ "summary_activation": null,
28
+ "summary_first_dropout": 0.1,
29
+ "summary_proj_to_labels": true,
30
+ "summary_type": "cls_index",
31
+ "summary_use_proj": true,
32
+ "tokenizer_class": "GPTPanguTokenizer",
33
+ "torch_dtype": "float32",
34
+ "transformers_version": "4.26.1",
35
+ "use_cache": false,
36
+ "vocab_size": 40000
37
+ }
configuration_gptpangu.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers.configuration_utils import PretrainedConfig
2
+
3
+
4
+ class GPTPanguConfig(PretrainedConfig):
5
+ model_type = "gpt_pangu"
6
+ keys_to_ignore_at_inference = ["past_key_values"]
7
+
8
+ def __init__(
9
+ self,
10
+ vocab_size=40000,
11
+ max_position_embeddings=1024,
12
+ hidden_size=1024,
13
+ intermediate_size=None,
14
+ num_layers=24,
15
+ num_heads=16,
16
+ activation_function="gelu",
17
+ resid_pdrop=0.1,
18
+ embd_pdrop=0.1,
19
+ attn_pdrop=0.1,
20
+ layer_norm_epsilon=1e-5,
21
+ scale_attn_weights=True,
22
+ initializer_range=0.02,
23
+ summary_type="cls_index",
24
+ summary_use_proj=True,
25
+ summary_activation=None,
26
+ summary_proj_to_labels=True,
27
+ summary_first_dropout=0.1,
28
+ use_cache=True,
29
+ # bos_token_id=9,
30
+ # eos_token_id=9,
31
+ **kwargs,
32
+ ):
33
+ self.vocab_size = vocab_size
34
+ self.max_position_embeddings = max_position_embeddings
35
+ self.hidden_size = hidden_size
36
+ self.intermediate_size = intermediate_size
37
+ self.num_layers = num_layers
38
+ self.num_heads = num_heads
39
+ self.activation_function = activation_function
40
+ self.resid_pdrop = resid_pdrop
41
+ self.embd_pdrop = embd_pdrop
42
+ self.attn_pdrop = attn_pdrop
43
+ self.layer_norm_epsilon = layer_norm_epsilon
44
+ self.scale_attn_weights = scale_attn_weights
45
+ self.initializer_range = initializer_range
46
+ self.summary_type = summary_type
47
+ self.summary_use_proj = summary_use_proj
48
+ self.summary_activation = summary_activation
49
+ self.summary_first_dropout = summary_first_dropout
50
+ self.summary_proj_to_labels = summary_proj_to_labels
51
+ self.use_cache = use_cache
52
+
53
+ # self.bos_token_id = bos_token_id
54
+ # self.eos_token_id = eos_token_id
55
+
56
+ super().__init__(**kwargs)
generation_config.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "transformers_version": "4.26.1",
4
+ "use_cache": false
5
+ }
modeling_gptpangu.py ADDED
@@ -0,0 +1,549 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """PyTorch PanguAlpha GPT2 Model"""
2
+ # from .configuration_gptpangu import GPTPanguConfig
3
+
4
+ from typing import Tuple
5
+ import math
6
+
7
+ import torch
8
+ from torch import nn
9
+
10
+ from transformers.activations import ACT2FN
11
+ from transformers.modeling_utils import PreTrainedModel
12
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
13
+
14
+ from transformers.utils import logging
15
+
16
+ logger = logging.get_logger(__name__)
17
+
18
+
19
+ class GPTPanguAttention(nn.Module):
20
+ def __init__(self, config):
21
+ super().__init__()
22
+
23
+ max_positions = config.max_position_embeddings
24
+ self.register_buffer(
25
+ "bias",
26
+ torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
27
+ 1, 1, max_positions, max_positions
28
+ ),
29
+ )
30
+ self.register_buffer("masked_bias", torch.tensor(-1e4))
31
+
32
+ self.embed_dim = config.hidden_size
33
+ self.num_heads = config.num_heads
34
+ self.head_dim = self.embed_dim // self.num_heads
35
+ if self.head_dim * self.num_heads != self.embed_dim:
36
+ raise ValueError(
37
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
38
+ )
39
+
40
+ self.scale_attn_weights = config.scale_attn_weights
41
+
42
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
43
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
44
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
45
+ self.c_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
46
+
47
+ self.attn_dropout = nn.Dropout(config.attn_pdrop)
48
+ self.resid_dropout = nn.Dropout(config.resid_pdrop)
49
+
50
+
51
+ def _attn(self, query, key, value, attention_mask=None, head_mask=None):
52
+ attn_weights = torch.matmul(query, key.transpose(-1, -2))
53
+
54
+ if self.scale_attn_weights:
55
+ attn_weights = attn_weights / (float(value.size(-1)) ** 0.5)
56
+
57
+ query_length, key_length = query.size(-2), key.size(-2)
58
+ causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
59
+ attn_weights = torch.where(causal_mask, attn_weights, self.masked_bias.to(attn_weights.dtype))
60
+
61
+ if attention_mask is not None:
62
+ # Apply the attention mask
63
+ attn_weights = attn_weights + attention_mask
64
+
65
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
66
+
67
+ # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise
68
+ attn_weights = attn_weights.type(value.dtype)
69
+ attn_weights = self.attn_dropout(attn_weights)
70
+
71
+ # Mask heads if we want to
72
+ if head_mask is not None:
73
+ attn_weights = attn_weights * head_mask
74
+
75
+ attn_output = torch.matmul(attn_weights, value)
76
+
77
+ return attn_output, attn_weights
78
+
79
+ def _split_heads(self, tensor, num_heads, attn_head_size):
80
+ """
81
+ Splits hidden_size dim into attn_head_size and num_heads
82
+ """
83
+ new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
84
+ tensor = tensor.view(*new_shape)
85
+ return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)
86
+
87
+ def _merge_heads(self, tensor, num_heads, attn_head_size):
88
+ """
89
+ Merges attn_head_size dim and num_attn_heads dim into hidden_size
90
+ """
91
+ tensor = tensor.permute(0, 2, 1, 3).contiguous()
92
+ new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
93
+ return tensor.view(new_shape)
94
+
95
+ def forward(
96
+ self,
97
+ hidden_states,
98
+ layer_past=None,
99
+ attention_mask=None,
100
+ head_mask=None,
101
+ custom_query=None,
102
+ use_cache=False,
103
+ output_attentions=False,
104
+ ):
105
+ query = self.q_proj(custom_query) if custom_query is not None else self.q_proj(hidden_states)
106
+ key = self.k_proj(hidden_states)
107
+ value = self.v_proj(hidden_states)
108
+
109
+ query = self._split_heads(query, self.num_heads, self.head_dim)
110
+ key = self._split_heads(key, self.num_heads, self.head_dim)
111
+ value = self._split_heads(value, self.num_heads, self.head_dim)
112
+
113
+ if layer_past is not None:
114
+ past_key, past_value = layer_past
115
+ key = torch.cat((past_key, key), dim=-2)
116
+ value = torch.cat((past_value, value), dim=-2)
117
+
118
+ if use_cache is True:
119
+ present = (key, value)
120
+ else:
121
+ present = None
122
+
123
+ attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
124
+
125
+ attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
126
+ attn_output = self.c_proj(attn_output)
127
+ attn_output = self.resid_dropout(attn_output)
128
+
129
+ outputs = (attn_output, present)
130
+ if output_attentions:
131
+ outputs += (attn_weights,)
132
+
133
+ return outputs # a, present, (attentions)
134
+
135
+
136
+ class GPTPanguMLP(nn.Module):
137
+ def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * hidden_size
138
+ super().__init__()
139
+ embed_dim = config.hidden_size
140
+ self.c_fc = nn.Linear(embed_dim, intermediate_size)
141
+ self.c_proj = nn.Linear(intermediate_size, embed_dim)
142
+ self.act = ACT2FN[config.activation_function]
143
+ self.dropout = nn.Dropout(config.resid_pdrop)
144
+
145
+ def forward(self, hidden_states):
146
+ hidden_states = self.c_fc(hidden_states)
147
+ hidden_states = self.act(hidden_states)
148
+ hidden_states = self.c_proj(hidden_states)
149
+ hidden_states = self.dropout(hidden_states)
150
+ return hidden_states
151
+
152
+
153
+ class GPTPanguBlock(nn.Module):
154
+ def __init__(self, config):
155
+ super().__init__()
156
+ hidden_size = config.hidden_size
157
+ inner_dim = config.intermediate_size if config.intermediate_size is not None else 4 * hidden_size
158
+
159
+ self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
160
+ self.attn = GPTPanguAttention(config)
161
+ self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
162
+ self.mlp = GPTPanguMLP(inner_dim, config)
163
+
164
+ def forward(
165
+ self,
166
+ hidden_states,
167
+ layer_past=None,
168
+ attention_mask=None,
169
+ head_mask=None,
170
+ custom_query=None,
171
+ use_cache=False,
172
+ output_attentions=False,
173
+ ):
174
+ residual = hidden_states
175
+ hidden_states = self.ln_1(hidden_states)
176
+ attn_outputs = self.attn(
177
+ hidden_states,
178
+ layer_past=layer_past,
179
+ attention_mask=attention_mask,
180
+ head_mask=head_mask,
181
+ custom_query=custom_query,
182
+ use_cache=use_cache,
183
+ output_attentions=output_attentions,
184
+ )
185
+ attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
186
+ outputs = attn_outputs[1:]
187
+ # residual connection
188
+ hidden_states = attn_output + residual
189
+
190
+ residual = hidden_states
191
+ hidden_states = self.ln_2(hidden_states)
192
+ feed_forward_hidden_states = self.mlp(hidden_states)
193
+ # residual connection
194
+ hidden_states = residual + feed_forward_hidden_states
195
+
196
+ if use_cache:
197
+ outputs = (hidden_states,) + outputs
198
+ else:
199
+ outputs = (hidden_states,) + outputs[1:]
200
+
201
+ return outputs # hidden_states, present, (attentions, cross_attentions)
202
+
203
+
204
+ class GPTPanguPreTrainedModel(PreTrainedModel):
205
+ """
206
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
207
+ models.
208
+ """
209
+
210
+ # config_class = GPTPanguConfig
211
+ base_model_prefix = "transformer"
212
+ supports_gradient_checkpointing = True
213
+
214
+ def __init__(self, *inputs, **kwargs):
215
+ super().__init__(*inputs, **kwargs)
216
+
217
+ def _init_weights(self, module):
218
+ """Initialize the weights."""
219
+ if isinstance(module, (nn.Linear,)):
220
+ # Slightly different from the TF version which uses truncated_normal for initialization
221
+ # cf https://github.com/pytorch/pytorch/pull/5617
222
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
223
+ if module.bias is not None:
224
+ module.bias.data.zero_()
225
+ elif isinstance(module, nn.Embedding):
226
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
227
+ if module.padding_idx is not None:
228
+ module.weight.data[module.padding_idx].zero_()
229
+ elif isinstance(module, nn.LayerNorm):
230
+ module.bias.data.zero_()
231
+ module.weight.data.fill_(1.0)
232
+
233
+ # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
234
+ # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
235
+ # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
236
+ # > -- GPT-2 :: https://openai.com/blog/better-language-models/
237
+ #
238
+ # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
239
+ for name, p in module.named_parameters():
240
+ if "c_proj" in name and "weight" in name:
241
+ # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
242
+ p.data.normal_(mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.num_layers)))
243
+
244
+ def _set_gradient_checkpointing(self, module, value=False):
245
+ if isinstance(module, GPTPanguModel):
246
+ module.gradient_checkpointing = value
247
+
248
+
249
+ class GPTPanguModel(GPTPanguPreTrainedModel):
250
+ def __init__(self, config):
251
+ super().__init__(config)
252
+
253
+ self.embed_dim = config.hidden_size
254
+
255
+ self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
256
+ self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
257
+ self.wqe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
258
+
259
+ self.drop = nn.Dropout(config.embd_pdrop)
260
+ self.h = nn.ModuleList([GPTPanguBlock(config) for _ in range(config.num_layers)])
261
+ self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
262
+
263
+ self.gradient_checkpointing = False
264
+ # Initialize weights and apply final processing
265
+ self.post_init()
266
+
267
+ def get_input_embeddings(self):
268
+ return self.wte
269
+
270
+ def set_input_embeddings(self, new_embeddings):
271
+ self.wte = new_embeddings
272
+
273
+ def forward(
274
+ self,
275
+ input_ids=None,
276
+ past_key_values=None,
277
+ attention_mask=None,
278
+ token_type_ids=None,
279
+ position_ids=None,
280
+ head_mask=None,
281
+ inputs_embeds=None,
282
+ use_cache=None,
283
+ output_attentions=None,
284
+ output_hidden_states=None,
285
+ return_dict=None,
286
+ ):
287
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
288
+ output_hidden_states = (
289
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
290
+ )
291
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
292
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
293
+
294
+ if input_ids is not None and inputs_embeds is not None:
295
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
296
+ elif input_ids is not None:
297
+ input_shape = input_ids.size()
298
+ input_ids = input_ids.view(-1, input_shape[-1])
299
+ batch_size = input_ids.shape[0]
300
+ elif inputs_embeds is not None:
301
+ input_shape = inputs_embeds.size()[:-1]
302
+ batch_size = inputs_embeds.shape[0]
303
+ else:
304
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
305
+
306
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
307
+
308
+ if token_type_ids is not None:
309
+ token_type_ids = token_type_ids.view(-1, input_shape[-1])
310
+ if position_ids is not None:
311
+ position_ids = position_ids.view(-1, input_shape[-1])
312
+
313
+ if past_key_values is None:
314
+ past_length = 0
315
+ past_key_values = tuple([None] * len(self.h))
316
+ else:
317
+ past_length = past_key_values[0][0].size(-2)
318
+ if position_ids is None:
319
+ position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
320
+ position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
321
+
322
+ # GPT2Attention mask.
323
+ if attention_mask is not None:
324
+ if batch_size <= 0:
325
+ raise ValueError("batch_size has to be defined and > 0")
326
+ attention_mask = attention_mask.view(batch_size, -1)
327
+ # We create a 3D attention mask from a 2D tensor mask.
328
+ # Sizes are [batch_size, 1, 1, to_seq_length]
329
+ # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
330
+ # this attention mask is more simple than the triangular masking of causal attention
331
+ # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
332
+ attention_mask = attention_mask[:, None, None, :]
333
+
334
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
335
+ # masked positions, this operation will create a tensor which is 0.0 for
336
+ # positions we want to attend and -10000.0 for masked positions.
337
+ # Since we are adding it to the raw scores before the softmax, this is
338
+ # effectively the same as removing these entirely.
339
+ attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
340
+ attention_mask = (1.0 - attention_mask) * -10000.0
341
+
342
+ # Prepare head mask if needed
343
+ # 1.0 in head_mask indicate we keep the head
344
+ # attention_probs has shape bsz x num_heads x N x N
345
+ # head_mask has shape n_layer x batch x num_heads x N x N
346
+ head_mask = self.get_head_mask(head_mask, self.config.num_layers)
347
+
348
+ if inputs_embeds is None:
349
+ inputs_embeds = self.wte(input_ids)
350
+ position_embeds = self.wpe(position_ids)
351
+ hidden_states = inputs_embeds + position_embeds
352
+
353
+ if token_type_ids is not None:
354
+ token_type_embeds = self.wte(token_type_ids)
355
+ hidden_states = hidden_states + token_type_embeds
356
+
357
+ hidden_states = self.drop(hidden_states)
358
+
359
+ output_shape = input_shape + (hidden_states.size(-1),)
360
+
361
+ # top attention custom query
362
+ last_layer_id = len(self.h) - 1
363
+ query_embeds = self.wqe(position_ids)
364
+
365
+ presents = () if use_cache else None
366
+ all_self_attentions = () if output_attentions else None
367
+ all_hidden_states = () if output_hidden_states else None
368
+ for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
369
+ # Final LayerNorm before last query layer
370
+ if i == last_layer_id:
371
+ hidden_states = self.ln_f(hidden_states)
372
+
373
+ if output_hidden_states:
374
+ all_hidden_states = all_hidden_states + (hidden_states,)
375
+
376
+ if self.gradient_checkpointing and self.training:
377
+
378
+ if use_cache:
379
+ logger.warning(
380
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
381
+ )
382
+ use_cache = False
383
+
384
+ def create_custom_forward(module):
385
+ def custom_forward(*inputs):
386
+ # None for past_key_value
387
+ return module(*inputs, use_cache, output_attentions)
388
+
389
+ return custom_forward
390
+
391
+ outputs = torch.utils.checkpoint.checkpoint(
392
+ create_custom_forward(block),
393
+ hidden_states=hidden_states,
394
+ layer_past=None,
395
+ attention_mask=attention_mask,
396
+ head_mask=head_mask[i],
397
+ # custom query
398
+ custom_query=query_embeds if i == last_layer_id else None,
399
+ )
400
+ else:
401
+ outputs = block(
402
+ hidden_states,
403
+ layer_past=layer_past,
404
+ attention_mask=attention_mask,
405
+ head_mask=head_mask[i],
406
+ # custom query
407
+ custom_query=query_embeds if i == last_layer_id else None,
408
+ use_cache=use_cache,
409
+ output_attentions=output_attentions,
410
+ )
411
+
412
+ hidden_states = outputs[0]
413
+ if use_cache is True:
414
+ presents = presents + (outputs[1],)
415
+
416
+ if output_attentions:
417
+ all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
418
+
419
+ hidden_states = hidden_states.view(*output_shape)
420
+ # Add last hidden state
421
+ if output_hidden_states:
422
+ all_hidden_states = all_hidden_states + (hidden_states,)
423
+
424
+ if not return_dict:
425
+ return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
426
+
427
+ return BaseModelOutputWithPast(
428
+ last_hidden_state=hidden_states,
429
+ past_key_values=presents,
430
+ hidden_states=all_hidden_states,
431
+ attentions=all_self_attentions,
432
+ )
433
+
434
+
435
+ class GPTPanguForCausalLM(GPTPanguPreTrainedModel):
436
+ def __init__(self, config):
437
+ super().__init__(config)
438
+ self.transformer = GPTPanguModel(config)
439
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
440
+
441
+ # Initialize weights and apply final processing
442
+ self.post_init()
443
+
444
+ def get_output_embeddings(self):
445
+ return self.lm_head
446
+
447
+ def set_output_embeddings(self, new_embeddings):
448
+ self.lm_head = new_embeddings
449
+
450
+ def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
451
+ token_type_ids = kwargs.get("token_type_ids", None)
452
+ # only last token for inputs_ids if past is defined in kwargs
453
+ if past:
454
+ input_ids = input_ids[:, -1].unsqueeze(-1)
455
+ if token_type_ids is not None:
456
+ token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
457
+
458
+ attention_mask = kwargs.get("attention_mask", None)
459
+ position_ids = kwargs.get("position_ids", None)
460
+
461
+ if attention_mask is not None and position_ids is None:
462
+ # create position_ids on the fly for batch generation
463
+ position_ids = attention_mask.long().cumsum(-1) - 1
464
+ position_ids.masked_fill_(attention_mask == 0, 1)
465
+ if past:
466
+ position_ids = position_ids[:, -1].unsqueeze(-1)
467
+ else:
468
+ position_ids = None
469
+ return {
470
+ "input_ids": input_ids,
471
+ "past_key_values": past,
472
+ "use_cache": kwargs.get("use_cache"),
473
+ "position_ids": position_ids,
474
+ "attention_mask": attention_mask,
475
+ "token_type_ids": token_type_ids,
476
+ }
477
+
478
+ def forward(
479
+ self,
480
+ input_ids=None,
481
+ past_key_values=None,
482
+ attention_mask=None,
483
+ token_type_ids=None,
484
+ position_ids=None,
485
+ head_mask=None,
486
+ inputs_embeds=None,
487
+ labels=None,
488
+ use_cache=None,
489
+ output_attentions=None,
490
+ output_hidden_states=None,
491
+ return_dict=None,
492
+ ):
493
+ r"""
494
+ labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
495
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
496
+ ``labels = input_ids`` Indices are selected in ``[-100, 0, ..., config.vocab_size]`` All labels set to
497
+ ``-100`` are ignored (masked), the loss is only computed for labels in ``[0, ..., config.vocab_size]``
498
+ """
499
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
500
+
501
+ transformer_outputs = self.transformer(
502
+ input_ids,
503
+ past_key_values=past_key_values,
504
+ attention_mask=attention_mask,
505
+ token_type_ids=token_type_ids,
506
+ position_ids=position_ids,
507
+ head_mask=head_mask,
508
+ inputs_embeds=inputs_embeds,
509
+ use_cache=use_cache,
510
+ output_attentions=output_attentions,
511
+ output_hidden_states=output_hidden_states,
512
+ return_dict=return_dict,
513
+ )
514
+ hidden_states = transformer_outputs[0]
515
+
516
+ lm_logits = self.lm_head(hidden_states)
517
+
518
+ loss = None
519
+ if labels is not None:
520
+ # Shift so that tokens < n predict n
521
+ shift_logits = lm_logits[..., :-1, :].contiguous()
522
+ shift_labels = labels[..., 1:].contiguous()
523
+ # Flatten the tokens
524
+ loss_fct = nn.CrossEntropyLoss()
525
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
526
+
527
+ if not return_dict:
528
+ output = (lm_logits,) + transformer_outputs[1:]
529
+ return ((loss,) + output) if loss is not None else output
530
+
531
+ return CausalLMOutputWithPast(
532
+ loss=loss,
533
+ logits=lm_logits,
534
+ past_key_values=transformer_outputs.past_key_values,
535
+ hidden_states=transformer_outputs.hidden_states,
536
+ attentions=transformer_outputs.attentions,
537
+ )
538
+
539
+ @staticmethod
540
+ def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
541
+ """
542
+ This function is used to re-order the :obj:`past_key_values` cache if
543
+ :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
544
+ called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
545
+ """
546
+ return tuple(
547
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
548
+ for layer_past in past
549
+ )
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0ed712e4f3990a8d902d262dea7aae2ada367f7b1e37e5dc5b979db09854fb6d
3
+ size 1406779991
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:31274bab401fb8fcc7d5aabbb3c7b4481347c7c2d1466ecb0edb1ac397113da5
3
+ size 3567