luoxindi commited on
Commit
6e92f5b
1 Parent(s): b087ce4

Delete llama2_7B/modeling_llama_copy.py

Browse files
Files changed (1) hide show
  1. llama2_7B/modeling_llama_copy.py +0 -1360
llama2_7B/modeling_llama_copy.py DELETED
@@ -1,1360 +0,0 @@
1
- # coding=utf-8
2
- # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
- #
4
- # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
- # and OPT implementations in this library. It has been modified from its
6
- # original forms to accommodate minor architectural differences compared
7
- # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
- #
9
- # Licensed under the Apache License, Version 2.0 (the "License");
10
- # you may not use this file except in compliance with the License.
11
- # You may obtain a copy of the License at
12
- #
13
- # http://www.apache.org/licenses/LICENSE-2.0
14
- #
15
- # Unless required by applicable law or agreed to in writing, software
16
- # distributed under the License is distributed on an "AS IS" BASIS,
17
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
- # See the License for the specific language governing permissions and
19
- # limitations under the License.
20
- """ PyTorch LLaMA model."""
21
- import math
22
- from typing import List, Optional, Tuple, Union
23
- import numpy as np
24
- import torch
25
- import torch.nn.functional as F
26
- import torch.utils.checkpoint
27
- from torch import nn
28
- from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
29
-
30
- from ...activations import ACT2FN
31
- from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
32
- from ...modeling_utils import PreTrainedModel
33
- from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
34
- from .configuration_llama import LlamaConfig
35
-
36
-
37
- logger = logging.get_logger(__name__)
38
-
39
- _CONFIG_FOR_DOC = "LlamaConfig"
40
-
41
-
42
- # Copied from transformers.models.bart.modeling_bart._make_causal_mask
43
- def _make_causal_mask(
44
- input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
45
- ):
46
- """
47
- Make causal mask used for bi-directional self-attention.
48
- """
49
- bsz, tgt_len = input_ids_shape
50
- mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
51
- mask_cond = torch.arange(mask.size(-1), device=device)
52
- mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
53
- mask = mask.to(dtype)
54
-
55
- if past_key_values_length > 0:
56
- mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
57
- return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
58
-
59
-
60
- # Copied from transformers.models.bart.modeling_bart._expand_mask
61
- def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
62
- """
63
- Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
64
- """
65
- bsz, src_len = mask.size()
66
- tgt_len = tgt_len if tgt_len is not None else src_len
67
-
68
- expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
69
-
70
- inverted_mask = 1.0 - expanded_mask
71
-
72
- return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
73
-
74
-
75
- class LlamaRMSNorm(nn.Module):
76
- def __init__(self, hidden_size, eps=1e-6):
77
- """
78
- LlamaRMSNorm is equivalent to T5LayerNorm
79
- """
80
- super().__init__()
81
- self.weight = nn.Parameter(torch.ones(hidden_size))
82
- self.variance_epsilon = eps
83
-
84
- def forward(self, hidden_states):
85
- input_dtype = hidden_states.dtype
86
- hidden_states = hidden_states.to(torch.float32)
87
- variance = hidden_states.pow(2).mean(-1, keepdim=True)
88
- hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
89
- return self.weight * hidden_states.to(input_dtype)
90
-
91
-
92
- class LlamaRotaryEmbedding(torch.nn.Module):
93
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
94
- super().__init__()
95
-
96
- self.dim = dim
97
- self.max_position_embeddings = max_position_embeddings
98
- self.base = base
99
- inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
100
- self.register_buffer("inv_freq", inv_freq, persistent=False)
101
-
102
- # Build here to make `torch.jit.trace` work.
103
- self._set_cos_sin_cache(
104
- seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
105
- )
106
-
107
- def _set_cos_sin_cache(self, seq_len, device, dtype):
108
- self.max_seq_len_cached = seq_len
109
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
110
-
111
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
112
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
113
- emb = torch.cat((freqs, freqs), dim=-1)
114
- self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
115
- self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
116
-
117
- def forward(self, x, seq_len=None):
118
- # x: [bs, num_attention_heads, seq_len, head_size]
119
- if seq_len > self.max_seq_len_cached:
120
- self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
121
-
122
- return (
123
- self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
124
- self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
125
- )
126
-
127
-
128
- class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
129
- """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
130
-
131
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
132
- self.scaling_factor = scaling_factor
133
- super().__init__(dim, max_position_embeddings, base, device)
134
-
135
- def _set_cos_sin_cache(self, seq_len, device, dtype):
136
- self.max_seq_len_cached = seq_len
137
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
138
- t = t / self.scaling_factor
139
-
140
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
141
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
142
- emb = torch.cat((freqs, freqs), dim=-1)
143
- self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
144
- self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
145
-
146
-
147
- class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):
148
- """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
149
-
150
- def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
151
- self.scaling_factor = scaling_factor
152
- super().__init__(dim, max_position_embeddings, base, device)
153
-
154
- def _set_cos_sin_cache(self, seq_len, device, dtype):
155
- self.max_seq_len_cached = seq_len
156
-
157
- if seq_len > self.max_position_embeddings:
158
- base = self.base * (
159
- (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
160
- ) ** (self.dim / (self.dim - 2))
161
- inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
162
- self.register_buffer("inv_freq", inv_freq, persistent=False)
163
-
164
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
165
-
166
- freqs = torch.einsum("i,j->ij", t, self.inv_freq)
167
- # Different from paper, but it uses a different permutation in order to obtain the same calculation
168
- emb = torch.cat((freqs, freqs), dim=-1)
169
- self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
170
- self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
171
-
172
-
173
- def rotate_half(x):
174
- """Rotates half the hidden dims of the input."""
175
- x1 = x[..., : x.shape[-1] // 2]
176
- x2 = x[..., x.shape[-1] // 2 :]
177
- return torch.cat((-x2, x1), dim=-1)
178
-
179
-
180
- def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
181
- # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
182
- cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
183
- sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
184
- cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
185
- sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
186
- q_embed = (q * cos) + (rotate_half(q) * sin)
187
- k_embed = (k * cos) + (rotate_half(k) * sin)
188
- return q_embed, k_embed
189
-
190
-
191
- class LlamaMLP(nn.Module):
192
- def __init__(self, config):
193
- super().__init__()
194
- self.config = config
195
- self.hidden_size = config.hidden_size
196
- self.intermediate_size = config.intermediate_size
197
- self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
198
- self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
199
- self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
200
- self.act_fn = ACT2FN[config.hidden_act]
201
-
202
- def forward(self, x):
203
- if self.config.pretraining_tp > 1:
204
- slice = self.intermediate_size // self.config.pretraining_tp
205
- gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
206
- up_proj_slices = self.up_proj.weight.split(slice, dim=0)
207
- down_proj_slices = self.down_proj.weight.split(slice, dim=1)
208
-
209
- gate_proj = torch.cat(
210
- [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
211
- )
212
- up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)
213
-
214
- intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
215
- down_proj = [
216
- F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
217
- ]
218
- down_proj = sum(down_proj)
219
- else:
220
- down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
221
-
222
- return down_proj
223
-
224
-
225
- def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
226
- """
227
- This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
228
- num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
229
- """
230
- batch, num_key_value_heads, slen, head_dim = hidden_states.shape
231
- if n_rep == 1:
232
- return hidden_states
233
- hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
234
- return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
235
-
236
-
237
- class LlamaAttention(nn.Module):
238
- """Multi-headed attention from 'Attention Is All You Need' paper"""
239
-
240
- def __init__(self, config: LlamaConfig):
241
- super().__init__()
242
- self.config = config
243
- self.hidden_size = config.hidden_size
244
- self.num_heads = config.num_attention_heads
245
- self.head_dim = self.hidden_size // self.num_heads
246
- self.num_key_value_heads = config.num_key_value_heads
247
- self.num_key_value_groups = self.num_heads // self.num_key_value_heads
248
- self.max_position_embeddings = config.max_position_embeddings
249
-
250
- if (self.head_dim * self.num_heads) != self.hidden_size:
251
- raise ValueError(
252
- f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
253
- f" and `num_heads`: {self.num_heads})."
254
- )
255
- self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
256
- self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
257
- self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
258
- self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
259
- self._init_rope()
260
-
261
- def _init_rope(self):
262
- if self.config.rope_scaling is None:
263
- self.rotary_emb = LlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
264
- else:
265
- scaling_type = self.config.rope_scaling["type"]
266
- scaling_factor = self.config.rope_scaling["factor"]
267
- if scaling_type == "linear":
268
- self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
269
- self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor
270
- )
271
- elif scaling_type == "dynamic":
272
- self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
273
- self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor
274
- )
275
- else:
276
- raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
277
-
278
- def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
279
- return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
280
-
281
- def forward(
282
- self,
283
- hidden_states: torch.Tensor,
284
- attention_mask: Optional[torch.Tensor] = None,
285
- position_ids: Optional[torch.LongTensor] = None,
286
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
287
- output_attentions: bool = False,
288
- use_cache: bool = False,
289
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
290
- bsz, q_len, _ = hidden_states.size()
291
-
292
- if self.config.pretraining_tp > 1:
293
- key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
294
- query_slices = self.q_proj.weight.split(
295
- (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
296
- )
297
- key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
298
- value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
299
-
300
- query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]
301
- query_states = torch.cat(query_states, dim=-1)
302
-
303
- key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
304
- key_states = torch.cat(key_states, dim=-1)
305
-
306
- value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
307
- value_states = torch.cat(value_states, dim=-1)
308
-
309
- else:
310
- query_states = self.q_proj(hidden_states)
311
- key_states = self.k_proj(hidden_states)
312
- value_states = self.v_proj(hidden_states)
313
-
314
- query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
315
- key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
316
- value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
317
-
318
- kv_seq_len = key_states.shape[-2]
319
- if past_key_value is not None:
320
- kv_seq_len += past_key_value[0].shape[-2]
321
- cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
322
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
323
-
324
- if past_key_value is not None:
325
- # reuse k, v, self_attention
326
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
327
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
328
-
329
- past_key_value = (key_states, value_states) if use_cache else None
330
-
331
- # repeat k/v heads if n_kv_heads < n_heads
332
- key_states = repeat_kv(key_states, self.num_key_value_groups)
333
- value_states = repeat_kv(value_states, self.num_key_value_groups)
334
-
335
- attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
336
-
337
- if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
338
- raise ValueError(
339
- f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
340
- f" {attn_weights.size()}"
341
- )
342
-
343
- if attention_mask is not None:
344
- if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
345
- raise ValueError(
346
- f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
347
- )
348
- attn_weights = attn_weights + attention_mask
349
-
350
- # upcast attention to fp32
351
- attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
352
- attn_output = torch.matmul(attn_weights, value_states)
353
-
354
- if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
355
- raise ValueError(
356
- f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
357
- f" {attn_output.size()}"
358
- )
359
-
360
- attn_output = attn_output.transpose(1, 2).contiguous()
361
- attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
362
-
363
- if self.config.pretraining_tp > 1:
364
- attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
365
- o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
366
- attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
367
- else:
368
- attn_output = self.o_proj(attn_output)
369
-
370
- if not output_attentions:
371
- attn_weights = None
372
-
373
- return attn_output, attn_weights, past_key_value
374
-
375
- class KGMLP(nn.Module):
376
- def __init__(
377
- self,
378
- hidden_size: int,
379
- intermediate_size: int,
380
- output_size: int,
381
- hidden_act: str,
382
- ):
383
- super().__init__()
384
- self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
385
- self.down_proj = nn.Linear(intermediate_size, output_size, bias=False)
386
- self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
387
- self.act_fn = ACT2FN[hidden_act]
388
-
389
- def forward(self, x):
390
- return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
391
-
392
-
393
- class LlamaDecoderLayer_1(nn.Module):
394
- def __init__(self, idx, config: LlamaConfig):
395
- super().__init__()
396
- self.hidden_size = config.hidden_size
397
- self.self_attn = LlamaAttention(config=config)
398
- self.mlp = LlamaMLP(config)
399
- self.store_value_mha = 0
400
- self.store_value_ffn = 0
401
- self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
402
- self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
403
-
404
- def init(self):
405
- pass
406
-
407
- def forward(
408
- self,
409
- hidden_states: torch.Tensor,
410
- words_ents_list = None,
411
- words_subtoken_map = None,
412
- attention_mask: Optional[torch.Tensor] = None,
413
- position_ids: Optional[torch.LongTensor] = None,
414
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
415
- output_attentions: Optional[bool] = False,
416
- use_cache: Optional[bool] = False,
417
- idx = None
418
- ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
419
- """
420
- Args:
421
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
422
- attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
423
- `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
424
- output_attentions (`bool`, *optional*):
425
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
426
- returned tensors for more detail.
427
- use_cache (`bool`, *optional*):
428
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
429
- (see `past_key_values`).
430
- past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
431
- """
432
-
433
- residual = hidden_states
434
- pre = hidden_states[-1][-1]
435
- hidden_states = self.input_layernorm(hidden_states)
436
- # Self Attention
437
- hidden_states, self_attn_weights, present_key_value = self.self_attn(
438
- hidden_states=hidden_states,
439
- attention_mask=attention_mask,
440
- position_ids=position_ids,
441
- past_key_value=past_key_value,
442
- output_attentions=output_attentions,
443
- use_cache=use_cache,
444
- )
445
- post = hidden_states[-1][-1]
446
- if hidden_states.size()[1] != 1:
447
- self.store_value_mha += (torch.cosine_similarity(pre, post, dim=0).item())
448
- hidden_states = residual + hidden_states
449
-
450
- # Fully Connected
451
- residual = hidden_states
452
- pre = hidden_states[-1][-1]
453
- hidden_states = self.post_attention_layernorm(hidden_states)
454
- hidden_states = self.mlp(hidden_states)
455
- post = hidden_states[-1][-1]
456
- if hidden_states.size()[1] != 1:
457
- self.store_value_ffn += (torch.cosine_similarity(pre, post, dim=0).item())
458
- hidden_states = residual + hidden_states
459
-
460
- outputs = (hidden_states,)
461
-
462
- if output_attentions:
463
- outputs += (self_attn_weights,)
464
-
465
- if use_cache:
466
- outputs += (present_key_value,)
467
-
468
- return outputs
469
-
470
-
471
-
472
- class LlamaDecoderLayer_2(nn.Module):
473
- def __init__(self, idx, config: LlamaConfig):
474
- super().__init__()
475
- self.config = config
476
- self.hidden_size = config.hidden_size
477
- self.self_attn = LlamaAttention(config=config)
478
- self.KG_infuded_module = None
479
- self.mlp = LlamaMLP(
480
- config
481
- )
482
- self.store_value_mha = 0
483
- self.store_value_ffn = 0
484
- self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
485
- self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
486
-
487
- def init(self):
488
- self.KG_infuded_module = KG_infuded_module(self.config)
489
- #self.KG_infuded_module = self.KG_infuded_module.half()
490
- self.KG_infuded_module = self.KG_infuded_module.cuda()
491
-
492
- def forward(
493
- self,
494
- hidden_states: torch.Tensor,
495
- attention_mask: Optional[torch.Tensor] = None,
496
- words_ents_list = None,
497
- words_subtoken_map = None,
498
- position_ids: Optional[torch.LongTensor] = None,
499
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
500
- output_attentions: Optional[bool] = False,
501
- use_cache: Optional[bool] = False,
502
- idx = None
503
- ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
504
- """
505
- Args:
506
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
507
- attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
508
- `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
509
- output_attentions (`bool`, *optional*):
510
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
511
- returned tensors for more detail.
512
- use_cache (`bool`, *optional*):
513
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
514
- (see `past_key_values`).
515
- past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
516
- """
517
-
518
- residual = hidden_states
519
- pre = hidden_states[-1][-1]
520
- hidden_states = self.input_layernorm(hidden_states)
521
- # Self Attention
522
- hidden_states, self_attn_weights, present_key_value = self.self_attn(
523
- hidden_states=hidden_states,
524
- attention_mask=attention_mask,
525
- position_ids=position_ids,
526
- past_key_value=past_key_value,
527
- output_attentions=output_attentions,
528
- use_cache=use_cache,
529
- )
530
- post = hidden_states[-1][-1]
531
- if hidden_states.size()[1] != 1:
532
- self.store_value_mha += (torch.cosine_similarity(pre, post, dim=0).item())
533
- hidden_states = residual + hidden_states
534
-
535
- residual = hidden_states
536
- pre = hidden_states[-1][-1]
537
- hidden_states = self.post_attention_layernorm(hidden_states)
538
- hidden_states = self.mlp(hidden_states)
539
- post = hidden_states[-1][-1]
540
- if hidden_states.size()[1] != 1:
541
- self.store_value_ffn += (torch.cosine_similarity(pre, post, dim=0).item())
542
- hidden_states = residual + hidden_states
543
- if (idx + 1) % 32 == 0:
544
- hidden_states = self.KG_infuded_module(hidden_states, words_ents_list, words_subtoken_map, None)
545
- outputs = (hidden_states,)
546
-
547
- if output_attentions:
548
- outputs += (self_attn_weights,)
549
-
550
- if use_cache:
551
- outputs += (present_key_value,)
552
-
553
- return outputs
554
-
555
- def forward_ffd(self, hidden_states):
556
- # Fully Connected
557
- residual = hidden_states[0]
558
- hidden_states[0] = self.post_attention_layernorm(hidden_states[0])
559
- hidden_states[0] = self.mlp(hidden_states[0])
560
- hidden_states[0] = residual + hidden_states[0]
561
-
562
- return hidden_states
563
-
564
-
565
- class KG_infuded_module(nn.Module):
566
- def __init__(self, config: LlamaConfig):
567
- super().__init__()
568
- embedding_path = "/data1/xdluo/alpaca-lora-main/data/kgs/wn_concept2vec.txt"
569
- self.concept_embed = None
570
- self.interlayer = 100
571
- self.knowledge_sentinel = nn.Embedding(1, self.interlayer).cuda()
572
- self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
573
- #self.convert_matrix_token = KGMLP(
574
- # hidden_size=config.hidden_size,
575
- # intermediate_size=2048,
576
- # output_size=self.interlayer,
577
- # hidden_act=config.hidden_act,
578
- #).cuda()
579
- self.convert_matrix_entity = KGMLP(
580
- hidden_size=100,
581
- intermediate_size=1024,
582
- output_size=4096,
583
- hidden_act=config.hidden_act,
584
- ).cuda()
585
- #self.convert_matrix_token = nn.Linear(config.hidden_size, self.interlayer, bias = False).cuda()
586
- #self.convert_matrix_entity = nn.Linear(dim, self.interlayer, bias = False).cuda()
587
- #self.convert_matrix = nn.Linear(config.hidden_size, dim, bias = False).cuda()
588
- self.MLP = nn.Linear(config.hidden_size + self.interlayer, config.hidden_size).cuda()
589
- self.act_fn = ACT2FN[config.hidden_act]
590
- self.alpha = nn.Parameter(torch.Tensor([0.5]))
591
- self.dim = self.interlayer
592
-
593
- def init(self, embedding_path):
594
- #embedding_path = "/data1/xdluo/alpaca-lora-main/data/kgs/wn_concept2vec.txt"
595
- #embedding_path = "/data1/xdluo/alpaca-lora-main/data/kgs/conceptnet/ent.npy"
596
- #embedding_path = "/data1/xdluo/alpaca-lora-main/data/kgs/wn18/rescal_embeds.npy"
597
- #embedding_path = "/data1/xdluo/Knower/output/wi_ent_embeds.npy"
598
- #embedding_path = embedding_path
599
- self.id2concept, self.concept2id, embedding_mat, (concept_size, dim) = self.read_conceptnet_embedding(embedding_path)
600
- self.concept_embed = nn.Embedding.from_pretrained(torch.from_numpy(embedding_mat)).cuda()
601
- #self.alpha = nn.Parameter(torch.Tensor([0.5]))
602
- #torch.nn.init.xavier_uniform_(self.convert_matrix_token.weight)
603
- #torch.nn.init.xavier_uniform_(self.convert_matrix_entity.weight)
604
- torch.nn.init.xavier_uniform_(self.knowledge_sentinel.weight)
605
-
606
- def read_concept_embedding(self, embedding_path):
607
- fin = open(embedding_path, encoding='utf-8')
608
- info = [line.strip() for line in fin]
609
- dim = len(info[0].split(' ')[1:])
610
- n_concept = len(info)
611
- embedding_mat = []
612
- id2concept, concept2id = [], {}
613
- # add padding concept into vocab
614
- #embedding_mat.append([0.0 for _ in range(dim)])
615
- for line in info:
616
- concept_name = line.split(' ')[0]
617
- embedding = [float(value_str) for value_str in line.split(' ')[1:]]
618
- assert len(embedding) == dim and not np.any(np.isnan(embedding))
619
- embedding_mat.append(embedding)
620
- concept2id[concept_name] = len(id2concept)
621
- id2concept.append(concept_name)
622
- embedding_mat = np.array(embedding_mat, dtype=np.float32)
623
- fin.close()
624
- return id2concept, concept2id, embedding_mat, (n_concept, dim)
625
-
626
- def read_conceptnet_embedding(self, embedding_path):
627
- ar_load = np.load(embedding_path)
628
- embedding_mat = np.array(ar_load, dtype=np.float32)
629
- return None, None, embedding_mat, (embedding_mat.shape[0], 100)
630
-
631
- def forward(self, output_hidden_states, words_ents_list, words_subtoken_map, input_ids):
632
- """
633
- Infused KG to the embeddings of output_hidden_states
634
-
635
- Args:
636
- output_hidden_states: Output of each decoder layer.
637
- size: [batch_size, seq_length, hidden_size_dim]
638
- """
639
- bsz, _, _ = output_hidden_states.size()
640
- output = None
641
- # 这里应该在考虑下,句子长度为1时说明在推理,前面的embedding已经保存到cache里了
642
- if output_hidden_states.size()[1] == 1:
643
- return output_hidden_states
644
- #residual = output_hidden_states
645
- residual = output_hidden_states
646
- output_hidden_states = self.input_layernorm(output_hidden_states)
647
- for i in range(bsz):
648
- hidden_state = output_hidden_states[i]
649
- #print("hidden is {}".format(torch.is_half(hidden_state)))
650
- #print(hidden_state.dtype == torch.float16)
651
- try:
652
- words_ents = words_ents_list[i].long().to(hidden_state.device)
653
- except:
654
- words_ents = words_ents_list[0]
655
- try:
656
- words_ents = words_ents.long().to(hidden_state.device)
657
- except:
658
- #print(words_ents_list)
659
- #print(i)
660
- pass
661
- # words_ents = words_ents_list[0].long().to(hidden_state.device)
662
- if len(words_ents) == 0:
663
- if output == None:
664
- output = hidden_state.unsqueeze(0)
665
- else:
666
- output = torch.cat((output, hidden_state.unsqueeze(0)), 0)
667
- continue
668
-
669
- pad_embed = torch.zeros_like(hidden_state[0]).unsqueeze(0)
670
- hidden_state = torch.cat((hidden_state, pad_embed), dim = 0)
671
-
672
- # words_ents size: [map_num, max_mapping_num]
673
- #words_ents = torch.LongTensor(words_ents_list[i])
674
- converted_words_ents = words_ents.masked_fill(words_ents == -1, 0)
675
- ents_embeds = self.concept_embed(converted_words_ents).to(hidden_state.device)
676
- #print(ents_embeds.dtype == torch.float16)
677
- #print("ents_embeds is {}".format(torch.is_half(ents_embeds)))
678
- #print(ents_embeds.requires_grad)
679
- # ents_embeds size: [map_num, top_k + 1, dim]
680
- knowledge_sentinel = self.knowledge_sentinel(torch.LongTensor([0]).to(hidden_state.device)).view(1, 1, -1).repeat(ents_embeds.size()[0], 1, 1)
681
- try:
682
- ents_embeds = torch.cat((ents_embeds, knowledge_sentinel), 1)
683
- ent_ori_embeds = ents_embeds
684
- except:
685
- print(ents_embeds)
686
- print(words_ents)
687
- ent_ori_embeds = ents_embeds
688
-
689
- ents_embeds = self.convert_matrix_entity(ents_embeds)
690
- #words_subtoken = torch.LongTensor(words_subtoken_map[i]).to(hidden_state.device)
691
- try:
692
- words_subtoken = words_subtoken_map[i].long().to(hidden_state.device)
693
- except:
694
- words_subtoken = words_subtoken_map[0].long().to(hidden_state.device)
695
- # words_subtoken_embeds size: [map_num, max_subtoken_num, hidden_size]
696
-
697
- # Avg pooling of each word(tokens)
698
- sub_token_num = words_subtoken.ne(-1).sum(1)
699
- """
700
- print("words_subtoken is {}".format(words_subtoken))
701
- index = words_subtoken.view(-1)
702
- print("index is {}".format(index))
703
- index_fixed = index.masked_fill(index == -1, hidden_state.size()[0] - 1)
704
- print("index_fixed is {}".format(index_fixed))
705
- b = torch.index_select(hidden_state, index=index_fixed, dim = 0)
706
- b = b.view(words_subtoken.size()[0], words_subtoken.size()[1], -1)
707
- """
708
- index_fixed = words_subtoken.masked_fill(words_subtoken == -1, hidden_state.size()[0] - 1)
709
- b = hidden_state[words_subtoken]
710
- b = torch.sum(b, dim = 1).squeeze()
711
- b = torch.div(b, sub_token_num.view(-1, 1))
712
- #b = self.convert_matrix_token(b).unsqueeze(2)
713
- b = b.unsqueeze(2)
714
- atten_weight = torch.bmm(ents_embeds, b)
715
-
716
- # Compute attention mask
717
- atten_ones = torch.ones([words_ents.size()[0], 1]).to(hidden_state.device)
718
- words_ents = torch.cat([words_ents, atten_ones], -1)
719
- attention_mask = torch.zeros_like(words_ents).to(hidden_state.device)
720
- #attention_mask = words_ents.masked_fill(words_ents != -1, value=torch.tensor(0))
721
- attention_mask = attention_mask.masked_fill(words_ents == -1, value=torch.tensor(-1e9))
722
- atten_weight = atten_weight.squeeze() + attention_mask
723
- # atten_weight size: [map_num, top_k + 1]
724
- attn_weights = nn.functional.softmax(atten_weight, dim=-1, dtype=torch.float32).to(ents_embeds.dtype).unsqueeze(1)
725
- attn_output = torch.bmm(attn_weights, ent_ori_embeds)
726
- attn_output = attn_output.repeat(1, words_subtoken.size()[1], 1).view(-1, self.dim)
727
- index_fixed = index_fixed.flatten()
728
- tmp = torch.zeros([hidden_state.size()[0], self.dim]).to(hidden_state.device).to(hidden_state.dtype)
729
- KG_infused = tmp.index_copy(0, index_fixed, attn_output)
730
- KG_infused = torch.cat((hidden_state, KG_infused), -1)[: -1, :]
731
- KG_infused = self.MLP(KG_infused)
732
- KG_infused = self.act_fn(KG_infused).unsqueeze(0)
733
-
734
- if output == None:
735
- output = KG_infused
736
- #print(output.size())
737
- else:
738
- output = torch.cat((output, KG_infused), 0)
739
- #print(output.size())
740
- # print(output.size())
741
- # print(output_hidden_states.size())
742
- assert output.size() == output_hidden_states.size()
743
- #output = torch.nn.functional.dropout(output, p=0.1, training=True)
744
- output = output * self.alpha + residual
745
- return output
746
-
747
- LLAMA_START_DOCSTRING = r"""
748
- This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
749
- library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
750
- etc.)
751
-
752
- This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
753
- Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
754
- and behavior.
755
-
756
- Parameters:
757
- config ([`LlamaConfig`]):
758
- Model configuration class with all the parameters of the model. Initializing with a config file does not
759
- load the weights associated with the model, only the configuration. Check out the
760
- [`~PreTrainedModel.from_pretrained`] method to load the model weights.
761
- """
762
-
763
-
764
- @add_start_docstrings(
765
- "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
766
- LLAMA_START_DOCSTRING,
767
- )
768
- class LlamaPreTrainedModel(PreTrainedModel):
769
- config_class = LlamaConfig
770
- base_model_prefix = "model"
771
- supports_gradient_checkpointing = True
772
- _no_split_modules = ["LlamaDecoderLayer"]
773
- _skip_keys_device_placement = "past_key_values"
774
-
775
- def _init_weights(self, module):
776
- std = self.config.initializer_range
777
- if isinstance(module, nn.Linear):
778
- module.weight.data.normal_(mean=0.0, std=std)
779
- if module.bias is not None:
780
- module.bias.data.zero_()
781
- elif isinstance(module, nn.Embedding):
782
- module.weight.data.normal_(mean=0.0, std=std)
783
- if module.padding_idx is not None:
784
- module.weight.data[module.padding_idx].zero_()
785
-
786
- def _set_gradient_checkpointing(self, module, value=False):
787
- if isinstance(module, LlamaModel):
788
- module.gradient_checkpointing = value
789
-
790
-
791
- LLAMA_INPUTS_DOCSTRING = r"""
792
- Args:
793
- input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
794
- Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
795
- it.
796
-
797
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
798
- [`PreTrainedTokenizer.__call__`] for details.
799
-
800
- [What are input IDs?](../glossary#input-ids)
801
- attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
802
- Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
803
-
804
- - 1 for tokens that are **not masked**,
805
- - 0 for tokens that are **masked**.
806
-
807
- [What are attention masks?](../glossary#attention-mask)
808
-
809
- Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
810
- [`PreTrainedTokenizer.__call__`] for details.
811
-
812
- If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
813
- `past_key_values`).
814
-
815
- If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
816
- and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
817
- information on the default strategy.
818
-
819
- - 1 indicates the head is **not masked**,
820
- - 0 indicates the head is **masked**.
821
- position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
822
- Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
823
- config.n_positions - 1]`.
824
-
825
- [What are position IDs?](../glossary#position-ids)
826
- past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
827
- Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
828
- `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
829
- `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
830
-
831
- Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
832
- blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
833
-
834
- If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
835
- don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
836
- `decoder_input_ids` of shape `(batch_size, sequence_length)`.
837
- inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
838
- Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
839
- is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
840
- model's internal embedding lookup matrix.
841
- use_cache (`bool`, *optional*):
842
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
843
- `past_key_values`).
844
- output_attentions (`bool`, *optional*):
845
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
846
- tensors for more detail.
847
- output_hidden_states (`bool`, *optional*):
848
- Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
849
- more detail.
850
- return_dict (`bool`, *optional*):
851
- Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
852
- """
853
-
854
-
855
- @add_start_docstrings(
856
- "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
857
- LLAMA_START_DOCSTRING,
858
- )
859
- class LlamaModel(LlamaPreTrainedModel):
860
- """
861
- Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
862
-
863
- Args:
864
- config: LlamaConfig
865
- """
866
-
867
- def __init__(self, config: LlamaConfig):
868
- super().__init__(config)
869
- self.padding_idx = config.pad_token_id
870
- self.vocab_size = config.vocab_size
871
-
872
-
873
- self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
874
- print(config.iskg)
875
- if config.iskg:
876
- self.layers = nn.ModuleList([LlamaDecoderLayer_2(idx, config) for idx in range(config.num_hidden_layers)])
877
- else:
878
- self.layers = nn.ModuleList([LlamaDecoderLayer_1(idx, config) for idx in range(config.num_hidden_layers)])
879
- #self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
880
- self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
881
-
882
- self.gradient_checkpointing = False
883
- # Initialize weights and apply final processing
884
- self.post_init()
885
-
886
- def get_input_embeddings(self):
887
- return self.embed_tokens
888
-
889
- def activate_KG_modules(self):
890
- for idx, decoder_layer in enumerate(self.layers):
891
- for name, param in decoder_layer.KG_infuded_module.named_parameters():
892
- param.requires_grad = False
893
- for idx, decoder_layer in enumerate(self.layers):
894
- if (idx + 1) % 32 == 0:
895
- for name, param in decoder_layer.KG_infuded_module.named_parameters():
896
- if name != 'concept_embed.weight':
897
- param.requires_grad = True
898
- def load_off(self):
899
- for i in range(32):
900
- if (i + 1) % 32 != 0:
901
- self.layers[i].KG_infuded_module.concept_embed = None
902
-
903
- def set_input_embeddings(self, value):
904
- self.embed_tokens = value
905
-
906
- # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
907
- def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
908
- # create causal mask
909
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
910
- combined_attention_mask = None
911
- if input_shape[-1] > 1:
912
- combined_attention_mask = _make_causal_mask(
913
- input_shape,
914
- inputs_embeds.dtype,
915
- device=inputs_embeds.device,
916
- past_key_values_length=past_key_values_length,
917
- )
918
-
919
- if attention_mask is not None:
920
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
921
- expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
922
- inputs_embeds.device
923
- )
924
- combined_attention_mask = (
925
- expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
926
- )
927
-
928
- return combined_attention_mask
929
-
930
- @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
931
- def forward(
932
- self,
933
- input_ids: torch.LongTensor = None,
934
- attention_mask: Optional[torch.Tensor] = None,
935
- position_ids: Optional[torch.LongTensor] = None,
936
- past_key_values: Optional[List[torch.FloatTensor]] = None,
937
- inputs_embeds: Optional[torch.FloatTensor] = None,
938
- words_ents_list = None,
939
- words_subtoken_map = None,
940
- use_cache: Optional[bool] = None,
941
- output_attentions: Optional[bool] = None,
942
- output_hidden_states: Optional[bool] = None,
943
- return_dict: Optional[bool] = None,
944
- ) -> Union[Tuple, BaseModelOutputWithPast]:
945
- try:
946
- if self.first:
947
- self.load_off()
948
- self.first = False
949
- except:
950
- pass
951
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
952
- output_hidden_states = (
953
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
954
- )
955
- use_cache = use_cache if use_cache is not None else self.config.use_cache
956
-
957
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
958
-
959
- # retrieve input_ids and inputs_embeds
960
- if input_ids is not None and inputs_embeds is not None:
961
- raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
962
- elif input_ids is not None:
963
- batch_size, seq_length = input_ids.shape
964
- elif inputs_embeds is not None:
965
- batch_size, seq_length, _ = inputs_embeds.shape
966
- else:
967
- raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
968
-
969
- seq_length_with_past = seq_length
970
- past_key_values_length = 0
971
-
972
- if past_key_values is not None:
973
- past_key_values_length = past_key_values[0][0].shape[2]
974
- seq_length_with_past = seq_length_with_past + past_key_values_length
975
-
976
- if position_ids is None:
977
- device = input_ids.device if input_ids is not None else inputs_embeds.device
978
- position_ids = torch.arange(
979
- past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
980
- )
981
- position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
982
- else:
983
- position_ids = position_ids.view(-1, seq_length).long()
984
-
985
- if inputs_embeds is None:
986
- inputs_embeds = self.embed_tokens(input_ids)
987
- # embed positions
988
- if attention_mask is None:
989
- attention_mask = torch.ones(
990
- (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
991
- )
992
- attention_mask = self._prepare_decoder_attention_mask(
993
- attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
994
- )
995
-
996
- hidden_states = inputs_embeds
997
-
998
- if self.gradient_checkpointing and self.training:
999
- if use_cache:
1000
- logger.warning_once(
1001
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1002
- )
1003
- use_cache = False
1004
-
1005
- # decoder layers
1006
- all_hidden_states = () if output_hidden_states else None
1007
- all_self_attns = () if output_attentions else None
1008
- next_decoder_cache = () if use_cache else None
1009
-
1010
- for idx, decoder_layer in enumerate(self.layers):
1011
- if output_hidden_states:
1012
- all_hidden_states += (hidden_states,)
1013
-
1014
- past_key_value = past_key_values[idx] if past_key_values is not None else None
1015
-
1016
- if self.gradient_checkpointing and self.training:
1017
-
1018
- def create_custom_forward(module):
1019
- def custom_forward(*inputs):
1020
- # None for past_key_value
1021
- return module(*inputs, past_key_value, output_attentions)
1022
-
1023
- return custom_forward
1024
-
1025
- layer_outputs = torch.utils.checkpoint.checkpoint(
1026
- create_custom_forward(decoder_layer),
1027
- hidden_states,
1028
- attention_mask,
1029
- position_ids,
1030
- )
1031
- else:
1032
- layer_outputs = decoder_layer(
1033
- hidden_states,
1034
- words_ents_list = words_ents_list,
1035
- words_subtoken_map = words_subtoken_map,
1036
- attention_mask=attention_mask,
1037
- position_ids=position_ids,
1038
- past_key_value=past_key_value,
1039
- output_attentions=output_attentions,
1040
- use_cache=use_cache,
1041
- idx=idx
1042
- )
1043
-
1044
- hidden_states = layer_outputs[0]
1045
-
1046
- if use_cache:
1047
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
1048
-
1049
- if output_attentions:
1050
- all_self_attns += (layer_outputs[1],)
1051
-
1052
- hidden_states = self.norm(hidden_states)
1053
-
1054
- # add hidden states from the last decoder layer
1055
- if output_hidden_states:
1056
- all_hidden_states += (hidden_states,)
1057
-
1058
- next_cache = next_decoder_cache if use_cache else None
1059
- if not return_dict:
1060
- return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1061
- return BaseModelOutputWithPast(
1062
- last_hidden_state=hidden_states,
1063
- past_key_values=next_cache,
1064
- hidden_states=all_hidden_states,
1065
- attentions=all_self_attns,
1066
- )
1067
-
1068
-
1069
- class LlamaForCausalLM(LlamaPreTrainedModel):
1070
- _tied_weights_keys = ["lm_head.weight"]
1071
-
1072
- def __init__(self, config):
1073
- super().__init__(config)
1074
- self.model = LlamaModel(config)
1075
- self.vocab_size = config.vocab_size
1076
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1077
-
1078
- # Initialize weights and apply final processing
1079
- self.post_init()
1080
-
1081
- def get_input_embeddings(self):
1082
- return self.model.embed_tokens
1083
-
1084
- def activate_KG_modules(self):
1085
- self.model.activate_KG_modules()
1086
-
1087
- def set_input_embeddings(self, value):
1088
- self.model.embed_tokens = value
1089
-
1090
- def get_output_embeddings(self):
1091
- return self.lm_head
1092
-
1093
- def set_output_embeddings(self, new_embeddings):
1094
- self.lm_head = new_embeddings
1095
-
1096
- def set_decoder(self, decoder):
1097
- self.model = decoder
1098
-
1099
- def get_decoder(self):
1100
- return self.model
1101
-
1102
- @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
1103
- @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1104
- def forward(
1105
- self,
1106
- input_ids: torch.LongTensor = None,
1107
- attention_mask: Optional[torch.Tensor] = None,
1108
- position_ids: Optional[torch.LongTensor] = None,
1109
- past_key_values: Optional[List[torch.FloatTensor]] = None,
1110
- inputs_embeds: Optional[torch.FloatTensor] = None,
1111
- words_ents_list = None,
1112
- words_subtoken_map = None,
1113
- labels: Optional[torch.LongTensor] = None,
1114
- use_cache: Optional[bool] = None,
1115
- output_attentions: Optional[bool] = None,
1116
- output_hidden_states: Optional[bool] = None,
1117
- return_dict: Optional[bool] = None,
1118
- ) -> Union[Tuple, CausalLMOutputWithPast]:
1119
- r"""
1120
- Args:
1121
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1122
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1123
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1124
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1125
-
1126
- Returns:
1127
-
1128
- Example:
1129
-
1130
- ```python
1131
- >>> from transformers import AutoTokenizer, LlamaForCausalLM
1132
-
1133
- >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1134
- >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1135
-
1136
- >>> prompt = "Hey, are you conscious? Can you talk to me?"
1137
- >>> inputs = tokenizer(prompt, return_tensors="pt")
1138
-
1139
- >>> # Generate
1140
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1141
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1142
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1143
- ```"""
1144
-
1145
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1146
- output_hidden_states = (
1147
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1148
- )
1149
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1150
-
1151
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1152
- outputs = self.model(
1153
- input_ids=input_ids,
1154
- attention_mask=attention_mask,
1155
- position_ids=position_ids,
1156
- past_key_values=past_key_values,
1157
- inputs_embeds=inputs_embeds,
1158
- words_ents_list = words_ents_list,
1159
- words_subtoken_map = words_subtoken_map,
1160
- use_cache=use_cache,
1161
- output_attentions=output_attentions,
1162
- output_hidden_states=output_hidden_states,
1163
- return_dict=return_dict,
1164
- )
1165
-
1166
- hidden_states = outputs[0]
1167
- if self.config.pretraining_tp > 1:
1168
- lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
1169
- logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
1170
- logits = torch.cat(logits, dim=-1)
1171
- else:
1172
- logits = self.lm_head(hidden_states)
1173
- logits = logits.float()
1174
-
1175
- loss = None
1176
- if labels is not None:
1177
- # Shift so that tokens < n predict n
1178
- shift_logits = logits[..., :-1, :].contiguous()
1179
- shift_labels = labels[..., 1:].contiguous()
1180
- # Flatten the tokens
1181
- loss_fct = CrossEntropyLoss()
1182
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
1183
- shift_labels = shift_labels.view(-1)
1184
- # Enable model parallelism
1185
- shift_labels = shift_labels.to(shift_logits.device)
1186
- loss = loss_fct(shift_logits, shift_labels)
1187
-
1188
- if not return_dict:
1189
- output = (logits,) + outputs[1:]
1190
- return (loss,) + output if loss is not None else output
1191
-
1192
- return CausalLMOutputWithPast(
1193
- loss=loss,
1194
- logits=logits,
1195
- past_key_values=outputs.past_key_values,
1196
- hidden_states=outputs.hidden_states,
1197
- attentions=outputs.attentions,
1198
- )
1199
-
1200
- def prepare_inputs_for_generation(
1201
- self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
1202
- ):
1203
- if past_key_values:
1204
- input_ids = input_ids[:, -1:]
1205
-
1206
- position_ids = kwargs.get("position_ids", None)
1207
- if attention_mask is not None and position_ids is None:
1208
- # create position_ids on the fly for batch generation
1209
- position_ids = attention_mask.long().cumsum(-1) - 1
1210
- position_ids.masked_fill_(attention_mask == 0, 1)
1211
- if past_key_values:
1212
- position_ids = position_ids[:, -1].unsqueeze(-1)
1213
-
1214
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1215
- if inputs_embeds is not None and past_key_values is None:
1216
- model_inputs = {"inputs_embeds": inputs_embeds}
1217
- else:
1218
- model_inputs = {"input_ids": input_ids}
1219
- model_inputs.update(
1220
- {
1221
- "position_ids": position_ids,
1222
- "words_ents_list": kwargs.get("words_ents_list"),
1223
- "words_subtoken_map": kwargs.get("words_subtoken_map"),
1224
- "past_key_values": past_key_values,
1225
- "use_cache": kwargs.get("use_cache"),
1226
- "attention_mask": attention_mask,
1227
- }
1228
- )
1229
- return model_inputs
1230
-
1231
- @staticmethod
1232
- def _reorder_cache(past_key_values, beam_idx):
1233
- reordered_past = ()
1234
- for layer_past in past_key_values:
1235
- reordered_past += (
1236
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
1237
- )
1238
- return reordered_past
1239
-
1240
-
1241
- @add_start_docstrings(
1242
- """
1243
- The LLaMa Model transformer with a sequence classification head on top (linear layer).
1244
-
1245
- [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1246
- (e.g. GPT-2) do.
1247
-
1248
- Since it does classification on the last token, it requires to know the position of the last token. If a
1249
- `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1250
- no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1251
- padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1252
- each row of the batch).
1253
- """,
1254
- LLAMA_START_DOCSTRING,
1255
- )
1256
- class LlamaForSequenceClassification(LlamaPreTrainedModel):
1257
- def __init__(self, config):
1258
- super().__init__(config)
1259
- self.num_labels = config.num_labels
1260
- self.model = LlamaModel(config)
1261
- self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1262
-
1263
- # Initialize weights and apply final processing
1264
- self.post_init()
1265
-
1266
- def get_input_embeddings(self):
1267
- return self.model.embed_tokens
1268
-
1269
- def set_input_embeddings(self, value):
1270
- self.model.embed_tokens = value
1271
-
1272
- @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
1273
- def forward(
1274
- self,
1275
- input_ids: torch.LongTensor = None,
1276
- attention_mask: Optional[torch.Tensor] = None,
1277
- position_ids: Optional[torch.LongTensor] = None,
1278
- past_key_values: Optional[List[torch.FloatTensor]] = None,
1279
- inputs_embeds: Optional[torch.FloatTensor] = None,
1280
- labels: Optional[torch.LongTensor] = None,
1281
- use_cache: Optional[bool] = None,
1282
- output_attentions: Optional[bool] = None,
1283
- output_hidden_states: Optional[bool] = None,
1284
- return_dict: Optional[bool] = None,
1285
- ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1286
- r"""
1287
- labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1288
- Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1289
- config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1290
- `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1291
- """
1292
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1293
-
1294
- transformer_outputs = self.model(
1295
- input_ids,
1296
- attention_mask=attention_mask,
1297
- position_ids=position_ids,
1298
- past_key_values=past_key_values,
1299
- inputs_embeds=inputs_embeds,
1300
- use_cache=use_cache,
1301
- output_attentions=output_attentions,
1302
- output_hidden_states=output_hidden_states,
1303
- return_dict=return_dict,
1304
- )
1305
- hidden_states = transformer_outputs[0]
1306
- logits = self.score(hidden_states)
1307
-
1308
- if input_ids is not None:
1309
- batch_size = input_ids.shape[0]
1310
- else:
1311
- batch_size = inputs_embeds.shape[0]
1312
-
1313
- if self.config.pad_token_id is None and batch_size != 1:
1314
- raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1315
- if self.config.pad_token_id is None:
1316
- sequence_lengths = -1
1317
- else:
1318
- if input_ids is not None:
1319
- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
1320
- logits.device
1321
- )
1322
- else:
1323
- sequence_lengths = -1
1324
-
1325
- pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1326
-
1327
- loss = None
1328
- if labels is not None:
1329
- labels = labels.to(logits.device)
1330
- if self.config.problem_type is None:
1331
- if self.num_labels == 1:
1332
- self.config.problem_type = "regression"
1333
- elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1334
- self.config.problem_type = "single_label_classification"
1335
- else:
1336
- self.config.problem_type = "multi_label_classification"
1337
-
1338
- if self.config.problem_type == "regression":
1339
- loss_fct = MSELoss()
1340
- if self.num_labels == 1:
1341
- loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1342
- else:
1343
- loss = loss_fct(pooled_logits, labels)
1344
- elif self.config.problem_type == "single_label_classification":
1345
- loss_fct = CrossEntropyLoss()
1346
- loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1347
- elif self.config.problem_type == "multi_label_classification":
1348
- loss_fct = BCEWithLogitsLoss()
1349
- loss = loss_fct(pooled_logits, labels)
1350
- if not return_dict:
1351
- output = (pooled_logits,) + transformer_outputs[1:]
1352
- return ((loss,) + output) if loss is not None else output
1353
-
1354
- return SequenceClassifierOutputWithPast(
1355
- loss=loss,
1356
- logits=pooled_logits,
1357
- past_key_values=transformer_outputs.past_key_values,
1358
- hidden_states=transformer_outputs.hidden_states,
1359
- attentions=transformer_outputs.attentions,
1360
- )