myownskyW7 commited on
Commit
5985c08
1 Parent(s): e19ed3c

init internlm-xcomposer-vl-7b

Browse files
added_tokens.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "<ITG_TOKEN>": 103168,
3
+ "<SOI_TOKEN>": 103169
4
+ }
config.json ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "performance",
3
+ "architectures": [
4
+ "InternLMXComposerForCausalLM"
5
+ ],
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_InternLM_XComposer.InternLMXComposerConfig",
8
+ "AutoModel": "modeling_InternLM_XComposer.InternLMXComposerForCausalLM",
9
+ "AutoModelForCausalLM": "modeling_InternLM_XComposer.InternLMXComposerForCausalLM"
10
+ },
11
+ "bias": true,
12
+ "bos_token_id": 1,
13
+ "device": "cuda",
14
+ "eos_token_id": 2,
15
+ "hidden_act": "silu",
16
+ "hidden_size": 4096,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 11008,
19
+ "intern_converted_llm": true,
20
+ "internlm_lora": {
21
+ "freeze": false,
22
+ "learn_param": [
23
+ "q",
24
+ "v",
25
+ "ffn"
26
+ ],
27
+ "lora_alpha": 256,
28
+ "lora_dropout": 0.05,
29
+ "lora_r": 256
30
+ },
31
+ "kqvo_bias": true,
32
+ "lora_cfg": {
33
+ "freeze": false,
34
+ "learn_param": [
35
+ "q",
36
+ "v",
37
+ "ffn"
38
+ ],
39
+ "lora_alpha": 256,
40
+ "lora_dropout": 0.05,
41
+ "lora_r": 256
42
+ },
43
+ "max_position_embeddings": 2048,
44
+ "model_type": "InternLMXComposer",
45
+ "num_attention_heads": 32,
46
+ "num_hidden_layers": 32,
47
+ "num_quant": 32,
48
+ "num_query_token": 64,
49
+ "pad_token_id": -1,
50
+ "rms_norm_eps": 1e-05,
51
+ "tie_word_embeddings": false,
52
+ "torch_dtype": "float32",
53
+ "transformers_version": "4.30.2",
54
+ "use_cache": true,
55
+ "vocab_size": 103170
56
+ }
configuration_InternLM_XComposer.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers.configuration_utils import PretrainedConfig
2
+ from transformers.utils import logging
3
+
4
+ logger = logging.get_logger(__name__)
5
+
6
+ INTERNLM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
7
+
8
+
9
+ class InternLMXComposerConfig(PretrainedConfig):
10
+
11
+ model_type = "InternLMXComposer"
12
+ _auto_class = "AutoConfig"
13
+
14
+ def __init__(
15
+ self,
16
+ vocab_size=103168,
17
+ hidden_size=4096,
18
+ intermediate_size=11008,
19
+ num_hidden_layers=32,
20
+ num_attention_heads=32,
21
+ hidden_act="silu",
22
+ max_position_embeddings=2048,
23
+ initializer_range=0.02,
24
+ rms_norm_eps=1e-5,
25
+ use_cache=True,
26
+ pad_token_id=-1,
27
+ bos_token_id=1,
28
+ eos_token_id=2,
29
+ tie_word_embeddings=False,
30
+ bias=True,
31
+ num_query_token=32,
32
+ num_quant=32,
33
+ intern_converted_llm=True,
34
+ kqvo_bias=True,
35
+ device='cuda',
36
+ internlm_lora=None,
37
+ **kwargs,
38
+ ):
39
+ self.vocab_size = vocab_size
40
+ self.max_position_embeddings = max_position_embeddings
41
+ self.hidden_size = hidden_size
42
+ self.intermediate_size = intermediate_size
43
+ self.num_hidden_layers = num_hidden_layers
44
+ self.num_attention_heads = num_attention_heads
45
+ self.hidden_act = hidden_act
46
+ self.initializer_range = initializer_range
47
+ self.rms_norm_eps = rms_norm_eps
48
+ self.use_cache = use_cache
49
+ self.bias = bias
50
+ self.num_query_token = num_query_token
51
+ self.num_quant = num_quant
52
+ self.internlm_lora = internlm_lora
53
+ self.kqvo_bias = kqvo_bias
54
+ self.intern_converted_llm = intern_converted_llm
55
+ self.device = device
56
+ super().__init__(
57
+ pad_token_id=pad_token_id,
58
+ bos_token_id=bos_token_id,
59
+ eos_token_id=eos_token_id,
60
+ tie_word_embeddings=tie_word_embeddings,
61
+ **kwargs,
62
+ )
modeling_InternLM.py ADDED
@@ -0,0 +1,1247 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from typing import List, Union
3
+ from typing import Optional, Tuple
4
+
5
+ import rotary_emb
6
+ import torch
7
+ import torch.utils.checkpoint
8
+ import torch.utils.checkpoint
9
+ from einops import rearrange
10
+ from flash_attn.layers.rotary import ApplyRotaryEmbQKV_ as LegacyApplyRotaryEmbQKV_
11
+ from torch import nn
12
+ from torch.nn import CrossEntropyLoss
13
+ from transformers.activations import ACT2FN
14
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
15
+ from transformers.modeling_utils import PreTrainedModel
16
+ from transformers.utils import logging
17
+
18
+ from .configuration_InternLM_XComposer import InternLMXComposerConfig
19
+ from .modeling_utils import LoRALinear
20
+
21
+ logger = logging.get_logger(__name__)
22
+
23
+ _CONFIG_FOR_DOC = "InternLMXComposerConfig"
24
+
25
+
26
+ class ApplyRotaryEmbQKV_(torch.autograd.Function):
27
+ """
28
+ ApplyRotaryEmbQKV_
29
+ """
30
+ @staticmethod
31
+ def forward(ctx, qkv, cos, sin, cos_k=None, sin_k=None):
32
+ """
33
+ qkv: (total, 3, nheads, headdim)
34
+ cos, sin: (seqlen, rotary_dim / 2)
35
+ cos_k, sin_k: (seqlen, rotary_dim / 2), optional
36
+ rotary_dim must be <= headdim
37
+ Apply rotary embedding *inplace* to the first rotary_dim of q and k.
38
+ """
39
+ _, three, _, headdim = qkv.shape
40
+ assert three == 3
41
+ rotary_seqlen, rotary_dim = cos.shape
42
+ rotary_dim *= 2
43
+ assert rotary_dim <= headdim
44
+ cos_k = cos if cos_k is None else cos_k
45
+ sin_k = sin if sin_k is None else sin_k
46
+ assert sin.shape == cos_k.shape == sin_k.shape == (rotary_seqlen,
47
+ rotary_dim // 2)
48
+ q1, q2 = qkv[:, 0, :, :rotary_dim].chunk(2, dim=-1)
49
+ rotary_emb.apply_rotary(q1, q2, rearrange(cos, "s d -> s 1 d"),
50
+ rearrange(sin, "s d -> s 1 d"), q1, q2, False)
51
+ k1, k2 = qkv[:, 1, :, :rotary_dim].chunk(2, dim=-1)
52
+ rotary_emb.apply_rotary(k1, k2, rearrange(cos_k, "s d -> s 1 d"),
53
+ rearrange(sin_k, "s d -> s 1 d"), k1, k2,
54
+ False)
55
+ ctx.save_for_backward(cos, sin, cos_k, sin_k)
56
+ return qkv
57
+
58
+ @staticmethod
59
+ def backward(ctx, dqkv):
60
+ cos, sin, cos_k, sin_k = ctx.saved_tensors
61
+ rotary_dim = cos.shape[-1]
62
+ rotary_dim *= 2
63
+ dq1, dq2 = dqkv[:, 0, :, :rotary_dim].chunk(2, dim=-1)
64
+ rotary_emb.apply_rotary(dq1, dq2, rearrange(cos, "s d -> s 1 d"),
65
+ rearrange(sin, "s d -> s 1 d"), dq1, dq2, True)
66
+ dk1, dk2 = dqkv[:, 1, :, :rotary_dim].chunk(2, dim=-1)
67
+ rotary_emb.apply_rotary(dk1, dk2, rearrange(cos_k, "s d -> s 1 d"),
68
+ rearrange(sin_k, "s d -> s 1 d"), dk1, dk2,
69
+ True)
70
+ return dqkv, None, None, None, None
71
+
72
+
73
+ class ConvertedInternLMRotaryEmbedding(torch.nn.Module):
74
+ def __init__(self, dim: int, base=10000, scale_base=0, device=None):
75
+ """ """
76
+ super().__init__()
77
+ # Generate and save the inverse frequency buffer (non trainable)
78
+ inv_freq = 1.0 / (base**(
79
+ torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
80
+ self.register_buffer("inv_freq", inv_freq)
81
+ self.scale_base = scale_base
82
+ scale = ((torch.arange(0, dim, 2, device=device, dtype=torch.float32) +
83
+ 0.4 * dim) / (1.4 * dim) if scale_base > 0 else None)
84
+ self.register_buffer("scale", scale)
85
+
86
+ self._seq_len_cached = 0
87
+ self._cos_cached = None
88
+ self._sin_cached = None
89
+ self._cos_k_cached = None
90
+ self._sin_k_cached = None
91
+
92
+ def _update_cos_sin_cache(self, x, indexes):
93
+ """x: (batch, seqlen, nheads, headdim) or (batch, seqlen, 3, nheads, headdim)"""
94
+ if not isinstance(indexes, int):
95
+ seqlen = indexes.max().item() + 1
96
+ else:
97
+ seqlen = indexes + 1 # eval_forward
98
+ # Reset the tables if the sequence length has changed,
99
+ # or if we're on a new device (possibly due to tracing for instance)
100
+ if seqlen > self._seq_len_cached or self._cos_cached.device != x.device or self._cos_cached.dtype != x.dtype:
101
+ self._seq_len_cached = seqlen
102
+ t = torch.arange(seqlen,
103
+ device=x.device,
104
+ dtype=self.inv_freq.dtype)
105
+ # Don't do einsum, it converts fp32 to fp16
106
+ # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
107
+ freqs = torch.outer(t, self.inv_freq.to(device=t.device))
108
+ if self.scale is None:
109
+ self._cos_cached = torch.cos(freqs).to(x.dtype)
110
+ self._sin_cached = torch.sin(freqs).to(x.dtype)
111
+ else:
112
+ power = (torch.arange(
113
+ seqlen, dtype=self.scale.dtype, device=self.scale.device) -
114
+ seqlen // 2) / self.scale_base
115
+ scale = self.scale.to(device=power.device)**rearrange(
116
+ power, "s -> s 1")
117
+ # We want the multiplication by scale to happen in fp32
118
+ self._cos_cached = (torch.cos(freqs) * scale).to(x.dtype)
119
+ self._sin_cached = (torch.sin(freqs) * scale).to(x.dtype)
120
+ self._cos_k_cached = (torch.cos(freqs) / scale).to(x.dtype)
121
+ self._sin_k_cached = (torch.sin(freqs) / scale).to(x.dtype)
122
+
123
+ def forward(self,
124
+ qkv: torch.Tensor,
125
+ indexes=0) -> Tuple[torch.Tensor, torch.Tensor]:
126
+ self._update_cos_sin_cache(qkv, indexes)
127
+ if self.scale is None:
128
+ return apply_rotary_emb_qkv_(qkv, self._cos_cached[indexes],
129
+ self._sin_cached[indexes]).to(
130
+ qkv.dtype)
131
+ else:
132
+ return apply_rotary_emb_qkv_(
133
+ qkv,
134
+ self._cos_cached[indexes],
135
+ self._sin_cached[indexes],
136
+ self._cos_k_cached[indexes],
137
+ self._sin_k_cached[indexes],
138
+ ).to(qkv.dtype)
139
+
140
+ def eval_forward(self, qkv, seqlen_offset=0):
141
+ """
142
+ seqlen_offset: can be used in generation where the qkv being passed in is only the last
143
+ token in the batch.
144
+ """
145
+ self._update_cos_sin_cache(qkv, seqlen_offset + qkv.shape[1])
146
+ if self.scale is None:
147
+ return legacy_apply_rotary_embed_qkv(
148
+ qkv, self._cos_cached[seqlen_offset:],
149
+ self._sin_cached[seqlen_offset:])
150
+ else:
151
+ return legacy_apply_rotary_embed_qkv(
152
+ qkv,
153
+ self._cos_cached[seqlen_offset:],
154
+ self._sin_cached[seqlen_offset:],
155
+ self._cos_k_cached[seqlen_offset:],
156
+ self._sin_k_cached[seqlen_offset:],
157
+ )
158
+
159
+
160
+ apply_rotary_emb_qkv_ = ApplyRotaryEmbQKV_.apply
161
+ legacy_apply_rotary_embed_qkv = LegacyApplyRotaryEmbQKV_.apply
162
+
163
+
164
+ class InternConvertedInternLMAttention(nn.Module):
165
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
166
+ def __init__(self, config: InternLMXComposerConfig):
167
+ super().__init__()
168
+ self.config = config
169
+ self.hidden_size = config.hidden_size
170
+ self.num_heads = config.num_attention_heads
171
+ self.head_dim = self.hidden_size // self.num_heads
172
+ self.max_position_embeddings = config.max_position_embeddings
173
+
174
+ if (self.head_dim * self.num_heads) != self.hidden_size:
175
+ raise ValueError(
176
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
177
+ f" and `num_heads`: {self.num_heads}).")
178
+ if config.lora_cfg is None:
179
+ self.q_proj = nn.Linear(self.hidden_size,
180
+ self.num_heads * self.head_dim,
181
+ bias=config.kqvo_bias)
182
+ self.k_proj = nn.Linear(self.hidden_size,
183
+ self.num_heads * self.head_dim,
184
+ bias=config.kqvo_bias)
185
+ self.v_proj = nn.Linear(self.hidden_size,
186
+ self.num_heads * self.head_dim,
187
+ bias=config.kqvo_bias)
188
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim,
189
+ self.hidden_size,
190
+ bias=config.kqvo_bias)
191
+
192
+ else:
193
+ lora_cfg = config.lora_cfg
194
+ if 'q' in lora_cfg['learn_param']:
195
+ self.q_proj = LoRALinear(self.hidden_size,
196
+ self.num_heads * self.head_dim,
197
+ bias=config.kqvo_bias,
198
+ **lora_cfg)
199
+ else:
200
+ self.q_proj = nn.Linear(
201
+ self.hidden_size,
202
+ self.num_heads * self.head_dim,
203
+ bias=config.kqvo_bias,
204
+ )
205
+ if 'k' in lora_cfg['learn_param']:
206
+ self.k_proj = LoRALinear(self.hidden_size,
207
+ self.num_heads * self.head_dim,
208
+ bias=config.kqvo_bias,
209
+ **lora_cfg)
210
+ else:
211
+ self.k_proj = nn.Linear(
212
+ self.hidden_size,
213
+ self.num_heads * self.head_dim,
214
+ bias=config.kqvo_bias,
215
+ )
216
+ if 'v' in lora_cfg['learn_param']:
217
+ self.v_proj = LoRALinear(self.hidden_size,
218
+ self.num_heads * self.head_dim,
219
+ bias=config.kqvo_bias,
220
+ **lora_cfg)
221
+ else:
222
+ self.v_proj = nn.Linear(
223
+ self.hidden_size,
224
+ self.num_heads * self.head_dim,
225
+ bias=config.kqvo_bias,
226
+ )
227
+
228
+ if 'o' in lora_cfg['learn_param']:
229
+ self.o_proj = LoRALinear(self.num_heads * self.head_dim,
230
+ self.hidden_size,
231
+ bias=config.kqvo_bias,
232
+ **lora_cfg)
233
+ else:
234
+ self.o_proj = nn.Linear(
235
+ self.num_heads * self.head_dim,
236
+ self.hidden_size,
237
+ bias=config.kqvo_bias,
238
+ )
239
+
240
+ self.rotary_emb = ConvertedInternLMRotaryEmbedding(self.head_dim)
241
+
242
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
243
+ return tensor.view(bsz, seq_len, self.num_heads,
244
+ self.head_dim).transpose(1, 2).contiguous()
245
+
246
+ def forward(
247
+ self,
248
+ hidden_states: torch.Tensor,
249
+ attention_mask: Optional[torch.Tensor] = None,
250
+ position_ids: Optional[torch.LongTensor] = None,
251
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
252
+ output_attentions: bool = False,
253
+ use_cache: bool = False,
254
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
255
+ Optional[Tuple[torch.Tensor]]]:
256
+ bsz, q_len, _ = hidden_states.size()
257
+
258
+ query_states = self.q_proj(hidden_states)
259
+ key_states = self.k_proj(hidden_states)
260
+ value_states = self.v_proj(hidden_states)
261
+
262
+ q = query_states
263
+ k = key_states
264
+ v = value_states
265
+
266
+ qkv = torch.cat([q, k, v], dim=2).contiguous()
267
+ qkv = qkv.view(bsz, q_len, -1)
268
+ qkv = rearrange(qkv,
269
+ "b s (three h d) -> b s three h d",
270
+ three=3,
271
+ d=self.head_dim)
272
+
273
+ if past_key_value is not None:
274
+ qkv = self.rotary_emb.eval_forward(
275
+ qkv, seqlen_offset=past_key_value[0].shape[2])
276
+ else:
277
+ qkv = self.rotary_emb.eval_forward(qkv)
278
+
279
+ query_states, key_states, value_states = qkv.unbind(2)
280
+ query_states = query_states.transpose(1, 2)
281
+ key_states = key_states.transpose(1, 2)
282
+ value_states = value_states.transpose(1, 2)
283
+
284
+ kv_seq_len = key_states.shape[-2]
285
+ if past_key_value is not None:
286
+ kv_seq_len += past_key_value[0].shape[-2]
287
+ # cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
288
+ # query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
289
+ # [bsz, nh, t, hd]
290
+
291
+ if past_key_value is not None:
292
+ # reuse k, v, self_attention
293
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
294
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
295
+
296
+ past_key_value = (key_states, value_states) if use_cache else None
297
+
298
+ attn_weights = torch.matmul(query_states, key_states.transpose(
299
+ 2, 3)) / math.sqrt(self.head_dim)
300
+
301
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
302
+ raise ValueError(
303
+ f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
304
+ f" {attn_weights.size()}")
305
+
306
+ if attention_mask is not None:
307
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
308
+ raise ValueError(
309
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
310
+ )
311
+ attn_weights = attn_weights + attention_mask
312
+ attn_weights = torch.max(
313
+ attn_weights,
314
+ torch.tensor(torch.finfo(attn_weights.dtype).min))
315
+
316
+ # upcast attention to fp32
317
+ attn_weights = nn.functional.softmax(attn_weights,
318
+ dim=-1,
319
+ dtype=torch.float32).to(
320
+ query_states.dtype)
321
+ attn_output = torch.matmul(attn_weights, value_states)
322
+
323
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
324
+ raise ValueError(
325
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
326
+ f" {attn_output.size()}")
327
+
328
+ attn_output = attn_output.transpose(1, 2)
329
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
330
+
331
+ attn_output = self.o_proj(attn_output)
332
+
333
+ if not output_attentions:
334
+ attn_weights = None
335
+
336
+ return attn_output, attn_weights, past_key_value
337
+
338
+
339
+ class ConvertedLoRALinear(nn.Linear):
340
+ def __init__(self,
341
+ in_features: int,
342
+ out_features: int,
343
+ bias: bool = True,
344
+ device=None,
345
+ dtype=None,
346
+ lora_r=8,
347
+ lora_alpha=16,
348
+ lora_dropout=0.05,
349
+ **kwargs) -> None:
350
+ super().__init__(in_features, out_features, bias, device, dtype)
351
+ self.lora_r = lora_r
352
+ self.lora_alpha = lora_alpha
353
+ if lora_dropout > 0.:
354
+ self.lora_dropout = nn.Dropout(p=lora_dropout)
355
+ else:
356
+ self.lora_dropout = lambda x: x
357
+ self.lora_scaling = self.lora_alpha / self.lora_r
358
+
359
+ self.lora_A = nn.Linear(in_features,
360
+ self.lora_r,
361
+ bias=False,
362
+ device=device,
363
+ dtype=dtype)
364
+ self.lora_B = nn.Linear(self.lora_r,
365
+ out_features,
366
+ bias=False,
367
+ device=device,
368
+ dtype=dtype)
369
+
370
+ self.reset_parameters()
371
+
372
+ def reset_parameters(self):
373
+ if hasattr(self, 'lora_A'):
374
+ # initialize A the same way as the default for nn.Linear and B to zero
375
+ nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
376
+ nn.init.zeros_(self.lora_B.weight)
377
+ # print ("lora weight init {} {}".format(torch.mean(self.lora_A.weight), torch.mean(self.lora_B.weight)))
378
+
379
+ def forward(self, x):
380
+ orig_type = x.dtype
381
+ res = super().forward(x)
382
+
383
+ dim = int(res.shape[-1] // 2)
384
+
385
+ r1 = res[..., :dim]
386
+ r2 = res[..., dim:]
387
+
388
+ r1 = r1.float()
389
+ r2 = r2.float()
390
+ x_ = x.float()
391
+
392
+ tmp = self.lora_B(self.lora_A(
393
+ self.lora_dropout(x_))) * self.lora_scaling
394
+ tmp1 = tmp[..., ::2]
395
+ tmp2 = tmp[..., 1::2]
396
+
397
+ r1 += tmp1
398
+ r2 += tmp2
399
+
400
+ r1 = r1.to(orig_type)
401
+ r2 = r2.to(orig_type)
402
+
403
+ res = torch.cat([r1, r2], -1)
404
+
405
+ # res += self.lora_B(self.lora_A(
406
+ # self.lora_dropout(x))) * self.lora_scaling
407
+ return res
408
+
409
+
410
+ # Copied from transformers.models.bart.modeling_bart._make_causal_mask
411
+ def _make_causal_mask(input_ids_shape: torch.Size,
412
+ dtype: torch.dtype,
413
+ device: torch.device,
414
+ past_key_values_length: int = 0):
415
+ """
416
+ Make causal mask used for bi-directional self-attention.
417
+ """
418
+ bsz, tgt_len = input_ids_shape
419
+ mask = torch.full((tgt_len, tgt_len),
420
+ torch.tensor(torch.finfo(dtype).min, device=device),
421
+ device=device)
422
+ mask_cond = torch.arange(mask.size(-1), device=device)
423
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
424
+ mask = mask.to(dtype)
425
+
426
+ if past_key_values_length > 0:
427
+ mask = torch.cat([
428
+ torch.zeros(
429
+ tgt_len, past_key_values_length, dtype=dtype, device=device),
430
+ mask
431
+ ],
432
+ dim=-1)
433
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len,
434
+ tgt_len + past_key_values_length)
435
+
436
+
437
+ # Copied from transformers.models.bart.modeling_bart._expand_mask
438
+ def _expand_mask(mask: torch.Tensor,
439
+ dtype: torch.dtype,
440
+ tgt_len: Optional[int] = None):
441
+ """
442
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
443
+ """
444
+ bsz, src_len = mask.size()
445
+ tgt_len = tgt_len if tgt_len is not None else src_len
446
+
447
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len,
448
+ src_len).to(dtype)
449
+
450
+ inverted_mask = 1.0 - expanded_mask
451
+
452
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool),
453
+ torch.finfo(dtype).min)
454
+
455
+
456
+ class InternLMRMSNorm(nn.Module):
457
+ def __init__(self, hidden_size, eps=1e-6):
458
+ """
459
+ InternLMRMSNorm is equivalent to T5LayerNorm
460
+ """
461
+ super().__init__()
462
+ self.weight = nn.Parameter(torch.ones(hidden_size))
463
+ self.variance_epsilon = eps
464
+
465
+ def forward(self, hidden_states):
466
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1,
467
+ keepdim=True)
468
+ hidden_states = hidden_states * torch.rsqrt(variance +
469
+ self.variance_epsilon)
470
+
471
+ # convert into half-precision if necessary
472
+ if self.weight.dtype in [torch.float16, torch.bfloat16]:
473
+ hidden_states = hidden_states.to(self.weight.dtype)
474
+
475
+ return self.weight * hidden_states
476
+
477
+
478
+ class InternLMRotaryEmbedding(torch.nn.Module):
479
+ def __init__(self,
480
+ dim,
481
+ max_position_embeddings=2048,
482
+ base=10000,
483
+ device=None):
484
+ super().__init__()
485
+ inv_freq = 1.0 / (base
486
+ **(torch.arange(0, dim, 2).float().to(device) / dim))
487
+ self.register_buffer("inv_freq", inv_freq)
488
+
489
+ # Build here to make `torch.jit.trace` work.
490
+ self.max_seq_len_cached = max_position_embeddings
491
+ t = torch.arange(self.max_seq_len_cached,
492
+ device=self.inv_freq.device,
493
+ dtype=self.inv_freq.dtype)
494
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
495
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
496
+ emb = torch.cat((freqs, freqs), dim=-1)
497
+ self.register_buffer("cos_cached",
498
+ emb.cos()[None, None, :, :],
499
+ persistent=False)
500
+ self.register_buffer("sin_cached",
501
+ emb.sin()[None, None, :, :],
502
+ persistent=False)
503
+
504
+ def forward(self, x, seq_len=None):
505
+ # x: [bs, num_attention_heads, seq_len, head_size]
506
+ # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
507
+ if seq_len > self.max_seq_len_cached:
508
+ self.max_seq_len_cached = seq_len
509
+ t = torch.arange(self.max_seq_len_cached,
510
+ device=x.device,
511
+ dtype=self.inv_freq.dtype)
512
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
513
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
514
+ emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
515
+ self.register_buffer("cos_cached",
516
+ emb.cos()[None, None, :, :],
517
+ persistent=False)
518
+ self.register_buffer("sin_cached",
519
+ emb.sin()[None, None, :, :],
520
+ persistent=False)
521
+ return (
522
+ self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
523
+ self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
524
+ )
525
+
526
+
527
+ def rotate_half(x):
528
+ """Rotates half the hidden dims of the input."""
529
+ x1 = x[..., :x.shape[-1] // 2]
530
+ x2 = x[..., x.shape[-1] // 2:]
531
+ return torch.cat((-x2, x1), dim=-1)
532
+
533
+
534
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
535
+ gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1]
536
+ gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3])
537
+ cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2,
538
+ gather_indices)
539
+ sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2,
540
+ gather_indices)
541
+ q_embed = (q * cos) + (rotate_half(q) * sin)
542
+ k_embed = (k * cos) + (rotate_half(k) * sin)
543
+ return q_embed, k_embed
544
+
545
+
546
+ class InternLMMLP(nn.Module):
547
+ def __init__(self, hidden_size: int, intermediate_size: int,
548
+ hidden_act: str, config: InternLMXComposerConfig):
549
+ super().__init__()
550
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
551
+ if config.lora_cfg is not None and 'ffn' in config.lora_cfg[
552
+ 'learn_param']:
553
+ lora_cfg = config.lora_cfg
554
+ self.down_proj = LoRALinear(intermediate_size,
555
+ hidden_size,
556
+ bias=False,
557
+ **lora_cfg)
558
+ self.up_proj = LoRALinear(hidden_size,
559
+ intermediate_size,
560
+ bias=False,
561
+ **lora_cfg)
562
+ else:
563
+ self.down_proj = nn.Linear(intermediate_size,
564
+ hidden_size,
565
+ bias=False)
566
+ self.up_proj = nn.Linear(hidden_size,
567
+ intermediate_size,
568
+ bias=False)
569
+ self.act_fn = ACT2FN[hidden_act]
570
+
571
+ def forward(self, x):
572
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
573
+
574
+
575
+ class InternLMAttention(nn.Module):
576
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
577
+ def __init__(self, config: InternLMXComposerConfig):
578
+ super().__init__()
579
+ self.config = config
580
+ self.hidden_size = config.hidden_size
581
+ self.num_heads = config.num_attention_heads
582
+ self.head_dim = self.hidden_size // self.num_heads
583
+ self.max_position_embeddings = config.max_position_embeddings
584
+
585
+ if (self.head_dim * self.num_heads) != self.hidden_size:
586
+ raise ValueError(
587
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
588
+ f" and `num_heads`: {self.num_heads}).")
589
+ if config.lora_cfg is None:
590
+ self.q_proj = nn.Linear(self.hidden_size,
591
+ self.num_heads * self.head_dim,
592
+ bias=False)
593
+ self.k_proj = nn.Linear(self.hidden_size,
594
+ self.num_heads * self.head_dim,
595
+ bias=False)
596
+ self.v_proj = nn.Linear(self.hidden_size,
597
+ self.num_heads * self.head_dim,
598
+ bias=False)
599
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim,
600
+ self.hidden_size,
601
+ bias=False)
602
+ else:
603
+ lora_cfg = config.lora_cfg
604
+ if 'q' in lora_cfg['learn_param']:
605
+ self.q_proj = LoRALinear(self.hidden_size,
606
+ self.num_heads * self.head_dim,
607
+ bias=False,
608
+ **lora_cfg)
609
+ else:
610
+ self.q_proj = nn.Linear(self.hidden_size,
611
+ self.num_heads * self.head_dim,
612
+ bias=False)
613
+
614
+ if 'k' in lora_cfg['learn_param']:
615
+ self.k_proj = LoRALinear(self.hidden_size,
616
+ self.num_heads * self.head_dim,
617
+ bias=False,
618
+ **lora_cfg)
619
+ else:
620
+ self.k_proj = nn.Linear(self.hidden_size,
621
+ self.num_heads * self.head_dim,
622
+ bias=False)
623
+
624
+ if 'v' in lora_cfg['learn_param']:
625
+ self.v_proj = LoRALinear(self.hidden_size,
626
+ self.num_heads * self.head_dim,
627
+ bias=False,
628
+ **lora_cfg)
629
+ else:
630
+ self.v_proj = nn.Linear(self.hidden_size,
631
+ self.num_heads * self.head_dim,
632
+ bias=False)
633
+
634
+ if 'o' in lora_cfg['learn_param']:
635
+ self.o_proj = LoRALinear(self.num_heads * self.head_dim,
636
+ self.hidden_size,
637
+ bias=False,
638
+ **lora_cfg)
639
+ else:
640
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim,
641
+ self.hidden_size,
642
+ bias=False)
643
+
644
+ self.rotary_emb = InternLMRotaryEmbedding(
645
+ self.head_dim,
646
+ max_position_embeddings=self.max_position_embeddings)
647
+
648
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
649
+ return tensor.view(bsz, seq_len, self.num_heads,
650
+ self.head_dim).transpose(1, 2).contiguous()
651
+
652
+ def forward(
653
+ self,
654
+ hidden_states: torch.Tensor,
655
+ attention_mask: Optional[torch.Tensor] = None,
656
+ position_ids: Optional[torch.LongTensor] = None,
657
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
658
+ output_attentions: bool = False,
659
+ use_cache: bool = False,
660
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
661
+ Optional[Tuple[torch.Tensor]]]:
662
+ bsz, q_len, _ = hidden_states.size()
663
+
664
+ query_states = self.q_proj(hidden_states).view(
665
+ bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
666
+ key_states = self.k_proj(hidden_states).view(
667
+ bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
668
+ value_states = self.v_proj(hidden_states).view(
669
+ bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
670
+
671
+ kv_seq_len = key_states.shape[-2]
672
+ if past_key_value is not None:
673
+ kv_seq_len += past_key_value[0].shape[-2]
674
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
675
+ query_states, key_states = apply_rotary_pos_emb(
676
+ query_states, key_states, cos, sin, position_ids)
677
+ # [bsz, nh, t, hd]
678
+
679
+ if past_key_value is not None:
680
+ # reuse k, v, self_attention
681
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
682
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
683
+
684
+ past_key_value = (key_states, value_states) if use_cache else None
685
+
686
+ attn_weights = torch.matmul(query_states, key_states.transpose(
687
+ 2, 3)) / math.sqrt(self.head_dim)
688
+
689
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
690
+ raise ValueError(
691
+ f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
692
+ f" {attn_weights.size()}")
693
+
694
+ if attention_mask is not None:
695
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
696
+ raise ValueError(
697
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
698
+ )
699
+ attn_weights = attn_weights + attention_mask
700
+ attn_weights = torch.max(
701
+ attn_weights,
702
+ torch.tensor(torch.finfo(attn_weights.dtype).min))
703
+
704
+ # upcast attention to fp32
705
+ attn_weights = nn.functional.softmax(attn_weights,
706
+ dim=-1,
707
+ dtype=torch.float32).to(
708
+ query_states.dtype)
709
+ attn_output = torch.matmul(attn_weights, value_states)
710
+
711
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
712
+ raise ValueError(
713
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
714
+ f" {attn_output.size()}")
715
+
716
+ attn_output = attn_output.transpose(1, 2)
717
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
718
+
719
+ attn_output = self.o_proj(attn_output)
720
+
721
+ if not output_attentions:
722
+ attn_weights = None
723
+
724
+ return attn_output, attn_weights, past_key_value
725
+
726
+
727
+ class InternLMDecoderLayer(nn.Module):
728
+ def __init__(self, config: InternLMXComposerConfig):
729
+ super().__init__()
730
+ self.hidden_size = config.hidden_size
731
+ if hasattr(config,
732
+ 'intern_converted_llm') and config.intern_converted_llm:
733
+ self.self_attn = InternConvertedInternLMAttention(config=config)
734
+ else:
735
+ self.self_attn = InternLMAttention(config=config)
736
+ self.mlp = InternLMMLP(
737
+ hidden_size=self.hidden_size,
738
+ intermediate_size=config.intermediate_size,
739
+ hidden_act=config.hidden_act,
740
+ config=config,
741
+ )
742
+ self.input_layernorm = InternLMRMSNorm(config.hidden_size,
743
+ eps=config.rms_norm_eps)
744
+ self.post_attention_layernorm = InternLMRMSNorm(
745
+ config.hidden_size, eps=config.rms_norm_eps)
746
+
747
+ def forward(
748
+ self,
749
+ hidden_states: torch.Tensor,
750
+ attention_mask: Optional[torch.Tensor] = None,
751
+ position_ids: Optional[torch.LongTensor] = None,
752
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
753
+ output_attentions: Optional[bool] = False,
754
+ use_cache: Optional[bool] = False,
755
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor,
756
+ torch.FloatTensor]]]:
757
+ """
758
+ Args:
759
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
760
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
761
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
762
+ output_attentions (`bool`, *optional*):
763
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
764
+ returned tensors for more detail.
765
+ use_cache (`bool`, *optional*):
766
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
767
+ (see `past_key_values`).
768
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
769
+ """
770
+
771
+ residual = hidden_states
772
+
773
+ hidden_states = self.input_layernorm(hidden_states)
774
+
775
+ # Self Attention
776
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
777
+ hidden_states=hidden_states,
778
+ attention_mask=attention_mask,
779
+ position_ids=position_ids,
780
+ past_key_value=past_key_value,
781
+ output_attentions=output_attentions,
782
+ use_cache=use_cache,
783
+ )
784
+ hidden_states = residual + hidden_states
785
+
786
+ # Fully Connected
787
+ residual = hidden_states
788
+ hidden_states = self.post_attention_layernorm(hidden_states)
789
+ hidden_states = self.mlp(hidden_states)
790
+ hidden_states = residual + hidden_states
791
+
792
+ outputs = (hidden_states, )
793
+
794
+ if output_attentions:
795
+ outputs += (self_attn_weights, )
796
+
797
+ if use_cache:
798
+ outputs += (present_key_value, )
799
+
800
+ return outputs
801
+
802
+
803
+ class InternLMPreTrainedModel(PreTrainedModel):
804
+ config_class = InternLMXComposerConfig
805
+ base_model_prefix = "model"
806
+ supports_gradient_checkpointing = True
807
+ _no_split_modules = ["InternLMDecoderLayer"]
808
+ _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
809
+
810
+ def _init_weights(self, module):
811
+ std = self.config.initializer_range
812
+ if isinstance(module, nn.Linear):
813
+ module.weight.data.normal_(mean=0.0, std=std)
814
+ if module.bias is not None:
815
+ module.bias.data.zero_()
816
+ elif isinstance(module, nn.Embedding):
817
+ module.weight.data.normal_(mean=0.0, std=std)
818
+ if module.padding_idx is not None:
819
+ module.weight.data[module.padding_idx].zero_()
820
+
821
+ def _set_gradient_checkpointing(self, module, value=False):
822
+ if isinstance(module, InternLMModel):
823
+ module.gradient_checkpointing = value
824
+
825
+
826
+ class InternLMModel(InternLMPreTrainedModel):
827
+ """
828
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`InternLMDecoderLayer`]
829
+
830
+ Args:
831
+ config: InternLMXComposerConfig
832
+ """
833
+ def __init__(self, config: InternLMXComposerConfig):
834
+ super().__init__(config)
835
+ self.padding_idx = config.pad_token_id
836
+ self.vocab_size = config.vocab_size
837
+
838
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size,
839
+ self.padding_idx)
840
+ self.layers = nn.ModuleList([
841
+ InternLMDecoderLayer(config)
842
+ for _ in range(config.num_hidden_layers)
843
+ ])
844
+ self.norm = InternLMRMSNorm(config.hidden_size,
845
+ eps=config.rms_norm_eps)
846
+
847
+ self.gradient_checkpointing = False
848
+ # Initialize weights and apply final processing
849
+ self.post_init()
850
+
851
+ def get_input_embeddings(self):
852
+ return self.embed_tokens
853
+
854
+ def set_input_embeddings(self, value):
855
+ self.embed_tokens = value
856
+
857
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
858
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape,
859
+ inputs_embeds, past_key_values_length):
860
+ # create causal mask
861
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
862
+ combined_attention_mask = None
863
+ if input_shape[-1] > 1:
864
+ combined_attention_mask = _make_causal_mask(
865
+ input_shape,
866
+ inputs_embeds.dtype,
867
+ device=inputs_embeds.device,
868
+ past_key_values_length=past_key_values_length,
869
+ )
870
+
871
+ if attention_mask is not None:
872
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
873
+ expanded_attn_mask = _expand_mask(attention_mask,
874
+ inputs_embeds.dtype,
875
+ tgt_len=input_shape[-1]).to(
876
+ inputs_embeds.device)
877
+ combined_attention_mask = (expanded_attn_mask
878
+ if combined_attention_mask is None else
879
+ expanded_attn_mask +
880
+ combined_attention_mask)
881
+
882
+ return combined_attention_mask
883
+
884
+ def forward(
885
+ self,
886
+ input_ids: torch.LongTensor = None,
887
+ attention_mask: Optional[torch.Tensor] = None,
888
+ position_ids: Optional[torch.LongTensor] = None,
889
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
890
+ inputs_embeds: Optional[torch.FloatTensor] = None,
891
+ query_embeds: Optional[torch.FloatTensor] = None,
892
+ use_cache: Optional[bool] = None,
893
+ output_attentions: Optional[bool] = None,
894
+ output_hidden_states: Optional[bool] = None,
895
+ return_dict: Optional[bool] = None,
896
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
897
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
898
+ output_hidden_states = (output_hidden_states
899
+ if output_hidden_states is not None else
900
+ self.config.output_hidden_states)
901
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
902
+
903
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
904
+
905
+ # retrieve input_ids and inputs_embeds
906
+ if input_ids is not None and inputs_embeds is not None:
907
+ raise ValueError(
908
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
909
+ )
910
+ elif input_ids is not None:
911
+ batch_size, seq_length = input_ids.shape
912
+ elif inputs_embeds is not None:
913
+ batch_size, seq_length, _ = inputs_embeds.shape
914
+ else:
915
+ raise ValueError(
916
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
917
+ )
918
+
919
+ if inputs_embeds is None:
920
+ inputs_embeds = self.embed_tokens(input_ids)
921
+ if query_embeds is not None:
922
+ inputs_embeds = torch.cat([query_embeds, inputs_embeds], dim=1)
923
+ batch_size, seq_length, _ = inputs_embeds.shape
924
+
925
+ seq_length_with_past = seq_length
926
+ past_key_values_length = 0
927
+
928
+ if past_key_values is not None:
929
+ past_key_values_length = past_key_values[0][0].shape[2]
930
+ seq_length_with_past = seq_length_with_past + past_key_values_length
931
+
932
+ if position_ids is None:
933
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
934
+ position_ids = torch.arange(past_key_values_length,
935
+ seq_length + past_key_values_length,
936
+ dtype=torch.long,
937
+ device=device)
938
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
939
+ else:
940
+ position_ids = position_ids.view(-1, seq_length).long()
941
+
942
+ # embed positions
943
+ if attention_mask is None:
944
+ attention_mask = torch.ones((batch_size, seq_length_with_past),
945
+ dtype=torch.bool,
946
+ device=inputs_embeds.device)
947
+ attention_mask = self._prepare_decoder_attention_mask(
948
+ attention_mask, (batch_size, seq_length), inputs_embeds,
949
+ past_key_values_length)
950
+
951
+ hidden_states = inputs_embeds
952
+
953
+ if self.gradient_checkpointing and self.training:
954
+ if use_cache:
955
+ logger.warning_once(
956
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
957
+ )
958
+ use_cache = False
959
+
960
+ # decoder layers
961
+ all_hidden_states = () if output_hidden_states else None
962
+ all_self_attns = () if output_attentions else None
963
+ next_decoder_cache = () if use_cache else None
964
+
965
+ for idx, decoder_layer in enumerate(self.layers):
966
+ if output_hidden_states:
967
+ all_hidden_states += (hidden_states, )
968
+
969
+ past_key_value = past_key_values[
970
+ idx] if past_key_values is not None else None
971
+
972
+ if self.gradient_checkpointing and self.training:
973
+
974
+ def create_custom_forward(module):
975
+ def custom_forward(*inputs):
976
+ # None for past_key_value
977
+ return module(*inputs, output_attentions, None)
978
+
979
+ return custom_forward
980
+
981
+ layer_outputs = torch.utils.checkpoint.checkpoint(
982
+ create_custom_forward(decoder_layer),
983
+ hidden_states,
984
+ attention_mask,
985
+ position_ids,
986
+ None,
987
+ )
988
+ else:
989
+ layer_outputs = decoder_layer(
990
+ hidden_states,
991
+ attention_mask=attention_mask,
992
+ position_ids=position_ids,
993
+ past_key_value=past_key_value,
994
+ output_attentions=output_attentions,
995
+ use_cache=use_cache,
996
+ )
997
+
998
+ hidden_states = layer_outputs[0]
999
+
1000
+ if use_cache:
1001
+ next_decoder_cache += (
1002
+ layer_outputs[2 if output_attentions else 1], )
1003
+
1004
+ if output_attentions:
1005
+ all_self_attns += (layer_outputs[1], )
1006
+
1007
+ hidden_states = self.norm(hidden_states)
1008
+
1009
+ # add hidden states from the last decoder layer
1010
+ if output_hidden_states:
1011
+ all_hidden_states += (hidden_states, )
1012
+
1013
+ next_cache = next_decoder_cache if use_cache else None
1014
+ if not return_dict:
1015
+ return tuple(
1016
+ v for v in
1017
+ [hidden_states, next_cache, all_hidden_states, all_self_attns]
1018
+ if v is not None)
1019
+ return BaseModelOutputWithPast(
1020
+ last_hidden_state=hidden_states,
1021
+ past_key_values=next_cache,
1022
+ hidden_states=all_hidden_states,
1023
+ attentions=all_self_attns,
1024
+ )
1025
+
1026
+
1027
+ class InternLMForCausalLM(InternLMPreTrainedModel):
1028
+ lora_cfg = None # init in MiniGPT4
1029
+
1030
+ def __init__(self, config):
1031
+ super().__init__(config)
1032
+ # TODO: find a way to explicitly initialize InternLM
1033
+ setattr(config, 'lora_cfg', self.lora_cfg)
1034
+
1035
+ if hasattr(config, 'kqvo_bias'):
1036
+ setattr(config, 'kqvo_bias', config.kqvo_bias)
1037
+ else:
1038
+ setattr(config, 'kqvo_bias', False)
1039
+ self.model = InternLMModel(config)
1040
+
1041
+ self.lm_head = nn.Linear(config.hidden_size,
1042
+ config.vocab_size,
1043
+ bias=False)
1044
+ if hasattr(config, 'ex_size'):
1045
+ self.ex_size = config.ex_size
1046
+ else:
1047
+ self.ex_size = 0
1048
+
1049
+ if hasattr(config, 'sp_id'):
1050
+ self.sp_id = config.sp_id
1051
+ else:
1052
+ self.sp_id = -1
1053
+
1054
+ # Initialize weights and apply final processing
1055
+ self.post_init()
1056
+
1057
+ @classmethod
1058
+ def from_pretrained(cls,
1059
+ pretrained_model_name_or_path,
1060
+ llm_cfg=None,
1061
+ *model_args,
1062
+ **kwargs):
1063
+ if llm_cfg:
1064
+ if 'torch_dtype' in kwargs:
1065
+ llm_cfg.torch_dtype = kwargs['torch_dtype']
1066
+ if 'load_in_8bit' in kwargs:
1067
+ llm_cfg.load_in_8bit = kwargs['load_in_8bit']
1068
+ if 'device_map' in kwargs:
1069
+ llm_cfg.device_map = kwargs['device_map']
1070
+ return cls._from_config(llm_cfg)
1071
+ else:
1072
+ return super().from_pretrained(pretrained_model_name_or_path,
1073
+ *model_args, **kwargs)
1074
+
1075
+ def get_input_embeddings(self):
1076
+ return self.model.embed_tokens
1077
+
1078
+ def set_input_embeddings(self, value):
1079
+ self.model.embed_tokens = value
1080
+
1081
+ def get_output_embeddings(self):
1082
+ return self.lm_head
1083
+
1084
+ def set_output_embeddings(self, new_embeddings):
1085
+ self.lm_head = new_embeddings
1086
+
1087
+ def set_decoder(self, decoder):
1088
+ self.model = decoder
1089
+
1090
+ def get_decoder(self):
1091
+ return self.model
1092
+
1093
+ def forward(
1094
+ self,
1095
+ input_ids: torch.LongTensor = None,
1096
+ attention_mask: Optional[torch.Tensor] = None,
1097
+ position_ids: Optional[torch.LongTensor] = None,
1098
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1099
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1100
+ query_embeds: Optional[torch.FloatTensor] = None,
1101
+ labels: Optional[torch.LongTensor] = None,
1102
+ use_cache: Optional[bool] = None,
1103
+ output_attentions: Optional[bool] = None,
1104
+ output_hidden_states: Optional[bool] = None,
1105
+ return_dict: Optional[bool] = None,
1106
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1107
+ r"""
1108
+ Args:
1109
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1110
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1111
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1112
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1113
+
1114
+ Returns:
1115
+
1116
+ Example:
1117
+
1118
+ ```python
1119
+ >>> from transformers import AutoTokenizer, InternLMForCausalLM
1120
+
1121
+ >>> model = InternLMForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1122
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1123
+
1124
+ >>> prompt = "Hey, are you consciours? Can you talk to me?"
1125
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1126
+
1127
+ >>> # Generate
1128
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1129
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1130
+ "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you."
1131
+ ```"""
1132
+
1133
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1134
+ output_hidden_states = (output_hidden_states
1135
+ if output_hidden_states is not None else
1136
+ self.config.output_hidden_states)
1137
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1138
+
1139
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1140
+ outputs = self.model(
1141
+ input_ids=input_ids,
1142
+ attention_mask=attention_mask,
1143
+ position_ids=position_ids,
1144
+ past_key_values=past_key_values,
1145
+ inputs_embeds=inputs_embeds,
1146
+ query_embeds=query_embeds,
1147
+ use_cache=use_cache,
1148
+ output_attentions=output_attentions,
1149
+ output_hidden_states=output_hidden_states,
1150
+ return_dict=return_dict,
1151
+ )
1152
+
1153
+ hidden_states = outputs[0]
1154
+ logits = self.lm_head(hidden_states)
1155
+
1156
+ loss = None
1157
+ if labels is not None:
1158
+ # Shift so that tokens < n predict n
1159
+ shift_logits = logits[..., :-1, :].contiguous()
1160
+ shift_labels = labels[..., 1:].contiguous()
1161
+ # Flatten the tokens
1162
+
1163
+ loss_fct = CrossEntropyLoss(reduce=False)
1164
+ loss_reduce = CrossEntropyLoss()
1165
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1166
+ shift_labels = shift_labels.view(-1)
1167
+ shift_labels = shift_labels.to(shift_logits.device)
1168
+ ###
1169
+ if self.sp_id >= 0:
1170
+ ori_mask = (shift_labels != self.sp_id).float()
1171
+ ori_mask = ori_mask * (shift_labels >= 0).float()
1172
+ local_mask = (shift_labels == self.sp_id).float()
1173
+ else:
1174
+ ori_mask = (shift_labels <
1175
+ self.config.vocab_size - self.ex_size).float()
1176
+ ori_mask = ori_mask * (shift_labels >= 0).float()
1177
+ local_mask = (shift_labels >=
1178
+ self.config.vocab_size - self.ex_size).float()
1179
+
1180
+ # Enable model parallelism
1181
+
1182
+ loss = loss_reduce(shift_logits, shift_labels)
1183
+
1184
+ loss_all = loss_fct(shift_logits, shift_labels)
1185
+ loss_o = (loss_all * ori_mask).sum() / ori_mask.sum()
1186
+ if torch.sum(local_mask) == 0:
1187
+ loss_l = loss_o * 0
1188
+ else:
1189
+ loss_l = (loss_all * local_mask).sum() / local_mask.sum()
1190
+
1191
+ if not return_dict:
1192
+ output = (logits, ) + outputs[1:]
1193
+ return (loss, ) + output if loss is not None else output
1194
+
1195
+ if (self.ex_size > 0 or self.sp_id >= 0) and labels is not None:
1196
+ return loss, loss_o, loss_l
1197
+
1198
+ return CausalLMOutputWithPast(
1199
+ loss=loss,
1200
+ logits=logits,
1201
+ past_key_values=outputs.past_key_values,
1202
+ hidden_states=outputs.hidden_states,
1203
+ attentions=outputs.attentions,
1204
+ )
1205
+
1206
+ def prepare_inputs_for_generation(self,
1207
+ input_ids,
1208
+ query_embeds=None,
1209
+ past_key_values=None,
1210
+ attention_mask=None,
1211
+ inputs_embeds=None,
1212
+ **kwargs):
1213
+ if past_key_values:
1214
+ input_ids = input_ids[:, -1:]
1215
+
1216
+ position_ids = kwargs.get("position_ids", None)
1217
+ if attention_mask is not None and position_ids is None:
1218
+ # create position_ids on the fly for batch generation
1219
+ position_ids = attention_mask.long().cumsum(-1) - 1
1220
+ position_ids.masked_fill_(attention_mask == 0, 1)
1221
+ if past_key_values:
1222
+ position_ids = position_ids[:, -1].unsqueeze(-1)
1223
+ query_embeds = None
1224
+
1225
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1226
+ if inputs_embeds is not None and past_key_values is None:
1227
+ model_inputs = {"inputs_embeds": inputs_embeds}
1228
+ else:
1229
+ model_inputs = {"input_ids": input_ids}
1230
+
1231
+ model_inputs.update({
1232
+ "position_ids": position_ids,
1233
+ "query_embeds": query_embeds,
1234
+ "past_key_values": past_key_values,
1235
+ "use_cache": kwargs.get("use_cache"),
1236
+ "attention_mask": attention_mask,
1237
+ })
1238
+ return model_inputs
1239
+
1240
+ @staticmethod
1241
+ def _reorder_cache(past_key_values, beam_idx):
1242
+ reordered_past = ()
1243
+ for layer_past in past_key_values:
1244
+ reordered_past += (tuple(
1245
+ past_state.index_select(0, beam_idx)
1246
+ for past_state in layer_past), )
1247
+ return reordered_past
modeling_InternLM_XComposer.py ADDED
@@ -0,0 +1,259 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import os
3
+ import sys
4
+
5
+ dir_path = os.path.dirname(os.path.realpath(__file__))
6
+ sys.path.insert(0, dir_path)
7
+
8
+ import contextlib
9
+
10
+ import torch.utils.checkpoint
11
+ from torch.nn import LayerNorm
12
+ from torchvision import transforms
13
+ from torchvision.transforms.functional import InterpolationMode
14
+ from PIL import Image
15
+
16
+ from .modeling_perceive_sampler import BertConfig, BertLMHeadModel
17
+ from .modeling_vit import *
18
+ from .modeling_InternLM import *
19
+ from .modeling_utils import *
20
+
21
+ from transformers.utils import logging
22
+ logger = logging.get_logger(__name__)
23
+
24
+
25
+ class InternLMXComposerForCausalLM(PreTrainedModel):
26
+ config_class = InternLMXComposerConfig
27
+ _auto_class = "AutoModelForCausalLM"
28
+
29
+ gen_config = dict(
30
+ num_beams=5,
31
+ do_sample=False,
32
+ min_length=1,
33
+ repetition_penalty=1.5,
34
+ length_penalty=1.0,
35
+ temperature=1.0,
36
+ max_new_tokens=200,
37
+ )
38
+
39
+ def __init__(self, config):
40
+ super().__init__(config)
41
+
42
+ print('Init VIT ... ', end='')
43
+ self.visual_encoder = create_eva_vit_g()
44
+ self.ln_vision = LayerNorm(self.visual_encoder.num_features)
45
+ print('Done')
46
+
47
+ print('Init Perceive Sampler ... ', end='')
48
+ with all_logging_disabled():
49
+ self.Qformer, self.query_tokens = self.init_qformer(
50
+ config.num_query_token, self.visual_encoder.num_features)
51
+ self.Qformer.bert.embeddings.word_embeddings = None
52
+ self.Qformer.bert.embeddings.position_embeddings = None
53
+ for layer in self.Qformer.bert.encoder.layer:
54
+ layer.output = None
55
+ layer.intermediate = None
56
+ self.Qformer.cls = None
57
+ print('Done')
58
+
59
+ print('Init InternLM ... ', end='')
60
+ self.flag_image_start = nn.Parameter(torch.zeros([1, 1, 4096]))
61
+ self.flag_image_end = nn.Parameter(torch.zeros([1, 1, 4096]))
62
+ self.flag_image_start.requires_grad = False
63
+ self.flag_image_end.requires_grad = False
64
+
65
+ internlm_lora = config.internlm_lora
66
+ self.internlm_lora = internlm_lora
67
+ setattr(InternLMForCausalLM, 'lora_cfg', internlm_lora)
68
+
69
+ if int(torch.__version__[0]) == 1:
70
+ self.internlm_model = InternLMForCausalLM._from_config(config).to(
71
+ torch.float16)
72
+ else:
73
+ assert int(torch.__version__[0]) == 2
74
+ # speed up init llm
75
+ with torch.device('meta'):
76
+ self.internlm_model = InternLMForCausalLM._from_config(config)
77
+ self.internlm_model.to_empty(device=config.device).to(torch.float16)
78
+ for n, m in self.internlm_model.named_modules():
79
+ if 'lora' in n:
80
+ m.float()
81
+
82
+ self.internlm_proj = nn.Linear(self.Qformer.config.hidden_size,
83
+ self.internlm_model.config.hidden_size)
84
+ print('Done')
85
+
86
+ self.vis_processor = transforms.Compose([
87
+ transforms.Resize((224, 224),
88
+ interpolation=InterpolationMode.BICUBIC),
89
+ transforms.ToTensor(),
90
+ transforms.Normalize((0.48145466, 0.4578275, 0.40821073),
91
+ (0.26862954, 0.26130258, 0.27577711)),
92
+ ])
93
+
94
+ self.tokenizer = None
95
+
96
+ @property
97
+ def eoh(self):
98
+ return self.tokenizer.decode(torch.Tensor([103027]),
99
+ skip_special_tokens=True)
100
+
101
+ @property
102
+ def eoa(self):
103
+ return self.tokenizer.decode(torch.Tensor([103028]),
104
+ skip_special_tokens=True)
105
+
106
+ def maybe_autocast(self, dtype=torch.float16):
107
+ # if on cpu, don't use autocast
108
+ # if on gpu, use autocast with dtype if provided, otherwise use torch.float16
109
+ enable_autocast = self.device != torch.device("cpu")
110
+
111
+ if enable_autocast:
112
+ return torch.cuda.amp.autocast(dtype=dtype)
113
+ else:
114
+ return contextlib.nullcontext()
115
+
116
+ @classmethod
117
+ def init_qformer(cls,
118
+ num_query_token,
119
+ vision_width,
120
+ cross_attention_freq=2,
121
+ pretrain=True):
122
+ encoder_config = BertConfig.from_pretrained("bert-base-uncased")
123
+ encoder_config.encoder_width = vision_width
124
+ # insert cross-attention layer every other block
125
+ encoder_config.add_cross_attention = True
126
+ encoder_config.cross_attention_freq = cross_attention_freq
127
+ encoder_config.query_length = num_query_token
128
+ if pretrain:
129
+ Qformer = BertLMHeadModel.from_pretrained("bert-base-uncased",
130
+ config=encoder_config)
131
+ else:
132
+ Qformer = BertLMHeadModel(config=encoder_config)
133
+ query_tokens = nn.Parameter(
134
+ torch.zeros(1, num_query_token, encoder_config.hidden_size))
135
+ query_tokens.data.normal_(mean=0.0,
136
+ std=encoder_config.initializer_range)
137
+ return Qformer, query_tokens
138
+
139
+ def encode_img(self, image):
140
+ if image is None:
141
+ return None
142
+ if isinstance(image, str):
143
+ image = Image.open(image).convert("RGB")
144
+ image = self.vis_processor(image).unsqueeze(0).to(self.device)
145
+ else:
146
+ assert isinstance(image, torch.Tensor)
147
+ device = image.device
148
+ with self.maybe_autocast():
149
+ image_embeds = self.ln_vision(
150
+ self.visual_encoder(image)).to(device)
151
+ image_atts = torch.ones(image_embeds.size()[:-1],
152
+ dtype=torch.long).to(device)
153
+ query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1,
154
+ -1)
155
+ query_output = self.Qformer.bert(
156
+ query_embeds=query_tokens,
157
+ encoder_hidden_states=image_embeds,
158
+ encoder_attention_mask=image_atts,
159
+ return_dict=True,
160
+ )
161
+ inputs_internlm = self.internlm_proj(query_output.last_hidden_state)
162
+ inputs_internlm = torch.cat([
163
+ self.flag_image_start.expand(inputs_internlm.shape[0], -1, -1),
164
+ inputs_internlm,
165
+ self.flag_image_end.expand(inputs_internlm.shape[0], -1, -1)
166
+ ],
167
+ dim=1)
168
+ return inputs_internlm
169
+
170
+ def encode_text(self, text, add_special_tokens=False):
171
+ text_token_ids = self.tokenizer(
172
+ text,
173
+ return_tensors='pt',
174
+ add_special_tokens=add_special_tokens,
175
+ ).input_ids.to(self.device)
176
+ text_embeds = self.internlm_model.model.embed_tokens(text_token_ids)
177
+ return text_embeds
178
+
179
+ def decode_text(self, out_embeds):
180
+ out_text = self.tokenizer.batch_decode(out_embeds,
181
+ skip_special_tokens=True)[0]
182
+ out_text = out_text.split(self.eoa)[0]
183
+ return out_text
184
+
185
+ def wrap_text(self, user_text, bot_text='', add_special=True):
186
+ if add_special:
187
+ eoh = self.eoh
188
+ else:
189
+ eoh = ''
190
+ text = f' <|User|>:{user_text} \n{eoh} <|Bot|>:{bot_text}'
191
+ return text
192
+
193
+ def get_gen_args(self, **kwargs):
194
+ new_kargs = copy.deepcopy(self.gen_config)
195
+ new_kargs.update(kwargs)
196
+ return new_kargs
197
+
198
+ def generate(self, text, image=None, **kwargs):
199
+ text_embeds = self.encode_text(text)
200
+ img_embeds = self.encode_img(image)
201
+ prompt_embeds = self.wrap_prompt(text_embeds, img_embeds)
202
+ out_embeds = self.internlm_model.generate(inputs_embeds=prompt_embeds,
203
+ **self.get_gen_args(**kwargs))
204
+ out_text = self.decode_text(out_embeds)
205
+ return out_text
206
+
207
+ def chat(self, text, image=None, history=None, **kwargs):
208
+ text_embeds = self.encode_text(text)
209
+ img_embeds = self.encode_img(image)
210
+ prompt_embeds = self.wrap_prompt(text_embeds,
211
+ img_embeds,
212
+ history=history)
213
+ out_embeds = self.internlm_model.generate(inputs_embeds=prompt_embeds,
214
+ **self.get_gen_args(**kwargs))
215
+ out_text = self.decode_text(out_embeds)
216
+
217
+ # trunc at eoh and eoa
218
+ clean_out_text_token_ids = self.tokenizer(
219
+ out_text, return_tensors='pt').input_ids.to(self.device)
220
+ clean_out_text_embeds = self.internlm_model.model.embed_tokens(
221
+ clean_out_text_token_ids)
222
+ clean_prompt_embeds = self.wrap_prompt(text_embeds,
223
+ img_embeds,
224
+ add_special=False)
225
+ cur_history = torch.cat([clean_prompt_embeds, clean_out_text_embeds],
226
+ dim=1)
227
+ if history is None:
228
+ history = []
229
+ history.append(cur_history)
230
+ return out_text, history
231
+
232
+ def wrap_prompt(self,
233
+ text_embeds,
234
+ img_embeds=None,
235
+ history=None,
236
+ add_special=True):
237
+ if add_special:
238
+ prompt_segs = [' <|User|>:', f'\n{self.eoh} <|Bot|>:']
239
+ else:
240
+ prompt_segs = [' <|User|>:', ' <|Bot|>:'] # used in wrap history
241
+ prompt_seg_embeds = []
242
+ for i, seg in enumerate(prompt_segs):
243
+ if history is not None:
244
+ add_special_tokens = False
245
+ else:
246
+ add_special_tokens = i == 0
247
+ seg_embeds = self.encode_text(
248
+ seg, add_special_tokens=add_special_tokens)
249
+ prompt_seg_embeds.append(seg_embeds)
250
+ if img_embeds is None:
251
+ img_embeds = text_embeds.new_empty(text_embeds.size(0), 0,
252
+ text_embeds.size(-1))
253
+ prompt_seg_embeds = [
254
+ prompt_seg_embeds[0], img_embeds, text_embeds, prompt_seg_embeds[1]
255
+ ]
256
+ prompt_embeds = torch.cat(prompt_seg_embeds, dim=1)
257
+ if history is not None:
258
+ prompt_embeds = torch.cat([*history, prompt_embeds], dim=1)
259
+ return prompt_embeds
modeling_perceive_sampler.py ADDED
@@ -0,0 +1,1193 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ * Copyright (c) 2023, salesforce.com, inc.
3
+ * All rights reserved.
4
+ * SPDX-License-Identifier: BSD-3-Clause
5
+ * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6
+ * By Junnan Li
7
+ * Based on huggingface code base
8
+ * https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert
9
+ """
10
+
11
+ import math
12
+ from typing import Tuple
13
+
14
+ import torch
15
+ import torch.utils.checkpoint
16
+ from torch import Tensor, device
17
+ from torch import nn
18
+ from torch.nn import CrossEntropyLoss
19
+ from transformers.activations import ACT2FN
20
+ from transformers.modeling_outputs import (
21
+ BaseModelOutputWithPastAndCrossAttentions,
22
+ BaseModelOutputWithPoolingAndCrossAttentions,
23
+ CausalLMOutputWithCrossAttentions,
24
+ MaskedLMOutput,
25
+ )
26
+ from transformers.modeling_utils import (
27
+ PreTrainedModel,
28
+ apply_chunking_to_forward,
29
+ find_pruneable_heads_and_indices,
30
+ prune_linear_layer,
31
+ )
32
+ from transformers.models.bert.configuration_bert import BertConfig
33
+ from transformers.utils import logging
34
+
35
+ logger = logging.get_logger(__name__)
36
+
37
+
38
+ class BertEmbeddings(nn.Module):
39
+ """Construct the embeddings from word and position embeddings."""
40
+ def __init__(self, config):
41
+ super().__init__()
42
+ self.word_embeddings = nn.Embedding(config.vocab_size,
43
+ config.hidden_size,
44
+ padding_idx=config.pad_token_id)
45
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings,
46
+ config.hidden_size)
47
+
48
+ # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
49
+ # any TensorFlow checkpoint file
50
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
51
+ eps=config.layer_norm_eps)
52
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
53
+
54
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
55
+ self.register_buffer(
56
+ "position_ids",
57
+ torch.arange(config.max_position_embeddings).expand((1, -1)))
58
+ self.position_embedding_type = getattr(config,
59
+ "position_embedding_type",
60
+ "absolute")
61
+
62
+ self.config = config
63
+
64
+ def forward(
65
+ self,
66
+ input_ids=None,
67
+ position_ids=None,
68
+ query_embeds=None,
69
+ past_key_values_length=0,
70
+ ):
71
+ if input_ids is not None:
72
+ seq_length = input_ids.size()[1]
73
+ else:
74
+ seq_length = 0
75
+
76
+ if position_ids is None:
77
+ position_ids = self.position_ids[:, past_key_values_length:
78
+ seq_length +
79
+ past_key_values_length].clone()
80
+
81
+ if input_ids is not None:
82
+ embeddings = self.word_embeddings(input_ids)
83
+ if self.position_embedding_type == "absolute":
84
+ position_embeddings = self.position_embeddings(position_ids)
85
+ embeddings = embeddings + position_embeddings
86
+
87
+ if query_embeds is not None:
88
+ embeddings = torch.cat((query_embeds, embeddings), dim=1)
89
+ else:
90
+ embeddings = query_embeds
91
+
92
+ embeddings = self.LayerNorm(embeddings)
93
+ embeddings = self.dropout(embeddings)
94
+ return embeddings
95
+
96
+
97
+ class BertSelfAttention(nn.Module):
98
+ def __init__(self, config, is_cross_attention):
99
+ super().__init__()
100
+ self.config = config
101
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
102
+ config, "embedding_size"):
103
+ raise ValueError(
104
+ "The hidden size (%d) is not a multiple of the number of attention "
105
+ "heads (%d)" %
106
+ (config.hidden_size, config.num_attention_heads))
107
+
108
+ self.num_attention_heads = config.num_attention_heads
109
+ self.attention_head_size = int(config.hidden_size /
110
+ config.num_attention_heads)
111
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
112
+
113
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
114
+ if is_cross_attention:
115
+ self.key = nn.Linear(config.encoder_width, self.all_head_size)
116
+ self.value = nn.Linear(config.encoder_width, self.all_head_size)
117
+ else:
118
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
119
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
120
+
121
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
122
+ self.position_embedding_type = getattr(config,
123
+ "position_embedding_type",
124
+ "absolute")
125
+ if (self.position_embedding_type == "relative_key"
126
+ or self.position_embedding_type == "relative_key_query"):
127
+ self.max_position_embeddings = config.max_position_embeddings
128
+ self.distance_embedding = nn.Embedding(
129
+ 2 * config.max_position_embeddings - 1,
130
+ self.attention_head_size)
131
+ self.save_attention = False
132
+
133
+ def save_attn_gradients(self, attn_gradients):
134
+ self.attn_gradients = attn_gradients
135
+
136
+ def get_attn_gradients(self):
137
+ return self.attn_gradients
138
+
139
+ def save_attention_map(self, attention_map):
140
+ self.attention_map = attention_map
141
+
142
+ def get_attention_map(self):
143
+ return self.attention_map
144
+
145
+ def transpose_for_scores(self, x):
146
+ new_x_shape = x.size()[:-1] + (
147
+ self.num_attention_heads,
148
+ self.attention_head_size,
149
+ )
150
+ x = x.view(*new_x_shape)
151
+ return x.permute(0, 2, 1, 3)
152
+
153
+ def forward(
154
+ self,
155
+ hidden_states,
156
+ attention_mask=None,
157
+ head_mask=None,
158
+ encoder_hidden_states=None,
159
+ encoder_attention_mask=None,
160
+ past_key_value=None,
161
+ output_attentions=False,
162
+ ):
163
+
164
+ # If this is instantiated as a cross-attention module, the keys
165
+ # and values come from an encoder; the attention mask needs to be
166
+ # such that the encoder's padding tokens are not attended to.
167
+ is_cross_attention = encoder_hidden_states is not None
168
+
169
+ if is_cross_attention:
170
+ key_layer = self.transpose_for_scores(
171
+ self.key(encoder_hidden_states))
172
+ value_layer = self.transpose_for_scores(
173
+ self.value(encoder_hidden_states))
174
+ attention_mask = encoder_attention_mask
175
+ elif past_key_value is not None:
176
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
177
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
178
+ key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
179
+ value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
180
+ else:
181
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
182
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
183
+
184
+ mixed_query_layer = self.query(hidden_states)
185
+
186
+ query_layer = self.transpose_for_scores(mixed_query_layer)
187
+
188
+ past_key_value = (key_layer, value_layer)
189
+
190
+ # Take the dot product between "query" and "key" to get the raw attention scores.
191
+ attention_scores = torch.matmul(query_layer,
192
+ key_layer.transpose(-1, -2))
193
+
194
+ if (self.position_embedding_type == "relative_key"
195
+ or self.position_embedding_type == "relative_key_query"):
196
+ seq_length = hidden_states.size()[1]
197
+ position_ids_l = torch.arange(seq_length,
198
+ dtype=torch.long,
199
+ device=hidden_states.device).view(
200
+ -1, 1)
201
+ position_ids_r = torch.arange(seq_length,
202
+ dtype=torch.long,
203
+ device=hidden_states.device).view(
204
+ 1, -1)
205
+ distance = position_ids_l - position_ids_r
206
+ positional_embedding = self.distance_embedding(
207
+ distance + self.max_position_embeddings - 1)
208
+ positional_embedding = positional_embedding.to(
209
+ dtype=query_layer.dtype) # fp16 compatibility
210
+
211
+ if self.position_embedding_type == "relative_key":
212
+ relative_position_scores = torch.einsum(
213
+ "bhld,lrd->bhlr", query_layer, positional_embedding)
214
+ attention_scores = attention_scores + relative_position_scores
215
+ elif self.position_embedding_type == "relative_key_query":
216
+ relative_position_scores_query = torch.einsum(
217
+ "bhld,lrd->bhlr", query_layer, positional_embedding)
218
+ relative_position_scores_key = torch.einsum(
219
+ "bhrd,lrd->bhlr", key_layer, positional_embedding)
220
+ attention_scores = (attention_scores +
221
+ relative_position_scores_query +
222
+ relative_position_scores_key)
223
+
224
+ attention_scores = attention_scores / math.sqrt(
225
+ self.attention_head_size)
226
+ if attention_mask is not None:
227
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
228
+ attention_scores = attention_scores + attention_mask
229
+
230
+ # Normalize the attention scores to probabilities.
231
+ attention_probs = nn.Softmax(dim=-1)(attention_scores)
232
+
233
+ if is_cross_attention and self.save_attention:
234
+ self.save_attention_map(attention_probs)
235
+ attention_probs.register_hook(self.save_attn_gradients)
236
+
237
+ # This is actually dropping out entire tokens to attend to, which might
238
+ # seem a bit unusual, but is taken from the original Transformer paper.
239
+ attention_probs_dropped = self.dropout(attention_probs)
240
+
241
+ # Mask heads if we want to
242
+ if head_mask is not None:
243
+ attention_probs_dropped = attention_probs_dropped * head_mask
244
+
245
+ context_layer = torch.matmul(attention_probs_dropped, value_layer)
246
+
247
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
248
+ new_context_layer_shape = context_layer.size()[:-2] + (
249
+ self.all_head_size, )
250
+ context_layer = context_layer.view(*new_context_layer_shape)
251
+
252
+ outputs = ((context_layer, attention_probs) if output_attentions else
253
+ (context_layer, ))
254
+
255
+ outputs = outputs + (past_key_value, )
256
+ return outputs
257
+
258
+
259
+ class BertSelfOutput(nn.Module):
260
+ def __init__(self, config):
261
+ super().__init__()
262
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
263
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
264
+ eps=config.layer_norm_eps)
265
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
266
+
267
+ def forward(self, hidden_states, input_tensor):
268
+ hidden_states = self.dense(hidden_states)
269
+ hidden_states = self.dropout(hidden_states)
270
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
271
+ return hidden_states
272
+
273
+
274
+ class BertAttention(nn.Module):
275
+ def __init__(self, config, is_cross_attention=False):
276
+ super().__init__()
277
+ self.self = BertSelfAttention(config, is_cross_attention)
278
+ self.output = BertSelfOutput(config)
279
+ self.pruned_heads = set()
280
+
281
+ def prune_heads(self, heads):
282
+ if len(heads) == 0:
283
+ return
284
+ heads, index = find_pruneable_heads_and_indices(
285
+ heads,
286
+ self.self.num_attention_heads,
287
+ self.self.attention_head_size,
288
+ self.pruned_heads,
289
+ )
290
+
291
+ # Prune linear layers
292
+ self.self.query = prune_linear_layer(self.self.query, index)
293
+ self.self.key = prune_linear_layer(self.self.key, index)
294
+ self.self.value = prune_linear_layer(self.self.value, index)
295
+ self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
296
+
297
+ # Update hyper params and store pruned heads
298
+ self.self.num_attention_heads = self.self.num_attention_heads - len(
299
+ heads)
300
+ self.self.all_head_size = (self.self.attention_head_size *
301
+ self.self.num_attention_heads)
302
+ self.pruned_heads = self.pruned_heads.union(heads)
303
+
304
+ def forward(
305
+ self,
306
+ hidden_states,
307
+ attention_mask=None,
308
+ head_mask=None,
309
+ encoder_hidden_states=None,
310
+ encoder_attention_mask=None,
311
+ past_key_value=None,
312
+ output_attentions=False,
313
+ ):
314
+ self_outputs = self.self(
315
+ hidden_states,
316
+ attention_mask,
317
+ head_mask,
318
+ encoder_hidden_states,
319
+ encoder_attention_mask,
320
+ past_key_value,
321
+ output_attentions,
322
+ )
323
+ attention_output = self.output(self_outputs[0], hidden_states)
324
+
325
+ outputs = (attention_output,
326
+ ) + self_outputs[1:] # add attentions if we output them
327
+ return outputs
328
+
329
+
330
+ class BertIntermediate(nn.Module):
331
+ def __init__(self, config):
332
+ super().__init__()
333
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
334
+ if isinstance(config.hidden_act, str):
335
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
336
+ else:
337
+ self.intermediate_act_fn = config.hidden_act
338
+
339
+ def forward(self, hidden_states):
340
+ hidden_states = self.dense(hidden_states)
341
+ hidden_states = self.intermediate_act_fn(hidden_states)
342
+ return hidden_states
343
+
344
+
345
+ class BertOutput(nn.Module):
346
+ def __init__(self, config):
347
+ super().__init__()
348
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
349
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
350
+ eps=config.layer_norm_eps)
351
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
352
+
353
+ def forward(self, hidden_states, input_tensor):
354
+ hidden_states = self.dense(hidden_states)
355
+ hidden_states = self.dropout(hidden_states)
356
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
357
+ return hidden_states
358
+
359
+
360
+ class BertLayer(nn.Module):
361
+ def __init__(self, config, layer_num):
362
+ super().__init__()
363
+ self.config = config
364
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
365
+ self.seq_len_dim = 1
366
+ self.attention = BertAttention(config)
367
+ self.layer_num = layer_num
368
+ if (self.config.add_cross_attention
369
+ and layer_num % self.config.cross_attention_freq == 0):
370
+ self.crossattention = BertAttention(
371
+ config, is_cross_attention=self.config.add_cross_attention)
372
+ self.has_cross_attention = True
373
+ else:
374
+ self.has_cross_attention = False
375
+ self.intermediate = BertIntermediate(config)
376
+ self.output = BertOutput(config)
377
+
378
+ self.intermediate_query = BertIntermediate(config)
379
+ self.output_query = BertOutput(config)
380
+
381
+ def forward(
382
+ self,
383
+ hidden_states,
384
+ attention_mask=None,
385
+ head_mask=None,
386
+ encoder_hidden_states=None,
387
+ encoder_attention_mask=None,
388
+ past_key_value=None,
389
+ output_attentions=False,
390
+ query_length=0,
391
+ ):
392
+ # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
393
+ self_attn_past_key_value = (past_key_value[:2]
394
+ if past_key_value is not None else None)
395
+ self_attention_outputs = self.attention(
396
+ hidden_states,
397
+ attention_mask,
398
+ head_mask,
399
+ output_attentions=output_attentions,
400
+ past_key_value=self_attn_past_key_value,
401
+ )
402
+ attention_output = self_attention_outputs[0]
403
+ outputs = self_attention_outputs[1:-1]
404
+
405
+ present_key_value = self_attention_outputs[-1]
406
+
407
+ if query_length > 0:
408
+ query_attention_output = attention_output[:, :query_length, :]
409
+
410
+ if self.has_cross_attention:
411
+ assert (
412
+ encoder_hidden_states is not None
413
+ ), "encoder_hidden_states must be given for cross-attention layers"
414
+ cross_attention_outputs = self.crossattention(
415
+ query_attention_output,
416
+ attention_mask,
417
+ head_mask,
418
+ encoder_hidden_states,
419
+ encoder_attention_mask,
420
+ output_attentions=output_attentions,
421
+ )
422
+ query_attention_output = cross_attention_outputs[0]
423
+ outputs = (
424
+ outputs + cross_attention_outputs[1:-1]
425
+ ) # add cross attentions if we output attention weights
426
+
427
+ layer_output = apply_chunking_to_forward(
428
+ self.feed_forward_chunk_query,
429
+ self.chunk_size_feed_forward,
430
+ self.seq_len_dim,
431
+ query_attention_output,
432
+ )
433
+ if attention_output.shape[1] > query_length:
434
+ layer_output_text = apply_chunking_to_forward(
435
+ self.feed_forward_chunk,
436
+ self.chunk_size_feed_forward,
437
+ self.seq_len_dim,
438
+ attention_output[:, query_length:, :],
439
+ )
440
+ layer_output = torch.cat([layer_output, layer_output_text],
441
+ dim=1)
442
+ else:
443
+ layer_output = apply_chunking_to_forward(
444
+ self.feed_forward_chunk,
445
+ self.chunk_size_feed_forward,
446
+ self.seq_len_dim,
447
+ attention_output,
448
+ )
449
+ outputs = (layer_output, ) + outputs
450
+
451
+ outputs = outputs + (present_key_value, )
452
+
453
+ return outputs
454
+
455
+ def feed_forward_chunk(self, attention_output):
456
+ intermediate_output = self.intermediate(attention_output)
457
+ layer_output = self.output(intermediate_output, attention_output)
458
+ return layer_output
459
+
460
+ def feed_forward_chunk_query(self, attention_output):
461
+ intermediate_output = self.intermediate_query(attention_output)
462
+ layer_output = self.output_query(intermediate_output, attention_output)
463
+ return layer_output
464
+
465
+
466
+ class BertEncoder(nn.Module):
467
+ def __init__(self, config):
468
+ super().__init__()
469
+ self.config = config
470
+ self.layer = nn.ModuleList(
471
+ [BertLayer(config, i) for i in range(config.num_hidden_layers)])
472
+
473
+ def forward(
474
+ self,
475
+ hidden_states,
476
+ attention_mask=None,
477
+ head_mask=None,
478
+ encoder_hidden_states=None,
479
+ encoder_attention_mask=None,
480
+ past_key_values=None,
481
+ use_cache=None,
482
+ output_attentions=False,
483
+ output_hidden_states=False,
484
+ return_dict=True,
485
+ query_length=0,
486
+ ):
487
+ all_hidden_states = () if output_hidden_states else None
488
+ all_self_attentions = () if output_attentions else None
489
+ all_cross_attentions = (() if output_attentions
490
+ and self.config.add_cross_attention else None)
491
+
492
+ next_decoder_cache = () if use_cache else None
493
+
494
+ for i in range(self.config.num_hidden_layers):
495
+ layer_module = self.layer[i]
496
+ if output_hidden_states:
497
+ all_hidden_states = all_hidden_states + (hidden_states, )
498
+
499
+ layer_head_mask = head_mask[i] if head_mask is not None else None
500
+ past_key_value = past_key_values[
501
+ i] if past_key_values is not None else None
502
+
503
+ if getattr(self.config, "gradient_checkpointing",
504
+ False) and self.training:
505
+
506
+ if use_cache:
507
+ logger.warn(
508
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
509
+ )
510
+ use_cache = False
511
+
512
+ def create_custom_forward(module):
513
+ def custom_forward(*inputs):
514
+ return module(*inputs, past_key_value,
515
+ output_attentions, query_length)
516
+
517
+ return custom_forward
518
+
519
+ layer_outputs = torch.utils.checkpoint.checkpoint(
520
+ create_custom_forward(layer_module),
521
+ hidden_states,
522
+ attention_mask,
523
+ layer_head_mask,
524
+ encoder_hidden_states,
525
+ encoder_attention_mask,
526
+ )
527
+ else:
528
+ layer_outputs = layer_module(
529
+ hidden_states,
530
+ attention_mask,
531
+ layer_head_mask,
532
+ encoder_hidden_states,
533
+ encoder_attention_mask,
534
+ past_key_value,
535
+ output_attentions,
536
+ query_length,
537
+ )
538
+
539
+ hidden_states = layer_outputs[0]
540
+ if use_cache:
541
+ next_decoder_cache += (layer_outputs[-1], )
542
+ if output_attentions:
543
+ all_self_attentions = all_self_attentions + (
544
+ layer_outputs[1], )
545
+ all_cross_attentions = all_cross_attentions + (
546
+ layer_outputs[2], )
547
+
548
+ if output_hidden_states:
549
+ all_hidden_states = all_hidden_states + (hidden_states, )
550
+
551
+ if not return_dict:
552
+ return tuple(v for v in [
553
+ hidden_states,
554
+ next_decoder_cache,
555
+ all_hidden_states,
556
+ all_self_attentions,
557
+ all_cross_attentions,
558
+ ] if v is not None)
559
+ return BaseModelOutputWithPastAndCrossAttentions(
560
+ last_hidden_state=hidden_states,
561
+ past_key_values=next_decoder_cache,
562
+ hidden_states=all_hidden_states,
563
+ attentions=all_self_attentions,
564
+ cross_attentions=all_cross_attentions,
565
+ )
566
+
567
+
568
+ class BertPooler(nn.Module):
569
+ def __init__(self, config):
570
+ super().__init__()
571
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
572
+ self.activation = nn.Tanh()
573
+
574
+ def forward(self, hidden_states):
575
+ # We "pool" the model by simply taking the hidden state corresponding
576
+ # to the first token.
577
+ first_token_tensor = hidden_states[:, 0]
578
+ pooled_output = self.dense(first_token_tensor)
579
+ pooled_output = self.activation(pooled_output)
580
+ return pooled_output
581
+
582
+
583
+ class BertPredictionHeadTransform(nn.Module):
584
+ def __init__(self, config):
585
+ super().__init__()
586
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
587
+ if isinstance(config.hidden_act, str):
588
+ self.transform_act_fn = ACT2FN[config.hidden_act]
589
+ else:
590
+ self.transform_act_fn = config.hidden_act
591
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
592
+ eps=config.layer_norm_eps)
593
+
594
+ def forward(self, hidden_states):
595
+ hidden_states = self.dense(hidden_states)
596
+ hidden_states = self.transform_act_fn(hidden_states)
597
+ hidden_states = self.LayerNorm(hidden_states)
598
+ return hidden_states
599
+
600
+
601
+ class BertLMPredictionHead(nn.Module):
602
+ def __init__(self, config):
603
+ super().__init__()
604
+ self.transform = BertPredictionHeadTransform(config)
605
+
606
+ # The output weights are the same as the input embeddings, but there is
607
+ # an output-only bias for each token.
608
+ self.decoder = nn.Linear(config.hidden_size,
609
+ config.vocab_size,
610
+ bias=False)
611
+
612
+ self.bias = nn.Parameter(torch.zeros(config.vocab_size))
613
+
614
+ # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
615
+ self.decoder.bias = self.bias
616
+
617
+ def forward(self, hidden_states):
618
+ hidden_states = self.transform(hidden_states)
619
+ hidden_states = self.decoder(hidden_states)
620
+ return hidden_states
621
+
622
+
623
+ class BertOnlyMLMHead(nn.Module):
624
+ def __init__(self, config):
625
+ super().__init__()
626
+ self.predictions = BertLMPredictionHead(config)
627
+
628
+ def forward(self, sequence_output):
629
+ prediction_scores = self.predictions(sequence_output)
630
+ return prediction_scores
631
+
632
+
633
+ class BertPreTrainedModel(PreTrainedModel):
634
+ """
635
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
636
+ models.
637
+ """
638
+
639
+ config_class = BertConfig
640
+ base_model_prefix = "bert"
641
+ _keys_to_ignore_on_load_missing = [r"position_ids"]
642
+
643
+ def _init_weights(self, module):
644
+ """Initialize the weights"""
645
+ if isinstance(module, (nn.Linear, nn.Embedding)):
646
+ # Slightly different from the TF version which uses truncated_normal for initialization
647
+ # cf https://github.com/pytorch/pytorch/pull/5617
648
+ module.weight.data.normal_(mean=0.0,
649
+ std=self.config.initializer_range)
650
+ elif isinstance(module, nn.LayerNorm):
651
+ module.bias.data.zero_()
652
+ module.weight.data.fill_(1.0)
653
+ if isinstance(module, nn.Linear) and module.bias is not None:
654
+ module.bias.data.zero_()
655
+
656
+
657
+ class BertModel(BertPreTrainedModel):
658
+ """
659
+ The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
660
+ cross-attention is added between the self-attention layers, following the architecture described in `Attention is
661
+ all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
662
+ Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
663
+ argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
664
+ input to the forward pass.
665
+ """
666
+ def __init__(self, config, add_pooling_layer=False):
667
+ super().__init__(config)
668
+ self.config = config
669
+
670
+ self.embeddings = BertEmbeddings(config)
671
+
672
+ self.encoder = BertEncoder(config)
673
+
674
+ self.pooler = BertPooler(config) if add_pooling_layer else None
675
+
676
+ self.init_weights()
677
+
678
+ def get_input_embeddings(self):
679
+ return self.embeddings.word_embeddings
680
+
681
+ def set_input_embeddings(self, value):
682
+ self.embeddings.word_embeddings = value
683
+
684
+ def _prune_heads(self, heads_to_prune):
685
+ """
686
+ Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
687
+ class PreTrainedModel
688
+ """
689
+ for layer, heads in heads_to_prune.items():
690
+ self.encoder.layer[layer].attention.prune_heads(heads)
691
+
692
+ def get_extended_attention_mask(
693
+ self,
694
+ attention_mask: Tensor,
695
+ input_shape: Tuple[int],
696
+ device: device,
697
+ is_decoder: bool,
698
+ has_query: bool = False,
699
+ ) -> Tensor:
700
+ """
701
+ Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
702
+
703
+ Arguments:
704
+ attention_mask (:obj:`torch.Tensor`):
705
+ Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
706
+ input_shape (:obj:`Tuple[int]`):
707
+ The shape of the input to the model.
708
+ device: (:obj:`torch.device`):
709
+ The device of the input to the model.
710
+
711
+ Returns:
712
+ :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.
713
+ """
714
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
715
+ # ourselves in which case we just need to make it broadcastable to all heads.
716
+ if attention_mask.dim() == 3:
717
+ extended_attention_mask = attention_mask[:, None, :, :]
718
+ elif attention_mask.dim() == 2:
719
+ # Provided a padding mask of dimensions [batch_size, seq_length]
720
+ # - if the model is a decoder, apply a causal mask in addition to the padding mask
721
+ # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
722
+ if is_decoder:
723
+ batch_size, seq_length = input_shape
724
+
725
+ seq_ids = torch.arange(seq_length, device=device)
726
+ causal_mask = (seq_ids[None, None, :].repeat(
727
+ batch_size, seq_length, 1) <= seq_ids[None, :, None])
728
+
729
+ # add a prefix ones mask to the causal mask
730
+ # causal and attention masks must have same type with pytorch version < 1.3
731
+ causal_mask = causal_mask.to(attention_mask.dtype)
732
+
733
+ if causal_mask.shape[1] < attention_mask.shape[1]:
734
+ prefix_seq_len = attention_mask.shape[
735
+ 1] - causal_mask.shape[1]
736
+ if has_query: # UniLM style attention mask
737
+ causal_mask = torch.cat(
738
+ [
739
+ torch.zeros(
740
+ (batch_size, prefix_seq_len, seq_length),
741
+ device=device,
742
+ dtype=causal_mask.dtype,
743
+ ),
744
+ causal_mask,
745
+ ],
746
+ axis=1,
747
+ )
748
+ causal_mask = torch.cat(
749
+ [
750
+ torch.ones(
751
+ (batch_size, causal_mask.shape[1],
752
+ prefix_seq_len),
753
+ device=device,
754
+ dtype=causal_mask.dtype,
755
+ ),
756
+ causal_mask,
757
+ ],
758
+ axis=-1,
759
+ )
760
+ extended_attention_mask = (causal_mask[:, None, :, :] *
761
+ attention_mask[:, None, None, :])
762
+ else:
763
+ extended_attention_mask = attention_mask[:, None, None, :]
764
+ else:
765
+ raise ValueError(
766
+ "Wrong shape for input_ids (shape {}) or attention_mask (shape {})"
767
+ .format(input_shape, attention_mask.shape))
768
+
769
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
770
+ # masked positions, this operation will create a tensor which is 0.0 for
771
+ # positions we want to attend and -10000.0 for masked positions.
772
+ # Since we are adding it to the raw scores before the softmax, this is
773
+ # effectively the same as removing these entirely.
774
+ extended_attention_mask = extended_attention_mask.to(
775
+ dtype=self.dtype) # fp16 compatibility
776
+ extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
777
+ return extended_attention_mask
778
+
779
+ def forward(
780
+ self,
781
+ input_ids=None,
782
+ attention_mask=None,
783
+ position_ids=None,
784
+ head_mask=None,
785
+ query_embeds=None,
786
+ encoder_hidden_states=None,
787
+ encoder_attention_mask=None,
788
+ past_key_values=None,
789
+ use_cache=None,
790
+ output_attentions=None,
791
+ output_hidden_states=None,
792
+ return_dict=None,
793
+ is_decoder=False,
794
+ ):
795
+ r"""
796
+ encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
797
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
798
+ the model is configured as a decoder.
799
+ encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
800
+ Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
801
+ the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
802
+ - 1 for tokens that are **not masked**,
803
+ - 0 for tokens that are **masked**.
804
+ past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
805
+ Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
806
+ If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
807
+ (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
808
+ instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
809
+ use_cache (:obj:`bool`, `optional`):
810
+ If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
811
+ decoding (see :obj:`past_key_values`).
812
+ """
813
+ output_attentions = (output_attentions if output_attentions is not None
814
+ else self.config.output_attentions)
815
+ output_hidden_states = (output_hidden_states
816
+ if output_hidden_states is not None else
817
+ self.config.output_hidden_states)
818
+ return_dict = (return_dict if return_dict is not None else
819
+ self.config.use_return_dict)
820
+
821
+ # use_cache = use_cache if use_cache is not None else self.config.use_cache
822
+
823
+ if input_ids is None:
824
+ assert (
825
+ query_embeds is not None
826
+ ), "You have to specify query_embeds when input_ids is None"
827
+
828
+ # past_key_values_length
829
+ past_key_values_length = (past_key_values[0][0].shape[2] -
830
+ self.config.query_length
831
+ if past_key_values is not None else 0)
832
+
833
+ query_length = query_embeds.shape[1] if query_embeds is not None else 0
834
+
835
+ embedding_output = self.embeddings(
836
+ input_ids=input_ids,
837
+ position_ids=position_ids,
838
+ query_embeds=query_embeds,
839
+ past_key_values_length=past_key_values_length,
840
+ )
841
+
842
+ input_shape = embedding_output.size()[:-1]
843
+ batch_size, seq_length = input_shape
844
+ device = embedding_output.device
845
+
846
+ if attention_mask is None:
847
+ attention_mask = torch.ones(
848
+ ((batch_size, seq_length + past_key_values_length)),
849
+ device=device)
850
+
851
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
852
+ # ourselves in which case we just need to make it broadcastable to all heads.
853
+ if is_decoder:
854
+ extended_attention_mask = self.get_extended_attention_mask(
855
+ attention_mask,
856
+ input_ids.shape,
857
+ device,
858
+ is_decoder,
859
+ has_query=(query_embeds is not None),
860
+ )
861
+ else:
862
+ extended_attention_mask = self.get_extended_attention_mask(
863
+ attention_mask, input_shape, device, is_decoder)
864
+
865
+ # If a 2D or 3D attention mask is provided for the cross-attention
866
+ # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
867
+ if encoder_hidden_states is not None:
868
+ if type(encoder_hidden_states) == list:
869
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[
870
+ 0].size()
871
+ else:
872
+ (
873
+ encoder_batch_size,
874
+ encoder_sequence_length,
875
+ _,
876
+ ) = encoder_hidden_states.size()
877
+ encoder_hidden_shape = (encoder_batch_size,
878
+ encoder_sequence_length)
879
+
880
+ if type(encoder_attention_mask) == list:
881
+ encoder_extended_attention_mask = [
882
+ self.invert_attention_mask(mask)
883
+ for mask in encoder_attention_mask
884
+ ]
885
+ elif encoder_attention_mask is None:
886
+ encoder_attention_mask = torch.ones(encoder_hidden_shape,
887
+ device=device)
888
+ encoder_extended_attention_mask = self.invert_attention_mask(
889
+ encoder_attention_mask)
890
+ else:
891
+ encoder_extended_attention_mask = self.invert_attention_mask(
892
+ encoder_attention_mask)
893
+ else:
894
+ encoder_extended_attention_mask = None
895
+
896
+ # Prepare head mask if needed
897
+ # 1.0 in head_mask indicate we keep the head
898
+ # attention_probs has shape bsz x n_heads x N x N
899
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
900
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
901
+ head_mask = self.get_head_mask(head_mask,
902
+ self.config.num_hidden_layers)
903
+
904
+ encoder_outputs = self.encoder(
905
+ embedding_output,
906
+ attention_mask=extended_attention_mask,
907
+ head_mask=head_mask,
908
+ encoder_hidden_states=encoder_hidden_states,
909
+ encoder_attention_mask=encoder_extended_attention_mask,
910
+ past_key_values=past_key_values,
911
+ use_cache=use_cache,
912
+ output_attentions=output_attentions,
913
+ output_hidden_states=output_hidden_states,
914
+ return_dict=return_dict,
915
+ query_length=query_length,
916
+ )
917
+ sequence_output = encoder_outputs[0]
918
+ pooled_output = (self.pooler(sequence_output)
919
+ if self.pooler is not None else None)
920
+
921
+ if not return_dict:
922
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
923
+
924
+ return BaseModelOutputWithPoolingAndCrossAttentions(
925
+ last_hidden_state=sequence_output,
926
+ pooler_output=pooled_output,
927
+ past_key_values=encoder_outputs.past_key_values,
928
+ hidden_states=encoder_outputs.hidden_states,
929
+ attentions=encoder_outputs.attentions,
930
+ cross_attentions=encoder_outputs.cross_attentions,
931
+ )
932
+
933
+
934
+ class BertLMHeadModel(BertPreTrainedModel):
935
+
936
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
937
+ _keys_to_ignore_on_load_missing = [
938
+ r"position_ids", r"predictions.decoder.bias"
939
+ ]
940
+
941
+ def __init__(self, config):
942
+ super().__init__(config)
943
+
944
+ self.bert = BertModel(config, add_pooling_layer=False)
945
+ self.cls = BertOnlyMLMHead(config)
946
+
947
+ self.init_weights()
948
+
949
+ def get_output_embeddings(self):
950
+ return self.cls.predictions.decoder
951
+
952
+ def set_output_embeddings(self, new_embeddings):
953
+ self.cls.predictions.decoder = new_embeddings
954
+
955
+ def forward(
956
+ self,
957
+ input_ids=None,
958
+ attention_mask=None,
959
+ position_ids=None,
960
+ head_mask=None,
961
+ query_embeds=None,
962
+ encoder_hidden_states=None,
963
+ encoder_attention_mask=None,
964
+ labels=None,
965
+ past_key_values=None,
966
+ use_cache=True,
967
+ output_attentions=None,
968
+ output_hidden_states=None,
969
+ return_dict=None,
970
+ return_logits=False,
971
+ is_decoder=True,
972
+ reduction="mean",
973
+ ):
974
+ r"""
975
+ encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
976
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
977
+ the model is configured as a decoder.
978
+ encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
979
+ Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
980
+ the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
981
+ - 1 for tokens that are **not masked**,
982
+ - 0 for tokens that are **masked**.
983
+ labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
984
+ Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
985
+ ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are
986
+ ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``
987
+ past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
988
+ Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
989
+ If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
990
+ (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
991
+ instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
992
+ use_cache (:obj:`bool`, `optional`):
993
+ If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
994
+ decoding (see :obj:`past_key_values`).
995
+ Returns:
996
+ Example::
997
+ >>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig
998
+ >>> import torch
999
+ >>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
1000
+ >>> config = BertConfig.from_pretrained("bert-base-cased")
1001
+ >>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)
1002
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
1003
+ >>> outputs = model(**inputs)
1004
+ >>> prediction_logits = outputs.logits
1005
+ """
1006
+ return_dict = (return_dict if return_dict is not None else
1007
+ self.config.use_return_dict)
1008
+ if labels is not None:
1009
+ use_cache = False
1010
+ if past_key_values is not None:
1011
+ query_embeds = None
1012
+
1013
+ outputs = self.bert(
1014
+ input_ids,
1015
+ attention_mask=attention_mask,
1016
+ position_ids=position_ids,
1017
+ head_mask=head_mask,
1018
+ query_embeds=query_embeds,
1019
+ encoder_hidden_states=encoder_hidden_states,
1020
+ encoder_attention_mask=encoder_attention_mask,
1021
+ past_key_values=past_key_values,
1022
+ use_cache=use_cache,
1023
+ output_attentions=output_attentions,
1024
+ output_hidden_states=output_hidden_states,
1025
+ return_dict=return_dict,
1026
+ is_decoder=is_decoder,
1027
+ )
1028
+
1029
+ sequence_output = outputs[0]
1030
+ if query_embeds is not None:
1031
+ sequence_output = outputs[0][:, query_embeds.shape[1]:, :]
1032
+
1033
+ prediction_scores = self.cls(sequence_output)
1034
+
1035
+ if return_logits:
1036
+ return prediction_scores[:, :-1, :].contiguous()
1037
+
1038
+ lm_loss = None
1039
+ if labels is not None:
1040
+ # we are doing next-token prediction; shift prediction scores and input ids by one
1041
+ shifted_prediction_scores = prediction_scores[:, :
1042
+ -1, :].contiguous()
1043
+ labels = labels[:, 1:].contiguous()
1044
+ loss_fct = CrossEntropyLoss(reduction=reduction,
1045
+ label_smoothing=0.1)
1046
+ lm_loss = loss_fct(
1047
+ shifted_prediction_scores.view(-1, self.config.vocab_size),
1048
+ labels.view(-1),
1049
+ )
1050
+ if reduction == "none":
1051
+ lm_loss = lm_loss.view(prediction_scores.size(0), -1).sum(1)
1052
+
1053
+ if not return_dict:
1054
+ output = (prediction_scores, ) + outputs[2:]
1055
+ return ((lm_loss, ) + output) if lm_loss is not None else output
1056
+
1057
+ return CausalLMOutputWithCrossAttentions(
1058
+ loss=lm_loss,
1059
+ logits=prediction_scores,
1060
+ past_key_values=outputs.past_key_values,
1061
+ hidden_states=outputs.hidden_states,
1062
+ attentions=outputs.attentions,
1063
+ cross_attentions=outputs.cross_attentions,
1064
+ )
1065
+
1066
+ def prepare_inputs_for_generation(self,
1067
+ input_ids,
1068
+ query_embeds,
1069
+ past=None,
1070
+ attention_mask=None,
1071
+ **model_kwargs):
1072
+ # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
1073
+ if attention_mask is None:
1074
+ attention_mask = input_ids.new_ones(input_ids.shape)
1075
+ query_mask = input_ids.new_ones(query_embeds.shape[:-1])
1076
+ attention_mask = torch.cat([query_mask, attention_mask], dim=-1)
1077
+
1078
+ # cut decoder_input_ids if past is used
1079
+ if past is not None:
1080
+ input_ids = input_ids[:, -1:]
1081
+
1082
+ return {
1083
+ "input_ids":
1084
+ input_ids,
1085
+ "query_embeds":
1086
+ query_embeds,
1087
+ "attention_mask":
1088
+ attention_mask,
1089
+ "past_key_values":
1090
+ past,
1091
+ "encoder_hidden_states":
1092
+ model_kwargs.get("encoder_hidden_states", None),
1093
+ "encoder_attention_mask":
1094
+ model_kwargs.get("encoder_attention_mask", None),
1095
+ "is_decoder":
1096
+ True,
1097
+ }
1098
+
1099
+ def _reorder_cache(self, past, beam_idx):
1100
+ reordered_past = ()
1101
+ for layer_past in past:
1102
+ reordered_past += (tuple(
1103
+ past_state.index_select(0, beam_idx)
1104
+ for past_state in layer_past), )
1105
+ return reordered_past
1106
+
1107
+
1108
+ class BertForMaskedLM(BertPreTrainedModel):
1109
+
1110
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1111
+ _keys_to_ignore_on_load_missing = [
1112
+ r"position_ids", r"predictions.decoder.bias"
1113
+ ]
1114
+
1115
+ def __init__(self, config):
1116
+ super().__init__(config)
1117
+
1118
+ self.bert = BertModel(config, add_pooling_layer=False)
1119
+ self.cls = BertOnlyMLMHead(config)
1120
+
1121
+ self.init_weights()
1122
+
1123
+ def get_output_embeddings(self):
1124
+ return self.cls.predictions.decoder
1125
+
1126
+ def set_output_embeddings(self, new_embeddings):
1127
+ self.cls.predictions.decoder = new_embeddings
1128
+
1129
+ def forward(
1130
+ self,
1131
+ input_ids=None,
1132
+ attention_mask=None,
1133
+ position_ids=None,
1134
+ head_mask=None,
1135
+ query_embeds=None,
1136
+ encoder_hidden_states=None,
1137
+ encoder_attention_mask=None,
1138
+ labels=None,
1139
+ output_attentions=None,
1140
+ output_hidden_states=None,
1141
+ return_dict=None,
1142
+ return_logits=False,
1143
+ is_decoder=False,
1144
+ ):
1145
+ r"""
1146
+ labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
1147
+ Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ...,
1148
+ config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored
1149
+ (masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``
1150
+ """
1151
+
1152
+ return_dict = (return_dict if return_dict is not None else
1153
+ self.config.use_return_dict)
1154
+
1155
+ outputs = self.bert(
1156
+ input_ids,
1157
+ attention_mask=attention_mask,
1158
+ position_ids=position_ids,
1159
+ head_mask=head_mask,
1160
+ query_embeds=query_embeds,
1161
+ encoder_hidden_states=encoder_hidden_states,
1162
+ encoder_attention_mask=encoder_attention_mask,
1163
+ output_attentions=output_attentions,
1164
+ output_hidden_states=output_hidden_states,
1165
+ return_dict=return_dict,
1166
+ is_decoder=is_decoder,
1167
+ )
1168
+
1169
+ if query_embeds is not None:
1170
+ sequence_output = outputs[0][:, query_embeds.shape[1]:, :]
1171
+ prediction_scores = self.cls(sequence_output)
1172
+
1173
+ if return_logits:
1174
+ return prediction_scores
1175
+
1176
+ masked_lm_loss = None
1177
+ if labels is not None:
1178
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
1179
+ masked_lm_loss = loss_fct(
1180
+ prediction_scores.view(-1, self.config.vocab_size),
1181
+ labels.view(-1))
1182
+
1183
+ if not return_dict:
1184
+ output = (prediction_scores, ) + outputs[2:]
1185
+ return (((masked_lm_loss, ) +
1186
+ output) if masked_lm_loss is not None else output)
1187
+
1188
+ return MaskedLMOutput(
1189
+ loss=masked_lm_loss,
1190
+ logits=prediction_scores,
1191
+ hidden_states=outputs.hidden_states,
1192
+ attentions=outputs.attentions,
1193
+ )
modeling_utils.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import math
3
+ import os
4
+ from contextlib import contextmanager
5
+
6
+ import timm.models.hub as timm_hub
7
+ import torch
8
+ import torch.distributed as dist
9
+ import torch.nn as nn
10
+
11
+
12
+ def is_dist_avail_and_initialized():
13
+ if not dist.is_available():
14
+ return False
15
+ if not dist.is_initialized():
16
+ return False
17
+ return True
18
+
19
+
20
+ def get_rank():
21
+ if not is_dist_avail_and_initialized():
22
+ return 0
23
+ return dist.get_rank()
24
+
25
+
26
+ def is_main_process():
27
+ return get_rank() == 0
28
+
29
+
30
+ def download_cached_file(url, check_hash=True, progress=False):
31
+ """
32
+ Download a file from a URL and cache it locally. If the file already exists, it is not downloaded again.
33
+ If distributed, only the main process downloads the file, and the other processes wait for the file to be downloaded.
34
+ """
35
+ def get_cached_file_path():
36
+ # a hack to sync the file path across processes
37
+ parts = torch.hub.urlparse(url)
38
+ filename = os.path.basename(parts.path)
39
+ cached_file = os.path.join(timm_hub.get_cache_dir(), filename)
40
+
41
+ return cached_file
42
+
43
+ if is_main_process():
44
+ timm_hub.download_cached_file(url, check_hash, progress)
45
+
46
+ if is_dist_avail_and_initialized():
47
+ dist.barrier()
48
+
49
+ return get_cached_file_path()
50
+
51
+
52
+ @contextmanager
53
+ def all_logging_disabled(highest_level=logging.CRITICAL):
54
+ """
55
+ A context manager that will prevent any logging messages
56
+ triggered during the body from being processed.
57
+ :param highest_level: the maximum logging level in use.
58
+ This would only need to be changed if a custom level greater than CRITICAL
59
+ is defined.
60
+ """
61
+ # two kind-of hacks here:
62
+ # * can't get the highest logging level in effect => delegate to the user
63
+ # * can't get the current module-level override => use an undocumented
64
+ # (but non-private!) interface
65
+
66
+ previous_level = logging.root.manager.disable
67
+
68
+ logging.disable(highest_level)
69
+
70
+ try:
71
+ yield
72
+ finally:
73
+ logging.disable(previous_level)
74
+
75
+
76
+ class LoRALinear(nn.Linear):
77
+ def __init__(self,
78
+ in_features: int,
79
+ out_features: int,
80
+ bias: bool = True,
81
+ device=None,
82
+ dtype=None,
83
+ lora_r=8,
84
+ lora_alpha=16,
85
+ lora_dropout=0.05,
86
+ **kwargs) -> None:
87
+ super().__init__(in_features, out_features, bias, device, dtype)
88
+ self.lora_r = lora_r
89
+ self.lora_alpha = lora_alpha
90
+ if lora_dropout > 0.:
91
+ self.lora_dropout = nn.Dropout(p=lora_dropout)
92
+ else:
93
+ self.lora_dropout = lambda x: x
94
+ self.lora_scaling = self.lora_alpha / self.lora_r
95
+
96
+ self.lora_A = nn.Linear(in_features,
97
+ self.lora_r,
98
+ bias=False,
99
+ device=device,
100
+ dtype=dtype)
101
+ self.lora_B = nn.Linear(self.lora_r,
102
+ out_features,
103
+ bias=False,
104
+ device=device,
105
+ dtype=dtype)
106
+
107
+ self.reset_parameters()
108
+
109
+ def reset_parameters(self):
110
+ if hasattr(self, 'lora_A'):
111
+ # initialize A the same way as the default for nn.Linear and B to zero
112
+ nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
113
+ nn.init.zeros_(self.lora_B.weight)
114
+ #print ("lora weight init {} {}".format(torch.mean(self.lora_A.weight), torch.mean(self.lora_B.weight)))
115
+
116
+ def forward(self, x):
117
+ orig_type = x.dtype
118
+ res = super().forward(x)
119
+ x = x.float()
120
+ res += self.lora_B(self.lora_A(
121
+ self.lora_dropout(x))) * self.lora_scaling
122
+ return res.to(orig_type)
modeling_vit.py ADDED
@@ -0,0 +1,535 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from functools import partial
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+ import torch.utils.checkpoint as checkpoint
8
+ from timm.models.layers import drop_path, to_2tuple, trunc_normal_
9
+
10
+ from .modeling_utils import download_cached_file
11
+
12
+
13
+ def _cfg(url='', **kwargs):
14
+ return {
15
+ 'url': url,
16
+ 'num_classes': 1000,
17
+ 'input_size': (3, 224, 224),
18
+ 'pool_size': None,
19
+ 'crop_pct': .9,
20
+ 'interpolation': 'bicubic',
21
+ 'mean': (0.5, 0.5, 0.5),
22
+ 'std': (0.5, 0.5, 0.5),
23
+ **kwargs
24
+ }
25
+
26
+
27
+ class DropPath(nn.Module):
28
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
29
+ """
30
+ def __init__(self, drop_prob=None):
31
+ super(DropPath, self).__init__()
32
+ self.drop_prob = drop_prob
33
+
34
+ def forward(self, x):
35
+ return drop_path(x, self.drop_prob, self.training)
36
+
37
+ def extra_repr(self) -> str:
38
+ return 'p={}'.format(self.drop_prob)
39
+
40
+
41
+ class Mlp(nn.Module):
42
+ def __init__(self,
43
+ in_features,
44
+ hidden_features=None,
45
+ out_features=None,
46
+ act_layer=nn.GELU,
47
+ drop=0.):
48
+ super().__init__()
49
+ out_features = out_features or in_features
50
+ hidden_features = hidden_features or in_features
51
+ self.fc1 = nn.Linear(in_features, hidden_features)
52
+ self.act = act_layer()
53
+ self.fc2 = nn.Linear(hidden_features, out_features)
54
+ self.drop = nn.Dropout(drop)
55
+
56
+ def forward(self, x):
57
+ x = self.fc1(x)
58
+ x = self.act(x)
59
+ # x = self.drop(x)
60
+ # commit this for the orignal BERT implement
61
+ x = self.fc2(x)
62
+ x = self.drop(x)
63
+ return x
64
+
65
+
66
+ class Attention(nn.Module):
67
+ def __init__(self,
68
+ dim,
69
+ num_heads=8,
70
+ qkv_bias=False,
71
+ qk_scale=None,
72
+ attn_drop=0.,
73
+ proj_drop=0.,
74
+ window_size=None,
75
+ attn_head_dim=None):
76
+ super().__init__()
77
+ self.num_heads = num_heads
78
+ head_dim = dim // num_heads
79
+ if attn_head_dim is not None:
80
+ head_dim = attn_head_dim
81
+ all_head_dim = head_dim * self.num_heads
82
+ self.scale = qk_scale or head_dim**-0.5
83
+
84
+ self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
85
+ if qkv_bias:
86
+ self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
87
+ self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
88
+ else:
89
+ self.q_bias = None
90
+ self.v_bias = None
91
+
92
+ if window_size:
93
+ self.window_size = window_size
94
+ self.num_relative_distance = (2 * window_size[0] -
95
+ 1) * (2 * window_size[1] - 1) + 3
96
+ self.relative_position_bias_table = nn.Parameter(
97
+ torch.zeros(self.num_relative_distance,
98
+ num_heads)) # 2*Wh-1 * 2*Ww-1, nH
99
+ # cls to token & token 2 cls & cls to cls
100
+
101
+ # get pair-wise relative position index for each token inside the window
102
+ coords_h = torch.arange(window_size[0])
103
+ coords_w = torch.arange(window_size[1])
104
+ coords = torch.stack(torch.meshgrid([coords_h,
105
+ coords_w])) # 2, Wh, Ww
106
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
107
+ relative_coords = coords_flatten[:, :,
108
+ None] - coords_flatten[:,
109
+ None, :] # 2, Wh*Ww, Wh*Ww
110
+ relative_coords = relative_coords.permute(
111
+ 1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
112
+ relative_coords[:, :,
113
+ 0] += window_size[0] - 1 # shift to start from 0
114
+ relative_coords[:, :, 1] += window_size[1] - 1
115
+ relative_coords[:, :, 0] *= 2 * window_size[1] - 1
116
+ relative_position_index = \
117
+ torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
118
+ relative_position_index[1:, 1:] = relative_coords.sum(
119
+ -1) # Wh*Ww, Wh*Ww
120
+ relative_position_index[0, 0:] = self.num_relative_distance - 3
121
+ relative_position_index[0:, 0] = self.num_relative_distance - 2
122
+ relative_position_index[0, 0] = self.num_relative_distance - 1
123
+
124
+ self.register_buffer("relative_position_index",
125
+ relative_position_index)
126
+ else:
127
+ self.window_size = None
128
+ self.relative_position_bias_table = None
129
+ self.relative_position_index = None
130
+
131
+ self.attn_drop = nn.Dropout(attn_drop)
132
+ self.proj = nn.Linear(all_head_dim, dim)
133
+ self.proj_drop = nn.Dropout(proj_drop)
134
+
135
+ def forward(self, x, rel_pos_bias=None):
136
+ B, N, C = x.shape
137
+ qkv_bias = None
138
+ if self.q_bias is not None:
139
+ qkv_bias = torch.cat(
140
+ (self.q_bias, torch.zeros_like(self.v_bias,
141
+ requires_grad=False),
142
+ self.v_bias))
143
+ # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
144
+ qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
145
+ qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
146
+ q, k, v = qkv[0], qkv[1], qkv[
147
+ 2] # make torchscript happy (cannot use tensor as tuple)
148
+
149
+ q = q * self.scale
150
+ attn = (q @ k.transpose(-2, -1))
151
+
152
+ if self.relative_position_bias_table is not None:
153
+ relative_position_bias = \
154
+ self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
155
+ self.window_size[0] * self.window_size[1] + 1,
156
+ self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH
157
+ relative_position_bias = relative_position_bias.permute(
158
+ 2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
159
+ attn = attn + relative_position_bias.unsqueeze(0)
160
+
161
+ if rel_pos_bias is not None:
162
+ attn = attn + rel_pos_bias
163
+
164
+ attn = attn.softmax(dim=-1)
165
+ attn = self.attn_drop(attn)
166
+
167
+ x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
168
+ x = self.proj(x)
169
+ x = self.proj_drop(x)
170
+ return x
171
+
172
+
173
+ class Block(nn.Module):
174
+ def __init__(self,
175
+ dim,
176
+ num_heads,
177
+ mlp_ratio=4.,
178
+ qkv_bias=False,
179
+ qk_scale=None,
180
+ drop=0.,
181
+ attn_drop=0.,
182
+ drop_path=0.,
183
+ init_values=None,
184
+ act_layer=nn.GELU,
185
+ norm_layer=nn.LayerNorm,
186
+ window_size=None,
187
+ attn_head_dim=None):
188
+ super().__init__()
189
+ self.norm1 = norm_layer(dim)
190
+ self.attn = Attention(dim,
191
+ num_heads=num_heads,
192
+ qkv_bias=qkv_bias,
193
+ qk_scale=qk_scale,
194
+ attn_drop=attn_drop,
195
+ proj_drop=drop,
196
+ window_size=window_size,
197
+ attn_head_dim=attn_head_dim)
198
+ # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
199
+ self.drop_path = DropPath(
200
+ drop_path) if drop_path > 0. else nn.Identity()
201
+ self.norm2 = norm_layer(dim)
202
+ mlp_hidden_dim = int(dim * mlp_ratio)
203
+ self.mlp = Mlp(in_features=dim,
204
+ hidden_features=mlp_hidden_dim,
205
+ act_layer=act_layer,
206
+ drop=drop)
207
+
208
+ if init_values is not None and init_values > 0:
209
+ self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),
210
+ requires_grad=True)
211
+ self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),
212
+ requires_grad=True)
213
+ else:
214
+ self.gamma_1, self.gamma_2 = None, None
215
+
216
+ def forward(self, x, rel_pos_bias=None):
217
+ if self.gamma_1 is None:
218
+ x = x + self.drop_path(
219
+ self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias))
220
+ x = x + self.drop_path(self.mlp(self.norm2(x)))
221
+ else:
222
+ x = x + self.drop_path(self.gamma_1 * self.attn(
223
+ self.norm1(x), rel_pos_bias=rel_pos_bias))
224
+ x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
225
+ return x
226
+
227
+
228
+ class PatchEmbed(nn.Module):
229
+ """ Image to Patch Embedding
230
+ """
231
+ def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
232
+ super().__init__()
233
+ img_size = to_2tuple(img_size)
234
+ patch_size = to_2tuple(patch_size)
235
+ num_patches = (img_size[1] // patch_size[1]) * (img_size[0] //
236
+ patch_size[0])
237
+ self.patch_shape = (img_size[0] // patch_size[0],
238
+ img_size[1] // patch_size[1])
239
+ self.img_size = img_size
240
+ self.patch_size = patch_size
241
+ self.num_patches = num_patches
242
+
243
+ self.proj = nn.Conv2d(in_chans,
244
+ embed_dim,
245
+ kernel_size=patch_size,
246
+ stride=patch_size)
247
+
248
+ def forward(self, x, **kwargs):
249
+ B, C, H, W = x.shape
250
+ # FIXME look at relaxing size constraints
251
+ assert H == self.img_size[0] and W == self.img_size[1], \
252
+ f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
253
+ x = self.proj(x).flatten(2).transpose(1, 2)
254
+ return x
255
+
256
+
257
+ class RelativePositionBias(nn.Module):
258
+ def __init__(self, window_size, num_heads):
259
+ super().__init__()
260
+ self.window_size = window_size
261
+ self.num_relative_distance = (2 * window_size[0] -
262
+ 1) * (2 * window_size[1] - 1) + 3
263
+ self.relative_position_bias_table = nn.Parameter(
264
+ torch.zeros(self.num_relative_distance,
265
+ num_heads)) # 2*Wh-1 * 2*Ww-1, nH
266
+ # cls to token & token 2 cls & cls to cls
267
+
268
+ # get pair-wise relative position index for each token inside the window
269
+ coords_h = torch.arange(window_size[0])
270
+ coords_w = torch.arange(window_size[1])
271
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
272
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
273
+ relative_coords = coords_flatten[:, :,
274
+ None] - coords_flatten[:,
275
+ None, :] # 2, Wh*Ww, Wh*Ww
276
+ relative_coords = relative_coords.permute(
277
+ 1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
278
+ relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0
279
+ relative_coords[:, :, 1] += window_size[1] - 1
280
+ relative_coords[:, :, 0] *= 2 * window_size[1] - 1
281
+ relative_position_index = \
282
+ torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
283
+ relative_position_index[1:,
284
+ 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
285
+ relative_position_index[0, 0:] = self.num_relative_distance - 3
286
+ relative_position_index[0:, 0] = self.num_relative_distance - 2
287
+ relative_position_index[0, 0] = self.num_relative_distance - 1
288
+
289
+ self.register_buffer("relative_position_index",
290
+ relative_position_index)
291
+
292
+ # trunc_normal_(self.relative_position_bias_table, std=.02)
293
+
294
+ def forward(self):
295
+ relative_position_bias = \
296
+ self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
297
+ self.window_size[0] * self.window_size[1] + 1,
298
+ self.window_size[0] * self.window_size[1] + 1, -1) # Wh*Ww,Wh*Ww,nH
299
+ return relative_position_bias.permute(
300
+ 2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
301
+
302
+
303
+ class VisionTransformer(nn.Module):
304
+ """ Vision Transformer with support for patch or hybrid CNN input stage
305
+ """
306
+ def __init__(self,
307
+ img_size=224,
308
+ patch_size=16,
309
+ in_chans=3,
310
+ num_classes=1000,
311
+ embed_dim=768,
312
+ depth=12,
313
+ num_heads=12,
314
+ mlp_ratio=4.,
315
+ qkv_bias=False,
316
+ qk_scale=None,
317
+ drop_rate=0.,
318
+ attn_drop_rate=0.,
319
+ drop_path_rate=0.,
320
+ norm_layer=nn.LayerNorm,
321
+ init_values=None,
322
+ use_abs_pos_emb=True,
323
+ use_rel_pos_bias=False,
324
+ use_shared_rel_pos_bias=False,
325
+ use_mean_pooling=True,
326
+ init_scale=0.001,
327
+ use_checkpoint=False):
328
+ super().__init__()
329
+ self.image_size = img_size
330
+ self.num_classes = num_classes
331
+ self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
332
+
333
+ self.patch_embed = PatchEmbed(img_size=img_size,
334
+ patch_size=patch_size,
335
+ in_chans=in_chans,
336
+ embed_dim=embed_dim)
337
+ num_patches = self.patch_embed.num_patches
338
+
339
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
340
+ if use_abs_pos_emb:
341
+ self.pos_embed = nn.Parameter(
342
+ torch.zeros(1, num_patches + 1, embed_dim))
343
+ else:
344
+ self.pos_embed = None
345
+ self.pos_drop = nn.Dropout(p=drop_rate)
346
+
347
+ if use_shared_rel_pos_bias:
348
+ self.rel_pos_bias = RelativePositionBias(
349
+ window_size=self.patch_embed.patch_shape, num_heads=num_heads)
350
+ else:
351
+ self.rel_pos_bias = None
352
+ self.use_checkpoint = use_checkpoint
353
+
354
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)
355
+ ] # stochastic depth decay rule
356
+ self.use_rel_pos_bias = use_rel_pos_bias
357
+ self.blocks = nn.ModuleList([
358
+ Block(dim=embed_dim,
359
+ num_heads=num_heads,
360
+ mlp_ratio=mlp_ratio,
361
+ qkv_bias=qkv_bias,
362
+ qk_scale=qk_scale,
363
+ drop=drop_rate,
364
+ attn_drop=attn_drop_rate,
365
+ drop_path=dpr[i],
366
+ norm_layer=norm_layer,
367
+ init_values=init_values,
368
+ window_size=self.patch_embed.patch_shape
369
+ if use_rel_pos_bias else None) for i in range(depth)
370
+ ])
371
+
372
+ if self.pos_embed is not None:
373
+ trunc_normal_(self.pos_embed, std=.02)
374
+ trunc_normal_(self.cls_token, std=.02)
375
+ self.apply(self._init_weights)
376
+ self.fix_init_weight()
377
+
378
+ def fix_init_weight(self):
379
+ def rescale(param, layer_id):
380
+ param.div_(math.sqrt(2.0 * layer_id))
381
+
382
+ for layer_id, layer in enumerate(self.blocks):
383
+ rescale(layer.attn.proj.weight.data, layer_id + 1)
384
+ rescale(layer.mlp.fc2.weight.data, layer_id + 1)
385
+
386
+ def _init_weights(self, m):
387
+ if isinstance(m, nn.Linear):
388
+ trunc_normal_(m.weight, std=.02)
389
+ if isinstance(m, nn.Linear) and m.bias is not None:
390
+ nn.init.constant_(m.bias, 0)
391
+ elif isinstance(m, nn.LayerNorm):
392
+ nn.init.constant_(m.bias, 0)
393
+ nn.init.constant_(m.weight, 1.0)
394
+
395
+ def get_classifier(self):
396
+ return self.head
397
+
398
+ def reset_classifier(self, num_classes, global_pool=''):
399
+ self.num_classes = num_classes
400
+ self.head = nn.Linear(
401
+ self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
402
+
403
+ def forward_features(self, x):
404
+ x = self.patch_embed(x)
405
+ batch_size, seq_len, _ = x.size()
406
+
407
+ cls_tokens = self.cls_token.expand(
408
+ batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
409
+ x = torch.cat((cls_tokens, x), dim=1)
410
+ if self.pos_embed is not None:
411
+ x = x + self.pos_embed
412
+ x = self.pos_drop(x)
413
+
414
+ rel_pos_bias = self.rel_pos_bias(
415
+ ) if self.rel_pos_bias is not None else None
416
+ for blk in self.blocks:
417
+ if self.use_checkpoint:
418
+ x = checkpoint.checkpoint(blk, x, rel_pos_bias)
419
+ else:
420
+ x = blk(x, rel_pos_bias)
421
+ return x
422
+
423
+ def forward(self, x):
424
+ x = self.forward_features(x)
425
+ # x = self.head(x)
426
+ return x
427
+
428
+ def get_intermediate_layers(self, x):
429
+ x = self.patch_embed(x)
430
+ batch_size, seq_len, _ = x.size()
431
+
432
+ cls_tokens = self.cls_token.expand(
433
+ batch_size, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
434
+ x = torch.cat((cls_tokens, x), dim=1)
435
+ if self.pos_embed is not None:
436
+ x = x + self.pos_embed
437
+ x = self.pos_drop(x)
438
+
439
+ features = []
440
+ rel_pos_bias = self.rel_pos_bias(
441
+ ) if self.rel_pos_bias is not None else None
442
+ for blk in self.blocks:
443
+ x = blk(x, rel_pos_bias)
444
+ features.append(x)
445
+
446
+ return features
447
+
448
+
449
+ def interpolate_pos_embed(model, checkpoint_model):
450
+ if 'pos_embed' in checkpoint_model:
451
+ pos_embed_checkpoint = checkpoint_model['pos_embed'].float()
452
+ embedding_size = pos_embed_checkpoint.shape[-1]
453
+ num_patches = model.patch_embed.num_patches
454
+ num_extra_tokens = model.pos_embed.shape[-2] - num_patches
455
+ # height (== width) for the checkpoint position embedding
456
+ orig_size = int(
457
+ (pos_embed_checkpoint.shape[-2] - num_extra_tokens)**0.5)
458
+ # height (== width) for the new position embedding
459
+ new_size = int(num_patches**0.5)
460
+ # class_token and dist_token are kept unchanged
461
+ if orig_size != new_size:
462
+ print("Position interpolate from %dx%d to %dx%d" %
463
+ (orig_size, orig_size, new_size, new_size))
464
+ extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
465
+ # only the position tokens are interpolated
466
+ pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
467
+ pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size,
468
+ embedding_size).permute(
469
+ 0, 3, 1, 2)
470
+ pos_tokens = torch.nn.functional.interpolate(pos_tokens,
471
+ size=(new_size,
472
+ new_size),
473
+ mode='bicubic',
474
+ align_corners=False)
475
+ pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
476
+ new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
477
+ checkpoint_model['pos_embed'] = new_pos_embed
478
+
479
+
480
+ def convert_weights_to_fp16(model: nn.Module):
481
+ """Convert applicable model parameters to fp16"""
482
+ def _convert_weights_to_fp16(l):
483
+ if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
484
+ l.weight.data = l.weight.data.half()
485
+ if l.bias is not None:
486
+ l.bias.data = l.bias.data.half()
487
+
488
+ model.apply(_convert_weights_to_fp16)
489
+
490
+
491
+ def convert_weights_to_fp32(model: nn.Module):
492
+ """Convert applicable model parameters to fp16"""
493
+ def _convert_weights_to_fp32(l):
494
+ if hasattr(l, 'weight') and l.weight is not None:
495
+ if l.weight.dtype == torch.float16:
496
+ l.weight = l.weight.to(torch.float32)
497
+ if hasattr(l, 'bias') and l.bias is not None:
498
+ if l.bias.dtype == torch.float16:
499
+ l.bias = l.bias.to(torch.float32)
500
+
501
+ model.apply(_convert_weights_to_fp32)
502
+
503
+
504
+ def create_eva_vit_g(img_size=224,
505
+ drop_path_rate=0.4,
506
+ use_checkpoint=False,
507
+ precision="fp16"):
508
+ model = VisionTransformer(
509
+ img_size=img_size,
510
+ patch_size=14,
511
+ use_mean_pooling=False,
512
+ embed_dim=1408,
513
+ depth=39,
514
+ num_heads=1408 // 88,
515
+ mlp_ratio=4.3637,
516
+ qkv_bias=True,
517
+ drop_path_rate=drop_path_rate,
518
+ norm_layer=partial(nn.LayerNorm, eps=1e-6),
519
+ use_checkpoint=use_checkpoint,
520
+ )
521
+ url = "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/eva_vit_g.pth"
522
+ cached_file = download_cached_file(url, check_hash=False, progress=True)
523
+ state_dict = torch.load(cached_file, map_location="cpu")
524
+ interpolate_pos_embed(model, state_dict)
525
+
526
+ incompatible_keys = model.load_state_dict(state_dict, strict=False)
527
+
528
+ if precision == "fp16":
529
+ convert_weights_to_fp16(model)
530
+
531
+ if precision == "fp32":
532
+ print('convert ViT weights to fp32')
533
+ convert_weights_to_fp32(model)
534
+
535
+ return model
pytorch_model-00001-of-00004.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:af0c2488322c736fc26124823ede6f41d474e36c0e427e31619e8125006eda2f
3
+ size 4943811896
pytorch_model-00002-of-00004.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:25164ec251335904c3a93d6f6e86108e18b2e3e5ae7e0b0097465312b4a45892
3
+ size 4977697573
pytorch_model-00003-of-00004.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c301b57f7cc16ca0e4f596722fd556087608d890b1228d9f9af632d3c1a51fd7
3
+ size 4977697701
pytorch_model-00004-of-00004.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:70f88474474beb56a61f80af62182175b6ff338e1ce85bb899ec657bb91c87c1
3
+ size 3678547170
pytorch_model.bin.index.json ADDED
The diff for this file is too large to render. See raw diff
 
special_tokens_map.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": "<s>",
3
+ "eos_token": "</s>",
4
+ "pad_token": "</s>",
5
+ "unk_token": "<unk>"
6
+ }
tokenization_InternLM_XComposer.py ADDED
@@ -0,0 +1,240 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
21
+ """Tokenization classes for IntermLM."""
22
+ import os
23
+ from shutil import copyfile
24
+ from typing import Any, Dict, List, Optional, Tuple
25
+
26
+ import sentencepiece as spm
27
+ from transformers.tokenization_utils import PreTrainedTokenizer
28
+ from transformers.utils import logging
29
+
30
+ logger = logging.get_logger(__name__)
31
+
32
+ VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"}
33
+
34
+ PRETRAINED_VOCAB_FILES_MAP = {}
35
+
36
+
37
+ class InternLMXComposerTokenizer(PreTrainedTokenizer):
38
+ """
39
+ Construct a InternLM tokenizer. Based on byte-level Byte-Pair-Encoding.
40
+
41
+ Args:
42
+ vocab_file (`str`):
43
+ Path to the vocabulary file.
44
+ """
45
+
46
+ vocab_files_names = VOCAB_FILES_NAMES
47
+ pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
48
+ model_input_names = ["input_ids", "attention_mask"]
49
+ _auto_class = "AutoTokenizer"
50
+
51
+ def __init__(
52
+ self,
53
+ vocab_file,
54
+ unk_token="<unk>",
55
+ bos_token="<s>",
56
+ eos_token="</s>",
57
+ pad_token="</s>",
58
+ sp_model_kwargs: Optional[Dict[str, Any]] = None,
59
+ add_bos_token=True,
60
+ add_eos_token=False,
61
+ decode_with_prefix_space=False,
62
+ clean_up_tokenization_spaces=False,
63
+ **kwargs,
64
+ ):
65
+ self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
66
+ super().__init__(
67
+ bos_token=bos_token,
68
+ eos_token=eos_token,
69
+ unk_token=unk_token,
70
+ pad_token=pad_token,
71
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
72
+ **kwargs,
73
+ )
74
+ self.vocab_file = vocab_file
75
+ self.add_bos_token = add_bos_token
76
+ self.add_eos_token = add_eos_token
77
+ self.decode_with_prefix_space = decode_with_prefix_space
78
+ self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
79
+ self.sp_model.Load(vocab_file)
80
+ self._no_prefix_space_tokens = None
81
+
82
+ """ Initialisation"""
83
+
84
+ @property
85
+ def no_prefix_space_tokens(self):
86
+ if self._no_prefix_space_tokens is None:
87
+ vocab = self.convert_ids_to_tokens(list(range(self.vocab_size)))
88
+ self._no_prefix_space_tokens = {i for i, tok in enumerate(vocab) if not tok.startswith("▁")}
89
+ return self._no_prefix_space_tokens
90
+
91
+ @property
92
+ def vocab_size(self):
93
+ """Returns vocab size"""
94
+ return self.sp_model.get_piece_size()
95
+
96
+ @property
97
+ def bos_token_id(self) -> Optional[int]:
98
+ return self.sp_model.bos_id()
99
+
100
+ @property
101
+ def eos_token_id(self) -> Optional[int]:
102
+ return self.sp_model.eos_id()
103
+
104
+ def get_vocab(self):
105
+ """Returns vocab as a dict"""
106
+ vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
107
+ vocab.update(self.added_tokens_encoder)
108
+ return vocab
109
+
110
+ def _tokenize(self, text):
111
+ """Returns a tokenized string."""
112
+ return self.sp_model.encode(text, out_type=str)
113
+
114
+ def _convert_token_to_id(self, token):
115
+ """Converts a token (str) in an id using the vocab."""
116
+ return self.sp_model.piece_to_id(token)
117
+
118
+ def _convert_id_to_token(self, index):
119
+ """Converts an index (integer) in a token (str) using the vocab."""
120
+ token = self.sp_model.IdToPiece(index)
121
+ return token
122
+
123
+ def _maybe_add_prefix_space(self, tokens, decoded):
124
+ if tokens and tokens[0] not in self.no_prefix_space_tokens:
125
+ return " " + decoded
126
+ else:
127
+ return decoded
128
+
129
+ def convert_tokens_to_string(self, tokens):
130
+ """Converts a sequence of tokens (string) in a single string."""
131
+ current_sub_tokens = []
132
+ out_string = ""
133
+ prev_is_special = False
134
+ for token in tokens:
135
+ # make sure that special tokens are not decoded using sentencepiece model
136
+ if token in self.all_special_tokens:
137
+ if not prev_is_special:
138
+ out_string += " "
139
+ out_string += self.sp_model.decode(current_sub_tokens) + token
140
+ prev_is_special = True
141
+ current_sub_tokens = []
142
+ else:
143
+ current_sub_tokens.append(token)
144
+ prev_is_special = False
145
+ out_string += self.sp_model.decode(current_sub_tokens)
146
+ out_string = self.clean_up_tokenization(out_string)
147
+ out_string = self._maybe_add_prefix_space(tokens=tokens, decoded=out_string)
148
+ return out_string[1:]
149
+
150
+ def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
151
+ """
152
+ Save the vocabulary and special tokens file to a directory.
153
+
154
+ Args:
155
+ save_directory (`str`):
156
+ The directory in which to save the vocabulary.
157
+
158
+ Returns:
159
+ `Tuple(str)`: Paths to the files saved.
160
+ """
161
+ if not os.path.isdir(save_directory):
162
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
163
+ return
164
+ out_vocab_file = os.path.join(
165
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
166
+ )
167
+
168
+ if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
169
+ copyfile(self.vocab_file, out_vocab_file)
170
+ elif not os.path.isfile(self.vocab_file):
171
+ with open(out_vocab_file, "wb") as fi:
172
+ content_spiece_model = self.sp_model.serialized_model_proto()
173
+ fi.write(content_spiece_model)
174
+
175
+ return (out_vocab_file,)
176
+
177
+ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
178
+ if self.add_bos_token:
179
+ bos_token_ids = [self.bos_token_id]
180
+ else:
181
+ bos_token_ids = []
182
+
183
+ output = bos_token_ids + token_ids_0
184
+
185
+ if token_ids_1 is not None:
186
+ output = output + token_ids_1
187
+
188
+ if self.add_eos_token:
189
+ output = output + [self.eos_token_id]
190
+
191
+ return output
192
+
193
+ def get_special_tokens_mask(
194
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
195
+ ) -> List[int]:
196
+ """
197
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
198
+ special tokens using the tokenizer `prepare_for_model` method.
199
+
200
+ Args:
201
+ token_ids_0 (`List[int]`):
202
+ List of IDs.
203
+ token_ids_1 (`List[int]`, *optional*):
204
+ Optional second list of IDs for sequence pairs.
205
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
206
+ Whether or not the token list is already formatted with special tokens for the model.
207
+
208
+ Returns:
209
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
210
+ """
211
+ if already_has_special_tokens:
212
+ return super().get_special_tokens_mask(
213
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
214
+ )
215
+
216
+ if token_ids_1 is None:
217
+ return [1] + ([0] * len(token_ids_0)) + [1]
218
+ return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1]
219
+
220
+ def create_token_type_ids_from_sequences(
221
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
222
+ ) -> List[int]:
223
+ """
224
+ Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make
225
+ use of token type ids, therefore a list of zeros is returned.
226
+
227
+ Args:
228
+ token_ids_0 (`List[int]`):
229
+ List of IDs.
230
+ token_ids_1 (`List[int]`, *optional*):
231
+ Optional second list of IDs for sequence pairs.
232
+
233
+ Returns:
234
+ `List[int]`: List of zeros.
235
+ """
236
+ eos = [self.eos_token_id]
237
+
238
+ if token_ids_1 is None:
239
+ return len(token_ids_0 + eos) * [0]
240
+ return len(token_ids_0 + eos + token_ids_1 + eos) * [0]
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:21ff673031fd4187f19721a86af4caa6a4deb1f3c2db284f763de3e53bd8f741
3
+ size 1658715
tokenizer_config.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoTokenizer": [
4
+ "tokenization_InternLM_XComposer.InternLMXComposerTokenizer",
5
+ null
6
+ ]
7
+ },
8
+ "bos_token": "<s>",
9
+ "clean_up_tokenization_spaces": false,
10
+ "eos_token": "</s>",
11
+ "model_max_length": 1000000000000000019884624838656,
12
+ "pad_token": "</s>",
13
+ "tokenizer_class": "InternLMXComposerTokenizer",
14
+ "unk_token": "<unk>"
15
+ }