stanrom commited on
Commit
296c480
1 Parent(s): 9373a31

Upload 9 files

Browse files
added_tokens.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "</box>": 103168,
3
+ "</ref>": 103169,
4
+ "<box>": 103170,
5
+ "<ref>": 103171
6
+ }
configuration_InternLM_XComposer.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ max_length=2048,
24
+ initializer_range=0.02,
25
+ rms_norm_eps=1e-5,
26
+ use_cache=True,
27
+ pad_token_id=-1,
28
+ bos_token_id=1,
29
+ eos_token_id=2,
30
+ tie_word_embeddings=False,
31
+ bias=True,
32
+ num_query_token=32,
33
+ num_quant=32,
34
+ intern_converted_llm=True,
35
+ kqvo_bias=True,
36
+ device='cuda',
37
+ internlm_lora=None,
38
+ # load_in_4bit=True,
39
+ **kwargs,
40
+ ):
41
+ self.vocab_size = vocab_size
42
+ self.max_length = max_length
43
+ self.max_position_embeddings = max_position_embeddings
44
+ self.hidden_size = hidden_size
45
+ self.intermediate_size = intermediate_size
46
+ self.num_hidden_layers = num_hidden_layers
47
+ self.num_attention_heads = num_attention_heads
48
+ self.hidden_act = hidden_act
49
+ self.initializer_range = initializer_range
50
+ self.rms_norm_eps = rms_norm_eps
51
+ self.use_cache = use_cache
52
+ self.bias = bias
53
+ self.num_query_token = num_query_token
54
+ self.num_quant = num_quant
55
+ self.internlm_lora = internlm_lora
56
+ self.kqvo_bias = kqvo_bias
57
+ self.intern_converted_llm = intern_converted_llm
58
+ self.device = device
59
+ super().__init__(
60
+ pad_token_id=pad_token_id,
61
+ bos_token_id=bos_token_id,
62
+ eos_token_id=eos_token_id,
63
+ tie_word_embeddings=tie_word_embeddings,
64
+ **kwargs,
65
+ )
gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
modeling_InternLM.py ADDED
@@ -0,0 +1,867 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from typing import List, Optional, Tuple, Union
3
+
4
+ import torch
5
+ import torch.utils.checkpoint
6
+ from einops import rearrange
7
+ from torch import nn
8
+ from torch.nn import CrossEntropyLoss
9
+ from transformers.activations import ACT2FN
10
+ from transformers.modeling_outputs import (BaseModelOutputWithPast,
11
+ CausalLMOutputWithPast)
12
+ from transformers.modeling_utils import PreTrainedModel
13
+ from transformers.utils import logging
14
+
15
+ from .configuration_InternLM_XComposer import InternLMXComposerConfig
16
+ from .modeling_utils import LoRALinear
17
+
18
+ logger = logging.get_logger(__name__)
19
+
20
+ _CONFIG_FOR_DOC = "InternLMXComposerConfig"
21
+
22
+
23
+ def rotary_embed(x1, x2, cos, sin, conj):
24
+ x1, x2 = x1.float(), x2.float()
25
+ if conj:
26
+ x1, x2 = x1 * cos + x2 * sin, x1 * sin + x2 * cos
27
+ else:
28
+ x1, x2 = x1 * cos - x2 * sin, x1 * sin + x2 * cos
29
+ return x1, x2
30
+
31
+
32
+ class LegacyApplyRotaryEmbQKV_(torch.autograd.Function):
33
+
34
+ @staticmethod
35
+ def forward(ctx, qkv, cos, sin, cos_k=None, sin_k=None, interleaved=False):
36
+ """
37
+ qkv: (batch_size, seqlen, 3, nheads, headdim)
38
+ cos, sin: (seqlen, rotary_dim / 2)
39
+ cos_k, sin_k: (seqlen, rotary_dim / 2), optional
40
+ interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of
41
+ 1st half and 2nd half (GPT-NeoX style).
42
+ rotary_dim must be <= headdim
43
+ Apply rotary embedding *inplace* to the first rotary_dim of q and k.
44
+ """
45
+ batch, seqlen, three, nheads, headdim = qkv.shape
46
+ assert three == 3
47
+ rotary_seqlen, rotary_dim = cos.shape
48
+ rotary_dim *= 2
49
+ assert rotary_dim <= headdim
50
+ assert seqlen <= rotary_seqlen
51
+ cos_k = cos if cos_k is None else cos_k
52
+ sin_k = sin if sin_k is None else sin_k
53
+ assert sin.shape == cos_k.shape == sin_k.shape == (rotary_seqlen, rotary_dim // 2)
54
+ q_ro = qkv[:, :, 0, :, :rotary_dim]
55
+ q1, q2 = q_ro.chunk(2, dim=-1) if not interleaved else (q_ro[..., ::2], q_ro[..., 1::2])
56
+ # rotary_emb.apply_rotary(q1, q2, rearrange(cos[:seqlen], 's d -> s 1 d'),
57
+ # rearrange(sin[:seqlen], 's d -> s 1 d'), q1, q2, False)
58
+ q1, q2 = rotary_embed(q1, q2, rearrange(cos[:seqlen], 's d -> s 1 d'), rearrange(sin[:seqlen], 's d -> s 1 d'), False)
59
+ qkv[:, :, 0, :, :rotary_dim] = torch.cat([q1, q2], dim=-1)
60
+ k_ro = qkv[:, :, 1, :, :rotary_dim]
61
+ k1, k2 = k_ro.chunk(2, dim=-1) if not interleaved else (k_ro[..., ::2], k_ro[..., 1::2])
62
+ # rotary_emb.apply_rotary(k1, k2, rearrange(cos_k[:seqlen], 's d -> s 1 d'),
63
+ # rearrange(sin_k[:seqlen], 's d -> s 1 d'), k1, k2, False)
64
+ k1, k2 = rotary_embed(k1, k2, rearrange(cos_k[:seqlen], 's d -> s 1 d'), rearrange(sin_k[:seqlen], 's d -> s 1 d'), False)
65
+ qkv[:, :, 1, :, :rotary_dim] = torch.cat([k1, k2], dim=-1)
66
+ ctx.save_for_backward(cos, sin, cos_k, sin_k)
67
+ ctx.interleaved = interleaved
68
+ return qkv
69
+
70
+ @staticmethod
71
+ def backward(ctx, dqkv):
72
+ cos, sin, cos_k, sin_k = ctx.saved_tensors
73
+ _, seqlen, _, _, headdim = dqkv.shape
74
+ rotary_dim = cos.shape[-1]
75
+ rotary_dim *= 2
76
+ dq_ro = dqkv[:, :, 0, :, :rotary_dim]
77
+ dq1, dq2 = (dq_ro.chunk(2, dim=-1) if not ctx.interleaved
78
+ else (dq_ro[..., ::2], dq_ro[..., 1::2]))
79
+ rotary_emb.apply_rotary(dq1, dq2, rearrange(cos[:seqlen], 's d -> s 1 d'),
80
+ rearrange(sin[:seqlen], 's d -> s 1 d'), dq1, dq2, True)
81
+ dk_ro = dqkv[:, :, 1, :, :rotary_dim]
82
+ dk1, dk2 = (dk_ro.chunk(2, dim=-1) if not ctx.interleaved
83
+ else (dk_ro[..., ::2], dk_ro[..., 1::2]))
84
+ rotary_emb.apply_rotary(dk1, dk2, rearrange(cos_k[:seqlen], 's d -> s 1 d'),
85
+ rearrange(sin_k[:seqlen], 's d -> s 1 d'), dk1, dk2, True)
86
+ return dqkv, None, None, None, None, None
87
+
88
+
89
+ class ConvertedInternLMRotaryEmbedding(torch.nn.Module):
90
+ def __init__(self, dim: int, base=10000, scale_base=0, device=None):
91
+ """ """
92
+ super().__init__()
93
+ # Generate and save the inverse frequency buffer (non trainable)
94
+ inv_freq = 1.0 / (base**(
95
+ torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
96
+ self.register_buffer("inv_freq", inv_freq)
97
+ self.scale_base = scale_base
98
+ scale = ((torch.arange(0, dim, 2, device=device, dtype=torch.float32) +
99
+ 0.4 * dim) / (1.4 * dim) if scale_base > 0 else None)
100
+ self.register_buffer("scale", scale)
101
+
102
+ self._seq_len_cached = 0
103
+ self._cos_cached = None
104
+ self._sin_cached = None
105
+ self._cos_k_cached = None
106
+ self._sin_k_cached = None
107
+
108
+ def _update_cos_sin_cache(self, x, indexes):
109
+ """x: (batch, seqlen, nheads, headdim) or (batch, seqlen, 3, nheads, headdim)"""
110
+ if not isinstance(indexes, int):
111
+ seqlen = indexes.max().item() + 1
112
+ else:
113
+ seqlen = indexes + 1 # eval_forward
114
+ # Reset the tables if the sequence length has changed,
115
+ # or if we're on a new device (possibly due to tracing for instance)
116
+ if seqlen > self._seq_len_cached or self._cos_cached.device != x.device or self._cos_cached.dtype != x.dtype:
117
+ self._seq_len_cached = seqlen
118
+ t = torch.arange(seqlen,
119
+ device=x.device,
120
+ dtype=self.inv_freq.dtype)
121
+ # Don't do einsum, it converts fp32 to fp16
122
+ # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
123
+ freqs = torch.outer(t, self.inv_freq.to(device=t.device))
124
+ if self.scale is None:
125
+ self._cos_cached = torch.cos(freqs).to(x.dtype)
126
+ self._sin_cached = torch.sin(freqs).to(x.dtype)
127
+ else:
128
+ power = (torch.arange(
129
+ seqlen, dtype=self.scale.dtype, device=self.scale.device) -
130
+ seqlen // 2) / self.scale_base
131
+ scale = self.scale.to(device=power.device)**rearrange(
132
+ power, "s -> s 1")
133
+ # We want the multiplication by scale to happen in fp32
134
+ self._cos_cached = (torch.cos(freqs) * scale).to(x.dtype)
135
+ self._sin_cached = (torch.sin(freqs) * scale).to(x.dtype)
136
+ self._cos_k_cached = (torch.cos(freqs) / scale).to(x.dtype)
137
+ self._sin_k_cached = (torch.sin(freqs) / scale).to(x.dtype)
138
+
139
+ def eval_forward(self, qkv, seqlen_offset=0):
140
+ """
141
+ seqlen_offset: can be used in generation where the qkv being passed in is only the last
142
+ token in the batch.
143
+ """
144
+ self._update_cos_sin_cache(qkv, seqlen_offset + qkv.shape[1])
145
+ if self.scale is None:
146
+ return legacy_apply_rotary_embed_qkv(
147
+ qkv, self._cos_cached[seqlen_offset:],
148
+ self._sin_cached[seqlen_offset:])
149
+ else:
150
+ return legacy_apply_rotary_embed_qkv(
151
+ qkv,
152
+ self._cos_cached[seqlen_offset:],
153
+ self._sin_cached[seqlen_offset:],
154
+ self._cos_k_cached[seqlen_offset:],
155
+ self._sin_k_cached[seqlen_offset:],
156
+ )
157
+
158
+
159
+ legacy_apply_rotary_embed_qkv = LegacyApplyRotaryEmbQKV_.apply
160
+
161
+
162
+ class InternConvertedInternLMAttention(nn.Module):
163
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
164
+ def __init__(self, config: InternLMXComposerConfig):
165
+ super().__init__()
166
+ self.config = config
167
+ self.hidden_size = config.hidden_size
168
+ self.num_heads = config.num_attention_heads
169
+ self.head_dim = self.hidden_size // self.num_heads
170
+ self.max_position_embeddings = config.max_position_embeddings
171
+
172
+ if (self.head_dim * self.num_heads) != self.hidden_size:
173
+ raise ValueError(
174
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
175
+ f" and `num_heads`: {self.num_heads}).")
176
+ self.q_proj = nn.Linear(self.hidden_size,
177
+ self.num_heads * self.head_dim,
178
+ bias=config.kqvo_bias)
179
+ self.k_proj = nn.Linear(self.hidden_size,
180
+ self.num_heads * self.head_dim,
181
+ bias=config.kqvo_bias)
182
+ self.v_proj = nn.Linear(self.hidden_size,
183
+ self.num_heads * self.head_dim,
184
+ bias=config.kqvo_bias)
185
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim,
186
+ self.hidden_size,
187
+ bias=config.kqvo_bias)
188
+
189
+ self.rotary_emb = ConvertedInternLMRotaryEmbedding(self.head_dim)
190
+
191
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
192
+ return tensor.view(bsz, seq_len, self.num_heads,
193
+ self.head_dim).transpose(1, 2).contiguous()
194
+
195
+ def forward(
196
+ self,
197
+ hidden_states: torch.Tensor,
198
+ attention_mask: Optional[torch.Tensor] = None,
199
+ position_ids: Optional[torch.LongTensor] = None,
200
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
201
+ output_attentions: bool = False,
202
+ use_cache: bool = False,
203
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
204
+ Optional[Tuple[torch.Tensor]]]:
205
+ bsz, q_len, _ = hidden_states.size()
206
+
207
+ query_states = self.q_proj(hidden_states)
208
+ key_states = self.k_proj(hidden_states)
209
+ value_states = self.v_proj(hidden_states)
210
+
211
+ q = query_states
212
+ k = key_states
213
+ v = value_states
214
+
215
+ qkv = torch.cat([q, k, v], dim=2).contiguous()
216
+ qkv = qkv.view(bsz, q_len, -1)
217
+ qkv = rearrange(qkv,
218
+ "b s (three h d) -> b s three h d",
219
+ three=3,
220
+ d=self.head_dim)
221
+
222
+ if past_key_value is not None:
223
+ qkv = self.rotary_emb.eval_forward(
224
+ qkv, seqlen_offset=past_key_value[0].shape[2])
225
+ else:
226
+ qkv = self.rotary_emb.eval_forward(qkv)
227
+
228
+ query_states, key_states, value_states = qkv.unbind(2)
229
+ query_states = query_states.transpose(1, 2)
230
+ key_states = key_states.transpose(1, 2)
231
+ value_states = value_states.transpose(1, 2)
232
+
233
+ kv_seq_len = key_states.shape[-2]
234
+ if past_key_value is not None:
235
+ kv_seq_len += past_key_value[0].shape[-2]
236
+ # cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
237
+ # query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
238
+ # [bsz, nh, t, hd]
239
+
240
+ if past_key_value is not None:
241
+ # reuse k, v, self_attention
242
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
243
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
244
+
245
+ past_key_value = (key_states, value_states) if use_cache else None
246
+
247
+ attn_weights = torch.matmul(query_states, key_states.transpose(
248
+ 2, 3)) / math.sqrt(self.head_dim)
249
+
250
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
251
+ raise ValueError(
252
+ f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
253
+ f" {attn_weights.size()}")
254
+
255
+ if attention_mask is not None:
256
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
257
+ raise ValueError(
258
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
259
+ )
260
+ attn_weights = attn_weights + attention_mask
261
+ attn_weights = torch.max(
262
+ attn_weights,
263
+ torch.tensor(torch.finfo(attn_weights.dtype).min))
264
+
265
+ # upcast attention to fp32
266
+ attn_weights = nn.functional.softmax(attn_weights,
267
+ dim=-1,
268
+ dtype=torch.float32).to(
269
+ query_states.dtype)
270
+ attn_output = torch.matmul(attn_weights, value_states)
271
+
272
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
273
+ raise ValueError(
274
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
275
+ f" {attn_output.size()}")
276
+
277
+ attn_output = attn_output.transpose(1, 2)
278
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
279
+
280
+ attn_output = self.o_proj(attn_output)
281
+
282
+ if not output_attentions:
283
+ attn_weights = None
284
+
285
+ return attn_output, attn_weights, past_key_value
286
+
287
+
288
+ # Copied from transformers.models.bart.modeling_bart._make_causal_mask
289
+ def _make_causal_mask(input_ids_shape: torch.Size,
290
+ dtype: torch.dtype,
291
+ device: torch.device,
292
+ past_key_values_length: int = 0):
293
+ """
294
+ Make causal mask used for bi-directional self-attention.
295
+ """
296
+ bsz, tgt_len = input_ids_shape
297
+ mask = torch.full((tgt_len, tgt_len),
298
+ torch.tensor(torch.finfo(dtype).min, device=device),
299
+ device=device)
300
+ mask_cond = torch.arange(mask.size(-1), device=device)
301
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
302
+ mask = mask.to(dtype)
303
+
304
+ if past_key_values_length > 0:
305
+ mask = torch.cat([
306
+ torch.zeros(
307
+ tgt_len, past_key_values_length, dtype=dtype, device=device),
308
+ mask
309
+ ],
310
+ dim=-1)
311
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len,
312
+ tgt_len + past_key_values_length)
313
+
314
+
315
+ # Copied from transformers.models.bart.modeling_bart._expand_mask
316
+ def _expand_mask(mask: torch.Tensor,
317
+ dtype: torch.dtype,
318
+ tgt_len: Optional[int] = None):
319
+ """
320
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
321
+ """
322
+ bsz, src_len = mask.size()
323
+ tgt_len = tgt_len if tgt_len is not None else src_len
324
+
325
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len,
326
+ src_len).to(dtype)
327
+
328
+ inverted_mask = 1.0 - expanded_mask
329
+
330
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool),
331
+ torch.finfo(dtype).min)
332
+
333
+
334
+ class InternLMRMSNorm(nn.Module):
335
+ def __init__(self, hidden_size, eps=1e-6):
336
+ """
337
+ InternLMRMSNorm is equivalent to T5LayerNorm
338
+ """
339
+ super().__init__()
340
+ self.weight = nn.Parameter(torch.ones(hidden_size))
341
+ self.variance_epsilon = eps
342
+
343
+ def forward(self, hidden_states):
344
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1,
345
+ keepdim=True)
346
+ hidden_states = hidden_states * torch.rsqrt(variance +
347
+ self.variance_epsilon)
348
+
349
+ # convert into half-precision if necessary
350
+ if self.weight.dtype in [torch.float16, torch.bfloat16]:
351
+ hidden_states = hidden_states.to(self.weight.dtype)
352
+
353
+ return self.weight * hidden_states
354
+
355
+ def rotate_half(x):
356
+ """Rotates half the hidden dims of the input."""
357
+ x1 = x[..., :x.shape[-1] // 2]
358
+ x2 = x[..., x.shape[-1] // 2:]
359
+ return torch.cat((-x2, x1), dim=-1)
360
+
361
+
362
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
363
+ gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1]
364
+ gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3])
365
+ cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2,
366
+ gather_indices)
367
+ sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2,
368
+ gather_indices)
369
+ q_embed = (q * cos) + (rotate_half(q) * sin)
370
+ k_embed = (k * cos) + (rotate_half(k) * sin)
371
+ return q_embed, k_embed
372
+
373
+
374
+ class InternLMMLP(nn.Module):
375
+ def __init__(self, hidden_size: int, intermediate_size: int,
376
+ hidden_act: str, config: InternLMXComposerConfig):
377
+ super().__init__()
378
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
379
+ self.down_proj = nn.Linear(intermediate_size,
380
+ hidden_size,
381
+ bias=False)
382
+ self.up_proj = nn.Linear(hidden_size,
383
+ intermediate_size,
384
+ bias=False)
385
+ self.act_fn = ACT2FN[hidden_act]
386
+
387
+ def forward(self, x):
388
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
389
+
390
+ class InternLMDecoderLayer(nn.Module):
391
+ def __init__(self, config: InternLMXComposerConfig):
392
+ super().__init__()
393
+ self.hidden_size = config.hidden_size
394
+ self.self_attn = InternConvertedInternLMAttention(config=config)
395
+ self.mlp = InternLMMLP(
396
+ hidden_size=self.hidden_size,
397
+ intermediate_size=config.intermediate_size,
398
+ hidden_act=config.hidden_act,
399
+ config=config,
400
+ )
401
+ self.input_layernorm = InternLMRMSNorm(config.hidden_size,
402
+ eps=config.rms_norm_eps)
403
+ self.post_attention_layernorm = InternLMRMSNorm(
404
+ config.hidden_size, eps=config.rms_norm_eps)
405
+
406
+ def forward(
407
+ self,
408
+ hidden_states: torch.Tensor,
409
+ attention_mask: Optional[torch.Tensor] = None,
410
+ position_ids: Optional[torch.LongTensor] = None,
411
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
412
+ output_attentions: Optional[bool] = False,
413
+ use_cache: Optional[bool] = False,
414
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor,
415
+ torch.FloatTensor]]]:
416
+ """
417
+ Args:
418
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
419
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
420
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
421
+ output_attentions (`bool`, *optional*):
422
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
423
+ returned tensors for more detail.
424
+ use_cache (`bool`, *optional*):
425
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
426
+ (see `past_key_values`).
427
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
428
+ """
429
+
430
+ residual = hidden_states
431
+
432
+ hidden_states = self.input_layernorm(hidden_states)
433
+
434
+ # Self Attention
435
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
436
+ hidden_states=hidden_states,
437
+ attention_mask=attention_mask,
438
+ position_ids=position_ids,
439
+ past_key_value=past_key_value,
440
+ output_attentions=output_attentions,
441
+ use_cache=use_cache,
442
+ )
443
+ hidden_states = residual + hidden_states
444
+
445
+ # Fully Connected
446
+ residual = hidden_states
447
+ hidden_states = self.post_attention_layernorm(hidden_states)
448
+ hidden_states = self.mlp(hidden_states)
449
+ hidden_states = residual + hidden_states
450
+
451
+ outputs = (hidden_states, )
452
+
453
+ if output_attentions:
454
+ outputs += (self_attn_weights, )
455
+
456
+ if use_cache:
457
+ outputs += (present_key_value, )
458
+
459
+ return outputs
460
+
461
+
462
+ class InternLMPreTrainedModel(PreTrainedModel):
463
+ config_class = InternLMXComposerConfig
464
+ base_model_prefix = "model"
465
+ supports_gradient_checkpointing = True
466
+ _no_split_modules = ["InternLMDecoderLayer"]
467
+ _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
468
+
469
+ def _init_weights(self, module):
470
+ std = self.config.initializer_range
471
+ if isinstance(module, nn.Linear):
472
+ module.weight.data.normal_(mean=0.0, std=std)
473
+ if module.bias is not None:
474
+ module.bias.data.zero_()
475
+ elif isinstance(module, nn.Embedding):
476
+ module.weight.data.normal_(mean=0.0, std=std)
477
+ if module.padding_idx is not None:
478
+ module.weight.data[module.padding_idx].zero_()
479
+
480
+ def _set_gradient_checkpointing(self, module, value=False):
481
+ if isinstance(module, InternLMModel):
482
+ module.gradient_checkpointing = value
483
+
484
+
485
+ class InternLMModel(InternLMPreTrainedModel):
486
+ """
487
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`InternLMDecoderLayer`]
488
+ Args:
489
+ config: InternLMXComposerConfig
490
+ """
491
+ def __init__(self, config: InternLMXComposerConfig):
492
+ super().__init__(config)
493
+ self.padding_idx = config.pad_token_id
494
+ self.vocab_size = config.vocab_size
495
+
496
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size,
497
+ self.padding_idx)
498
+ self.layers = nn.ModuleList([
499
+ InternLMDecoderLayer(config)
500
+ for _ in range(config.num_hidden_layers)
501
+ ])
502
+ self.norm = InternLMRMSNorm(config.hidden_size,
503
+ eps=config.rms_norm_eps)
504
+
505
+ self.gradient_checkpointing = False
506
+ # Initialize weights and apply final processing
507
+ self.post_init()
508
+
509
+ def get_input_embeddings(self):
510
+ return self.embed_tokens
511
+
512
+ def set_input_embeddings(self, value):
513
+ self.embed_tokens = value
514
+
515
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
516
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape,
517
+ inputs_embeds, past_key_values_length):
518
+ # create causal mask
519
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
520
+ combined_attention_mask = None
521
+ if input_shape[-1] > 1:
522
+ combined_attention_mask = _make_causal_mask(
523
+ input_shape,
524
+ inputs_embeds.dtype,
525
+ device=inputs_embeds.device,
526
+ past_key_values_length=past_key_values_length,
527
+ )
528
+
529
+ if attention_mask is not None:
530
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
531
+ expanded_attn_mask = _expand_mask(attention_mask,
532
+ inputs_embeds.dtype,
533
+ tgt_len=input_shape[-1]).to(
534
+ inputs_embeds.device)
535
+ combined_attention_mask = (expanded_attn_mask
536
+ if combined_attention_mask is None else
537
+ expanded_attn_mask +
538
+ combined_attention_mask)
539
+
540
+ return combined_attention_mask
541
+
542
+ def forward(
543
+ self,
544
+ input_ids: torch.LongTensor = None,
545
+ attention_mask: Optional[torch.Tensor] = None,
546
+ position_ids: Optional[torch.LongTensor] = None,
547
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
548
+ inputs_embeds: Optional[torch.FloatTensor] = None,
549
+ query_embeds: Optional[torch.FloatTensor] = None,
550
+ use_cache: Optional[bool] = None,
551
+ output_attentions: Optional[bool] = None,
552
+ output_hidden_states: Optional[bool] = None,
553
+ return_dict: Optional[bool] = None,
554
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
555
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
556
+ output_hidden_states = (output_hidden_states
557
+ if output_hidden_states is not None else
558
+ self.config.output_hidden_states)
559
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
560
+
561
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
562
+
563
+ # retrieve input_ids and inputs_embeds
564
+ if input_ids is not None and inputs_embeds is not None:
565
+ raise ValueError(
566
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
567
+ )
568
+ elif input_ids is not None:
569
+ batch_size, seq_length = input_ids.shape
570
+ elif inputs_embeds is not None:
571
+ batch_size, seq_length, _ = inputs_embeds.shape
572
+ else:
573
+ raise ValueError(
574
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
575
+ )
576
+
577
+ if inputs_embeds is None:
578
+ input_ids[input_ids==-1] = 2
579
+ inputs_embeds = self.embed_tokens(input_ids)
580
+ if query_embeds is not None:
581
+ inputs_embeds = torch.cat([query_embeds, inputs_embeds], dim=1)
582
+ batch_size, seq_length, _ = inputs_embeds.shape
583
+
584
+ seq_length_with_past = seq_length
585
+ past_key_values_length = 0
586
+
587
+ if past_key_values is not None:
588
+ past_key_values_length = past_key_values[0][0].shape[2]
589
+ seq_length_with_past = seq_length_with_past + past_key_values_length
590
+
591
+ if position_ids is None:
592
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
593
+ position_ids = torch.arange(past_key_values_length,
594
+ seq_length + past_key_values_length,
595
+ dtype=torch.long,
596
+ device=device)
597
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
598
+ else:
599
+ position_ids = position_ids.view(-1, seq_length).long()
600
+
601
+ # embed positions
602
+ if attention_mask is None:
603
+ attention_mask = torch.ones((batch_size, seq_length_with_past),
604
+ dtype=torch.bool,
605
+ device=inputs_embeds.device)
606
+ attention_mask = self._prepare_decoder_attention_mask(
607
+ attention_mask, (batch_size, seq_length), inputs_embeds,
608
+ past_key_values_length)
609
+
610
+ hidden_states = inputs_embeds
611
+
612
+ if self.gradient_checkpointing and self.training:
613
+ if use_cache:
614
+ logger.warning_once(
615
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
616
+ )
617
+ use_cache = False
618
+
619
+ # decoder layers
620
+ all_hidden_states = () if output_hidden_states else None
621
+ all_self_attns = () if output_attentions else None
622
+ next_decoder_cache = () if use_cache else None
623
+
624
+ for idx, decoder_layer in enumerate(self.layers):
625
+ if output_hidden_states:
626
+ all_hidden_states += (hidden_states, )
627
+
628
+ past_key_value = past_key_values[
629
+ idx] if past_key_values is not None else None
630
+
631
+ if self.gradient_checkpointing and self.training:
632
+
633
+ def create_custom_forward(module):
634
+ def custom_forward(*inputs):
635
+ # None for past_key_value
636
+ return module(*inputs, output_attentions, None)
637
+
638
+ return custom_forward
639
+
640
+ layer_outputs = torch.utils.checkpoint.checkpoint(
641
+ create_custom_forward(decoder_layer),
642
+ hidden_states,
643
+ attention_mask,
644
+ position_ids,
645
+ None,
646
+ )
647
+ else:
648
+ layer_outputs = decoder_layer(
649
+ hidden_states,
650
+ attention_mask=attention_mask,
651
+ position_ids=position_ids,
652
+ past_key_value=past_key_value,
653
+ output_attentions=output_attentions,
654
+ use_cache=use_cache,
655
+ )
656
+
657
+ hidden_states = layer_outputs[0]
658
+
659
+ if use_cache:
660
+ next_decoder_cache += (
661
+ layer_outputs[2 if output_attentions else 1], )
662
+
663
+ if output_attentions:
664
+ all_self_attns += (layer_outputs[1], )
665
+
666
+ hidden_states = self.norm(hidden_states)
667
+
668
+ # add hidden states from the last decoder layer
669
+ if output_hidden_states:
670
+ all_hidden_states += (hidden_states, )
671
+
672
+ next_cache = next_decoder_cache if use_cache else None
673
+ if not return_dict:
674
+ return tuple(
675
+ v for v in
676
+ [hidden_states, next_cache, all_hidden_states, all_self_attns]
677
+ if v is not None)
678
+ return BaseModelOutputWithPast(
679
+ last_hidden_state=hidden_states,
680
+ past_key_values=next_cache,
681
+ hidden_states=all_hidden_states,
682
+ attentions=all_self_attns,
683
+ )
684
+
685
+
686
+ class InternLMForCausalLM(InternLMPreTrainedModel):
687
+
688
+ def __init__(self, config):
689
+ super().__init__(config)
690
+ # TODO: find a way to explicitly initialize InternLM
691
+
692
+ if hasattr(config, 'kqvo_bias'):
693
+ setattr(config, 'kqvo_bias', config.kqvo_bias)
694
+ else:
695
+ setattr(config, 'kqvo_bias', False)
696
+ self.model = InternLMModel(config)
697
+
698
+ self.lm_head = nn.Linear(config.hidden_size,
699
+ config.vocab_size,
700
+ bias=False)
701
+
702
+ # Initialize weights and apply final processing
703
+ self.post_init()
704
+
705
+ @classmethod
706
+ def from_pretrained(cls,
707
+ pretrained_model_name_or_path,
708
+ llm_cfg=None,
709
+ *model_args,
710
+ **kwargs):
711
+ if llm_cfg:
712
+ if 'torch_dtype' in kwargs:
713
+ llm_cfg.torch_dtype = kwargs['torch_dtype']
714
+ if 'load_in_8bit' in kwargs:
715
+ llm_cfg.load_in_8bit = kwargs['load_in_8bit']
716
+ if 'device_map' in kwargs:
717
+ llm_cfg.device_map = kwargs['device_map']
718
+ return cls._from_config(llm_cfg)
719
+ else:
720
+ return super().from_pretrained(pretrained_model_name_or_path,
721
+ *model_args, **kwargs)
722
+
723
+ def get_input_embeddings(self):
724
+ return self.model.embed_tokens
725
+
726
+ def set_input_embeddings(self, value):
727
+ self.model.embed_tokens = value
728
+
729
+ def get_output_embeddings(self):
730
+ return self.lm_head
731
+
732
+ def set_output_embeddings(self, new_embeddings):
733
+ self.lm_head = new_embeddings
734
+
735
+ def set_decoder(self, decoder):
736
+ self.model = decoder
737
+
738
+ def get_decoder(self):
739
+ return self.model
740
+
741
+ def forward(
742
+ self,
743
+ input_ids: torch.LongTensor = None,
744
+ attention_mask: Optional[torch.Tensor] = None,
745
+ position_ids: Optional[torch.LongTensor] = None,
746
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
747
+ inputs_embeds: Optional[torch.FloatTensor] = None,
748
+ query_embeds: Optional[torch.FloatTensor] = None,
749
+ labels: Optional[torch.LongTensor] = None,
750
+ use_cache: Optional[bool] = None,
751
+ output_attentions: Optional[bool] = None,
752
+ output_hidden_states: Optional[bool] = None,
753
+ return_dict: Optional[bool] = None,
754
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
755
+ r"""
756
+ Args:
757
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
758
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
759
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
760
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
761
+ Returns:
762
+ Example:
763
+ ```python
764
+ >>> from transformers import AutoTokenizer, InternLMForCausalLM
765
+ >>> model = InternLMForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
766
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
767
+ >>> prompt = "Hey, are you consciours? Can you talk to me?"
768
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
769
+ >>> # Generate
770
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
771
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
772
+ "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you."
773
+ ```"""
774
+
775
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
776
+ output_hidden_states = (output_hidden_states
777
+ if output_hidden_states is not None else
778
+ self.config.output_hidden_states)
779
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
780
+
781
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
782
+ outputs = self.model(
783
+ input_ids=input_ids,
784
+ attention_mask=attention_mask,
785
+ position_ids=position_ids,
786
+ past_key_values=past_key_values,
787
+ inputs_embeds=inputs_embeds,
788
+ query_embeds=query_embeds,
789
+ use_cache=use_cache,
790
+ output_attentions=output_attentions,
791
+ output_hidden_states=output_hidden_states,
792
+ return_dict=return_dict,
793
+ )
794
+
795
+ hidden_states = outputs[0]
796
+ logits = self.lm_head(hidden_states)
797
+
798
+ loss = None
799
+ if labels is not None:
800
+ # Shift so that tokens < n predict n
801
+ shift_logits = logits[..., :-1, :].contiguous()
802
+ shift_labels = labels[..., 1:].contiguous()
803
+ # Flatten the tokens
804
+
805
+ loss_fct = CrossEntropyLoss()
806
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
807
+ shift_labels = shift_labels.view(-1)
808
+ shift_labels = shift_labels.to(shift_logits.device)
809
+
810
+ # Enable model parallelism
811
+
812
+ loss = loss_fct(shift_logits, shift_labels)
813
+
814
+ if not return_dict:
815
+ output = (logits, ) + outputs[1:]
816
+ return (loss, ) + output if loss is not None else output
817
+
818
+ return CausalLMOutputWithPast(
819
+ loss=loss,
820
+ logits=logits,
821
+ past_key_values=outputs.past_key_values,
822
+ hidden_states=outputs.hidden_states,
823
+ attentions=outputs.attentions,
824
+ )
825
+
826
+ def prepare_inputs_for_generation(self,
827
+ input_ids,
828
+ query_embeds=None,
829
+ past_key_values=None,
830
+ attention_mask=None,
831
+ inputs_embeds=None,
832
+ **kwargs):
833
+ if past_key_values:
834
+ input_ids = input_ids[:, -1:]
835
+
836
+ position_ids = kwargs.get("position_ids", None)
837
+ if attention_mask is not None and position_ids is None:
838
+ # create position_ids on the fly for batch generation
839
+ position_ids = attention_mask.long().cumsum(-1) - 1
840
+ position_ids.masked_fill_(attention_mask == 0, 1)
841
+ if past_key_values:
842
+ position_ids = position_ids[:, -1].unsqueeze(-1)
843
+ query_embeds = None
844
+
845
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
846
+ if inputs_embeds is not None and past_key_values is None:
847
+ model_inputs = {"inputs_embeds": inputs_embeds}
848
+ else:
849
+ model_inputs = {"input_ids": input_ids}
850
+
851
+ model_inputs.update({
852
+ "position_ids": position_ids,
853
+ "query_embeds": query_embeds,
854
+ "past_key_values": past_key_values,
855
+ "use_cache": kwargs.get("use_cache"),
856
+ "attention_mask": attention_mask,
857
+ })
858
+ return model_inputs
859
+
860
+ @staticmethod
861
+ def _reorder_cache(past_key_values, beam_idx):
862
+ reordered_past = ()
863
+ for layer_past in past_key_values:
864
+ reordered_past += (tuple(
865
+ past_state.index_select(0, beam_idx.to(past_state.device))
866
+ for past_state in layer_past), )
867
+ return reordered_past
modeling_InternLM_XComposer.py ADDED
@@ -0,0 +1,220 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ import torch.nn as nn
12
+ from torch.nn import LayerNorm
13
+ from torchvision import transforms
14
+ from torchvision.transforms.functional import InterpolationMode
15
+ from PIL import Image
16
+
17
+ from .modeling_vit import *
18
+ from .modeling_InternLM import *
19
+ from .modeling_utils import *
20
+ from .resampler import create_resampler
21
+
22
+ from transformers.utils import logging
23
+ logger = logging.get_logger(__name__)
24
+
25
+
26
+ class InternLMXComposerForCausalLM(PreTrainedModel):
27
+ config_class = InternLMXComposerConfig
28
+ _auto_class = "AutoModelForCausalLM"
29
+
30
+ gen_config = dict(
31
+ num_beams=5,
32
+ do_sample=True,
33
+ min_length=1,
34
+ repetition_penalty=1.5,
35
+ length_penalty=1.0,
36
+ temperature=1.0,
37
+ max_new_tokens=500,
38
+ )
39
+
40
+ def __init__(self, config):
41
+ super().__init__(config)
42
+
43
+ self.max_length = config.max_length
44
+ print (f'Set max length to {self.max_length}')
45
+ print('Init VIT ... ', end='')
46
+ self.visual_encoder = create_eva_vit_g(img_size=448)
47
+ self.ln_vision = nn.Identity()
48
+ self.supports_gradient_checkpointing = True
49
+ print('Done')
50
+ print('Init Perceive Sampler ... ', end='')
51
+ with all_logging_disabled():
52
+ self.Qformer = create_resampler(num_query_token=256)
53
+ print('Done')
54
+
55
+ print('Init InternLM ... ', end='')
56
+ self.flag_image_start = nn.Parameter(torch.zeros([1, 1, 4096]))
57
+ self.flag_image_end = nn.Parameter(torch.zeros([1, 1, 4096]))
58
+ self.flag_image_start.requires_grad = False
59
+ self.flag_image_end.requires_grad = False
60
+
61
+
62
+ if int(torch.__version__[0]) == 1:
63
+ self.internlm_model = InternLMForCausalLM._from_config(config).to(
64
+ torch.float16)
65
+ else:
66
+ assert int(torch.__version__[0]) == 2
67
+ # speed up init llm
68
+ with torch.device('meta'):
69
+ self.internlm_model = InternLMForCausalLM._from_config(config)
70
+ # self.internlm_model.to_empty(device=config.device).to(torch.float16)
71
+ # self.internlm_model.tie_weights()
72
+ # self.internlm_model.to(config.device)
73
+
74
+ self.internlm_proj = nn.Linear(4096,
75
+ self.internlm_model.config.hidden_size)
76
+ print('Done')
77
+
78
+ self.vis_processor = transforms.Compose([
79
+ transforms.Resize((448, 448),
80
+ interpolation=InterpolationMode.BICUBIC),
81
+ transforms.ToTensor(),
82
+ transforms.Normalize((0.48145466, 0.4578275, 0.40821073),
83
+ (0.26862954, 0.26130258, 0.27577711)),
84
+ ])
85
+
86
+ self.tokenizer = None
87
+
88
+ @property
89
+ def eoh(self):
90
+ return '<TOKENS_UNUSED_0>'
91
+
92
+ @property
93
+ def eoa(self):
94
+ return '<TOKENS_UNUSED_1>'
95
+
96
+ def get_input_embeddings(self):
97
+ return self.internlm_model.get_input_embeddings()
98
+
99
+ def _set_gradient_checkpointing(self, module, value=False):
100
+ if value:
101
+ self.internlm_model.apply(
102
+ partial(self.internlm_model._set_gradient_checkpointing, value=True)
103
+ )
104
+
105
+
106
+ def encode_img(self, image):
107
+ if image is None:
108
+ return None
109
+ if isinstance(image, str):
110
+ image = Image.open(image).convert("RGB")
111
+ image = self.vis_processor(image).unsqueeze(0).to(self.device)
112
+ else:
113
+ assert isinstance(image, torch.Tensor)
114
+ device = image.device
115
+ image_embeds = self.ln_vision(
116
+ self.visual_encoder(image)).to(device)
117
+ image_atts = torch.ones(image_embeds.size()[:-1],
118
+ dtype=torch.long).to(device)
119
+ query_output = self.Qformer(image_embeds)
120
+ inputs_internlm = self.internlm_proj(query_output)
121
+
122
+ inputs_internlm = torch.cat([
123
+ self.flag_image_start.expand(inputs_internlm.shape[0], -1, -1),
124
+ inputs_internlm,
125
+ self.flag_image_end.expand(inputs_internlm.shape[0], -1, -1)
126
+ ],
127
+ dim=1)
128
+ return inputs_internlm
129
+
130
+ def encode_text(self, text, add_special_tokens=False):
131
+ text_token_ids = self.tokenizer(
132
+ text,
133
+ return_tensors='pt',
134
+ add_special_tokens=add_special_tokens,
135
+ ).input_ids.to(self.device)
136
+ text_embeds = self.internlm_model.model.embed_tokens(text_token_ids)
137
+ return text_embeds
138
+
139
+ def decode_text(self, out_embeds):
140
+ out_text = self.tokenizer.batch_decode(out_embeds,
141
+ skip_special_tokens=True)[0]
142
+ out_text = out_text.split(self.eoa)[0]
143
+ return out_text
144
+
145
+ def wrap_text(self, user_text, bot_text='', add_special=True):
146
+ if add_special:
147
+ eoh = self.eoh
148
+ else:
149
+ eoh = ''
150
+ text = f'<|User|>:{user_text}{eoh}\n<|Bot|>:{bot_text}'
151
+ return text
152
+
153
+ def get_gen_args(self, **kwargs):
154
+ new_kargs = copy.deepcopy(self.gen_config)
155
+ new_kargs.update(kwargs)
156
+ return new_kargs
157
+
158
+ def generate(self, text, image=None, **kwargs):
159
+ text_embeds = self.encode_text(text)
160
+ img_embeds = self.encode_img(image)
161
+ prompt_embeds = self.wrap_prompt(text_embeds, img_embeds)
162
+ out_embeds = self.internlm_model.generate(inputs_embeds=prompt_embeds,
163
+ **self.get_gen_args(**kwargs))
164
+ out_text = self.decode_text(out_embeds)
165
+ return out_text
166
+
167
+ def chat(self, text, image=None, history=None, **kwargs):
168
+ text_embeds = self.encode_text(text)
169
+ img_embeds = self.encode_img(image)
170
+ prompt_embeds = self.wrap_prompt(text_embeds,
171
+ img_embeds,
172
+ history=history)
173
+ out_embeds = self.internlm_model.generate(inputs_embeds=prompt_embeds,
174
+ **self.get_gen_args(**kwargs))
175
+ out_text = self.decode_text(out_embeds)
176
+
177
+ # trunc at eoh and eoa
178
+ clean_out_text_token_ids = self.tokenizer(
179
+ out_text, return_tensors='pt').input_ids.to(self.device)
180
+ clean_out_text_embeds = self.internlm_model.model.embed_tokens(
181
+ clean_out_text_token_ids)
182
+ clean_prompt_embeds = self.wrap_prompt(text_embeds,
183
+ img_embeds,
184
+ add_special=False)
185
+ cur_history = torch.cat([clean_prompt_embeds, clean_out_text_embeds],
186
+ dim=1)
187
+ if history is None:
188
+ history = []
189
+ history.append(cur_history)
190
+ return out_text, history
191
+
192
+ def wrap_prompt(self,
193
+ text_embeds,
194
+ img_embeds=None,
195
+ history=None,
196
+ add_special=True):
197
+ if add_special:
198
+ prompt_segs = ['<|User|>:', f'{self.eoh}\n<|Bot|>:']
199
+ else:
200
+ prompt_segs = ['<|User|>:', '<|Bot|>:'] # used in wrap history
201
+ prompt_seg_embeds = []
202
+ for i, seg in enumerate(prompt_segs):
203
+ if history is not None:
204
+ add_special_tokens = False
205
+ else:
206
+ add_special_tokens = i == 0
207
+ seg_embeds = self.encode_text(
208
+ seg, add_special_tokens=add_special_tokens)
209
+ prompt_seg_embeds.append(seg_embeds)
210
+ if img_embeds is None:
211
+ img_embeds = text_embeds.new_empty(text_embeds.size(0), 0,
212
+ text_embeds.size(-1))
213
+ prompt_seg_embeds = [
214
+ prompt_seg_embeds[0], img_embeds, text_embeds, prompt_seg_embeds[1]
215
+ ]
216
+ prompt_embeds = torch.cat(prompt_seg_embeds, dim=1)
217
+ if history is not None:
218
+ prompt_embeds = torch.cat([*history, prompt_embeds], dim=1)
219
+ return prompt_embeds
220
+
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
resampler.py ADDED
@@ -0,0 +1,180 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ import torch
4
+ import torch.nn as nn
5
+ from torch.nn import functional as F
6
+ from timm.models.layers import trunc_normal_
7
+ from functools import partial
8
+ import math
9
+ import numpy as np
10
+
11
+ def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
12
+ """
13
+ grid_size: int of the grid height and width
14
+ return:
15
+ pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
16
+ """
17
+ grid_h = np.arange(grid_size, dtype=np.float32)
18
+ grid_w = np.arange(grid_size, dtype=np.float32)
19
+ grid = np.meshgrid(grid_w, grid_h) # here w goes first
20
+ grid = np.stack(grid, axis=0)
21
+
22
+ grid = grid.reshape([2, 1, grid_size, grid_size])
23
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
24
+ if cls_token:
25
+ pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
26
+ return pos_embed
27
+
28
+
29
+ def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
30
+ assert embed_dim % 2 == 0
31
+
32
+ # use half of dimensions to encode grid_h
33
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
34
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
35
+
36
+ emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
37
+ return emb
38
+
39
+
40
+ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
41
+ """
42
+ embed_dim: output dimension for each position
43
+ pos: a list of positions to be encoded: size (M,)
44
+ out: (M, D)
45
+ """
46
+ assert embed_dim % 2 == 0
47
+ omega = np.arange(embed_dim // 2, dtype=np.float32)
48
+ omega /= embed_dim / 2.
49
+ omega = 1. / 10000**omega # (D/2,)
50
+
51
+ pos = pos.reshape(-1) # (M,)
52
+ out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product
53
+
54
+ emb_sin = np.sin(out) # (M, D/2)
55
+ emb_cos = np.cos(out) # (M, D/2)
56
+
57
+ emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
58
+ return emb
59
+
60
+ def interpolate_pos_embed(model, checkpoint_model):
61
+ if 'pos_embed' in checkpoint_model:
62
+ pos_embed_checkpoint = checkpoint_model['pos_embed']
63
+ embedding_size = pos_embed_checkpoint.shape[-1]
64
+ num_patches = model.patch_embed.num_patches
65
+ num_extra_tokens = model.pos_embed.shape[-2] - num_patches
66
+ # height (== width) for the checkpoint position embedding
67
+ orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
68
+ # height (== width) for the new position embedding
69
+ new_size = int(num_patches ** 0.5)
70
+ # class_token and dist_token are kept unchanged
71
+ if orig_size != new_size:
72
+ print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
73
+ extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
74
+ # only the position tokens are interpolated
75
+ pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
76
+ pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
77
+ pos_tokens = torch.nn.functional.interpolate(
78
+ pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
79
+ pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
80
+ new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
81
+ checkpoint_model['pos_embed'] = new_pos_embed
82
+ def get_abs_pos(abs_pos, tgt_size):
83
+ # abs_pos: L, C
84
+ # tgt_size: M
85
+ # return: M, C
86
+ src_size = int(math.sqrt(abs_pos.size(0)))
87
+ tgt_size = int(math.sqrt(tgt_size))
88
+ dtype = abs_pos.dtype
89
+
90
+ if src_size != tgt_size:
91
+ return F.interpolate(
92
+ abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2),
93
+ size=(tgt_size, tgt_size),
94
+ mode="bicubic",
95
+ align_corners=False,
96
+ ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype)
97
+ else:
98
+ return abs_pos
99
+
100
+ class Resampler(nn.Module):
101
+ """
102
+ A 2D perceiver-resampler network with one cross attention layers by
103
+ (grid_size**2) learnable queries and 2d sincos pos_emb
104
+ Outputs:
105
+ A tensor with the shape of (grid_size**2, embed_dim)
106
+ """
107
+
108
+ def __init__(
109
+ self,
110
+ grid_size,
111
+ embed_dim,
112
+ num_heads,
113
+ kv_dim=None,
114
+ norm_layer=partial(nn.LayerNorm, eps=1e-6)
115
+ ):
116
+ super().__init__()
117
+ self.num_queries = grid_size ** 2
118
+ self.embed_dim = embed_dim
119
+ self.num_heads = num_heads
120
+
121
+ self.pos_embed = nn.Parameter(
122
+ torch.from_numpy(get_2d_sincos_pos_embed(embed_dim, grid_size)).float()
123
+ ).requires_grad_(False)
124
+
125
+ self.query = nn.Parameter(torch.zeros(self.num_queries, embed_dim))
126
+ trunc_normal_(self.query, std=.02)
127
+
128
+ if kv_dim is not None and kv_dim != embed_dim:
129
+ self.kv_proj = nn.Linear(kv_dim, embed_dim, bias=False)
130
+ else:
131
+ self.kv_proj = nn.Identity()
132
+
133
+ self.attn = nn.MultiheadAttention(embed_dim, num_heads)
134
+ self.ln_q = norm_layer(embed_dim)
135
+ self.ln_kv = norm_layer(embed_dim)
136
+
137
+ self.ln_post = norm_layer(embed_dim)
138
+
139
+ self.apply(self._init_weights)
140
+
141
+ def _init_weights(self, m):
142
+ if isinstance(m, nn.Linear):
143
+ trunc_normal_(m.weight, std=.02)
144
+ if isinstance(m, nn.Linear) and m.bias is not None:
145
+ nn.init.constant_(m.bias, 0)
146
+ elif isinstance(m, nn.LayerNorm):
147
+ nn.init.constant_(m.bias, 0)
148
+ nn.init.constant_(m.weight, 1.0)
149
+
150
+ def forward(self, x, attn_mask=None):
151
+
152
+ pos_embed = get_abs_pos(self.pos_embed, x.size(1))
153
+
154
+ x = self.kv_proj(x)
155
+ x = self.ln_kv(x).permute(1, 0, 2)
156
+ k = x.clone()
157
+ k[1:] = x[1:] + pos_embed.unsqueeze(1)
158
+
159
+ N = x.shape[1]
160
+ q = self.ln_q(self.query)
161
+ out = self.attn(
162
+ self._repeat(q, N) + self.pos_embed.unsqueeze(1),
163
+ k,
164
+ x,
165
+ attn_mask=attn_mask)[0]
166
+ out = self.ln_post(out.permute(1, 0, 2))
167
+ return out
168
+
169
+ def _repeat(self, query, N: int):
170
+ return query.unsqueeze(1).repeat(1, N, 1)
171
+
172
+
173
+ def create_resampler(num_query_token=32, vision_width=1408,):
174
+ attn_pool = Resampler(
175
+ grid_size=int(math.sqrt(num_query_token)),
176
+ embed_dim=4096,
177
+ num_heads=4096 // 128,
178
+ kv_dim=vision_width,
179
+ )
180
+ return attn_pool
tokenization_InternLM_XComposer.py ADDED
@@ -0,0 +1,241 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ self.vocab_file = vocab_file
67
+ self.add_bos_token = add_bos_token
68
+ self.add_eos_token = add_eos_token
69
+ self.decode_with_prefix_space = decode_with_prefix_space
70
+ self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
71
+ self.sp_model.Load(vocab_file)
72
+ self._no_prefix_space_tokens = None
73
+ super().__init__(
74
+ bos_token=bos_token,
75
+ eos_token=eos_token,
76
+ unk_token=unk_token,
77
+ pad_token=pad_token,
78
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
79
+ **kwargs,
80
+ )
81
+
82
+
83
+ """ Initialisation"""
84
+
85
+ @property
86
+ def no_prefix_space_tokens(self):
87
+ if self._no_prefix_space_tokens is None:
88
+ vocab = self.convert_ids_to_tokens(list(range(self.vocab_size)))
89
+ self._no_prefix_space_tokens = {i for i, tok in enumerate(vocab) if not tok.startswith("▁")}
90
+ return self._no_prefix_space_tokens
91
+
92
+ @property
93
+ def vocab_size(self):
94
+ """Returns vocab size"""
95
+ return self.sp_model.get_piece_size()
96
+
97
+ @property
98
+ def bos_token_id(self) -> Optional[int]:
99
+ return self.sp_model.bos_id()
100
+
101
+ @property
102
+ def eos_token_id(self) -> Optional[int]:
103
+ return self.sp_model.eos_id()
104
+
105
+ def get_vocab(self):
106
+ """Returns vocab as a dict"""
107
+ vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
108
+ vocab.update(self.added_tokens_encoder)
109
+ return vocab
110
+
111
+ def _tokenize(self, text):
112
+ """Returns a tokenized string."""
113
+ return self.sp_model.encode(text, out_type=str)
114
+
115
+ def _convert_token_to_id(self, token):
116
+ """Converts a token (str) in an id using the vocab."""
117
+ return self.sp_model.piece_to_id(token)
118
+
119
+ def _convert_id_to_token(self, index):
120
+ """Converts an index (integer) in a token (str) using the vocab."""
121
+ token = self.sp_model.IdToPiece(index)
122
+ return token
123
+
124
+ def _maybe_add_prefix_space(self, tokens, decoded):
125
+ if tokens and tokens[0] not in self.no_prefix_space_tokens:
126
+ return " " + decoded
127
+ else:
128
+ return decoded
129
+
130
+ def convert_tokens_to_string(self, tokens):
131
+ """Converts a sequence of tokens (string) in a single string."""
132
+ current_sub_tokens = []
133
+ out_string = ""
134
+ prev_is_special = False
135
+ for token in tokens:
136
+ # make sure that special tokens are not decoded using sentencepiece model
137
+ if token in self.all_special_tokens:
138
+ if not prev_is_special:
139
+ out_string += " "
140
+ out_string += self.sp_model.decode(current_sub_tokens) + token
141
+ prev_is_special = True
142
+ current_sub_tokens = []
143
+ else:
144
+ current_sub_tokens.append(token)
145
+ prev_is_special = False
146
+ out_string += self.sp_model.decode(current_sub_tokens)
147
+ out_string = self.clean_up_tokenization(out_string)
148
+ out_string = self._maybe_add_prefix_space(tokens=tokens, decoded=out_string)
149
+ return out_string[1:]
150
+
151
+ def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
152
+ """
153
+ Save the vocabulary and special tokens file to a directory.
154
+
155
+ Args:
156
+ save_directory (`str`):
157
+ The directory in which to save the vocabulary.
158
+
159
+ Returns:
160
+ `Tuple(str)`: Paths to the files saved.
161
+ """
162
+ if not os.path.isdir(save_directory):
163
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
164
+ return
165
+ out_vocab_file = os.path.join(
166
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
167
+ )
168
+
169
+ if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
170
+ copyfile(self.vocab_file, out_vocab_file)
171
+ elif not os.path.isfile(self.vocab_file):
172
+ with open(out_vocab_file, "wb") as fi:
173
+ content_spiece_model = self.sp_model.serialized_model_proto()
174
+ fi.write(content_spiece_model)
175
+
176
+ return (out_vocab_file,)
177
+
178
+ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
179
+ if self.add_bos_token:
180
+ bos_token_ids = [self.bos_token_id]
181
+ else:
182
+ bos_token_ids = []
183
+
184
+ output = bos_token_ids + token_ids_0
185
+
186
+ if token_ids_1 is not None:
187
+ output = output + token_ids_1
188
+
189
+ if self.add_eos_token:
190
+ output = output + [self.eos_token_id]
191
+
192
+ return output
193
+
194
+ def get_special_tokens_mask(
195
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
196
+ ) -> List[int]:
197
+ """
198
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
199
+ special tokens using the tokenizer `prepare_for_model` method.
200
+
201
+ Args:
202
+ token_ids_0 (`List[int]`):
203
+ List of IDs.
204
+ token_ids_1 (`List[int]`, *optional*):
205
+ Optional second list of IDs for sequence pairs.
206
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
207
+ Whether or not the token list is already formatted with special tokens for the model.
208
+
209
+ Returns:
210
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
211
+ """
212
+ if already_has_special_tokens:
213
+ return super().get_special_tokens_mask(
214
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
215
+ )
216
+
217
+ if token_ids_1 is None:
218
+ return [1] + ([0] * len(token_ids_0)) + [1]
219
+ return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1]
220
+
221
+ def create_token_type_ids_from_sequences(
222
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
223
+ ) -> List[int]:
224
+ """
225
+ Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make
226
+ use of token type ids, therefore a list of zeros is returned.
227
+
228
+ Args:
229
+ token_ids_0 (`List[int]`):
230
+ List of IDs.
231
+ token_ids_1 (`List[int]`, *optional*):
232
+ Optional second list of IDs for sequence pairs.
233
+
234
+ Returns:
235
+ `List[int]`: List of zeros.
236
+ """
237
+ eos = [self.eos_token_id]
238
+
239
+ if token_ids_1 is None:
240
+ return len(token_ids_0 + eos) * [0]
241
+ return len(token_ids_0 + eos + token_ids_1 + eos) * [0]