AoEiuV020 commited on
Commit
19e82e7
·
1 Parent(s): 19ba095
LMConfig.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+ from typing import List
3
+
4
+
5
+ class LMConfig(PretrainedConfig):
6
+ model_type = "minimind"
7
+
8
+ def __init__(
9
+ self,
10
+ dim: int = 512,
11
+ n_layers: int = 8,
12
+ n_heads: int = 8,
13
+ n_kv_heads: int = 2,
14
+ vocab_size: int = 6400,
15
+ hidden_dim: int = None,
16
+ multiple_of: int = 64,
17
+ norm_eps: float = 1e-5,
18
+ max_seq_len: int = 8192,
19
+ rope_theta: int = 1e6,
20
+ dropout: float = 0.0,
21
+ flash_attn: bool = True,
22
+ ####################################################
23
+ # Here are the specific configurations of MOE
24
+ # When use_moe is false, the following is invalid
25
+ ####################################################
26
+ use_moe: bool = False,
27
+ ####################################################
28
+ num_experts_per_tok: int = 2,
29
+ n_routed_experts: int = 4,
30
+ n_shared_experts: bool = True,
31
+ scoring_func: str = 'softmax',
32
+ aux_loss_alpha: float = 0.1,
33
+ seq_aux: bool = True,
34
+ norm_topk_prob: bool = True,
35
+ **kwargs,
36
+ ):
37
+ self.dim = dim
38
+ self.n_layers = n_layers
39
+ self.n_heads = n_heads
40
+ self.n_kv_heads = n_kv_heads
41
+ self.vocab_size = vocab_size
42
+ self.hidden_dim = hidden_dim
43
+ self.multiple_of = multiple_of
44
+ self.norm_eps = norm_eps
45
+ self.max_seq_len = max_seq_len
46
+ self.rope_theta = rope_theta
47
+ self.dropout = dropout
48
+ self.flash_attn = flash_attn
49
+ ####################################################
50
+ # Here are the specific configurations of MOE
51
+ # When use_moe is false, the following is invalid
52
+ ####################################################
53
+ self.use_moe = use_moe
54
+ self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量
55
+ self.n_routed_experts = n_routed_experts # 总的专家数量
56
+ self.n_shared_experts = n_shared_experts # 共享专家
57
+ self.scoring_func = scoring_func # 评分函数,默认为'softmax'
58
+ self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数
59
+ self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失
60
+ self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率
61
+ super().__init__(**kwargs)
config.json ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "MiniMindLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "LMConfig.LMConfig",
7
+ "AutoModelForCausalLM": "model.MiniMindLM"
8
+ },
9
+ "aux_loss_alpha": 0.1,
10
+ "dim": 512,
11
+ "dropout": 0.0,
12
+ "flash_attn": true,
13
+ "hidden_dim": 1408,
14
+ "max_seq_len": 512,
15
+ "model_type": "minimind",
16
+ "multiple_of": 64,
17
+ "n_heads": 8,
18
+ "n_kv_heads": 2,
19
+ "n_layers": 8,
20
+ "n_routed_experts": 4,
21
+ "n_shared_experts": true,
22
+ "norm_eps": 1e-05,
23
+ "norm_topk_prob": true,
24
+ "num_experts_per_tok": 2,
25
+ "rope_theta": 1000000.0,
26
+ "scoring_func": "softmax",
27
+ "seq_aux": true,
28
+ "torch_dtype": "float32",
29
+ "transformers_version": "4.48.0",
30
+ "use_moe": false,
31
+ "vocab_size": 6400
32
+ }
generation_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "transformers_version": "4.48.0"
4
+ }
model.py ADDED
@@ -0,0 +1,385 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import struct
3
+ import inspect
4
+ import time
5
+
6
+ from .LMConfig import LMConfig
7
+ from typing import Any, Optional, Tuple, List, Union
8
+ import numpy as np
9
+ import torch
10
+ import torch.nn.functional as F
11
+ from torch import nn
12
+ from transformers import PreTrainedModel
13
+ from transformers.modeling_outputs import CausalLMOutputWithPast
14
+
15
+
16
+ class RMSNorm(torch.nn.Module):
17
+ def __init__(self, dim: int, eps: float = 1e-6):
18
+ super().__init__()
19
+ self.eps = eps
20
+ self.weight = nn.Parameter(torch.ones(dim))
21
+
22
+ def _norm(self, x):
23
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
24
+
25
+ def forward(self, x):
26
+ return self.weight * self._norm(x.float()).type_as(x)
27
+
28
+
29
+ def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
30
+ freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
31
+ t = torch.arange(end, device=freqs.device) # type: ignore
32
+ freqs = torch.outer(t, freqs).float() # type: ignore
33
+ pos_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
34
+ return pos_cis
35
+
36
+
37
+ def apply_rotary_emb(xq, xk, pos_cis):
38
+ def unite_shape(pos_cis, x):
39
+ ndim = x.ndim
40
+ assert 0 <= 1 < ndim
41
+ assert pos_cis.shape == (x.shape[1], x.shape[-1])
42
+ shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
43
+ return pos_cis.view(*shape)
44
+
45
+ xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
46
+ xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
47
+ pos_cis = unite_shape(pos_cis, xq_)
48
+ xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
49
+ xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
50
+ return xq_out.type_as(xq), xk_out.type_as(xk)
51
+
52
+
53
+ def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
54
+ """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
55
+ bs, slen, n_kv_heads, head_dim = x.shape
56
+ if n_rep == 1:
57
+ return x
58
+ return (
59
+ x[:, :, :, None, :]
60
+ .expand(bs, slen, n_kv_heads, n_rep, head_dim)
61
+ .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
62
+ )
63
+
64
+
65
+ class Attention(nn.Module):
66
+ def __init__(self, args: LMConfig):
67
+ super().__init__()
68
+ self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
69
+ assert args.n_heads % self.n_kv_heads == 0
70
+ self.n_local_heads = args.n_heads
71
+ self.n_local_kv_heads = self.n_kv_heads
72
+ self.n_rep = self.n_local_heads // self.n_local_kv_heads
73
+ self.head_dim = args.dim // args.n_heads
74
+ self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
75
+ self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
76
+ self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
77
+ self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
78
+ self.attn_dropout = nn.Dropout(args.dropout)
79
+ self.resid_dropout = nn.Dropout(args.dropout)
80
+ self.dropout = args.dropout
81
+ self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
82
+ # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
83
+ mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
84
+ mask = torch.triu(mask, diagonal=1)
85
+ self.register_buffer("mask", mask, persistent=False)
86
+
87
+ def forward(self,
88
+ x: torch.Tensor,
89
+ pos_cis: torch.Tensor,
90
+ past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
91
+ use_cache=False):
92
+ bsz, seq_len, _ = x.shape
93
+ xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
94
+ xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
95
+ xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
96
+ xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
97
+
98
+ xq, xk = apply_rotary_emb(xq, xk, pos_cis)
99
+ # kv_cache实现
100
+ if past_key_value is not None:
101
+ xk = torch.cat([past_key_value[0], xk], dim=1)
102
+ xv = torch.cat([past_key_value[1], xv], dim=1)
103
+ past_kv = (xk, xv) if use_cache else None
104
+
105
+ xq, xk, xv = (
106
+ xq.transpose(1, 2),
107
+ repeat_kv(xk, self.n_rep).transpose(1, 2),
108
+ repeat_kv(xv, self.n_rep).transpose(1, 2)
109
+ )
110
+ if self.flash and seq_len != 1:
111
+ dropout_p = self.dropout if self.training else 0.0
112
+ output = F.scaled_dot_product_attention(
113
+ xq, xk, xv,
114
+ attn_mask=None,
115
+ dropout_p=dropout_p,
116
+ is_causal=True
117
+ )
118
+ else:
119
+ scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
120
+ scores += self.mask[:, :, :seq_len, :seq_len]
121
+ scores = F.softmax(scores.float(), dim=-1).type_as(xq)
122
+ scores = self.attn_dropout(scores)
123
+ output = scores @ xv
124
+
125
+ output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
126
+ output = self.resid_dropout(self.wo(output))
127
+ return output, past_kv
128
+
129
+
130
+ class FeedForward(nn.Module):
131
+ def __init__(self, config: LMConfig):
132
+ super().__init__()
133
+ if config.hidden_dim is None:
134
+ hidden_dim = 4 * config.dim
135
+ hidden_dim = int(2 * hidden_dim / 3)
136
+ config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
137
+ self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
138
+ self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
139
+ self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)
140
+ self.dropout = nn.Dropout(config.dropout)
141
+
142
+ def forward(self, x):
143
+ return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
144
+
145
+
146
+ class MoEGate(nn.Module):
147
+ def __init__(self, config: LMConfig):
148
+ super().__init__()
149
+ self.config = config
150
+ self.top_k = config.num_experts_per_tok
151
+ self.n_routed_experts = config.n_routed_experts
152
+
153
+ self.scoring_func = config.scoring_func
154
+ self.alpha = config.aux_loss_alpha
155
+ self.seq_aux = config.seq_aux
156
+
157
+ self.norm_topk_prob = config.norm_topk_prob
158
+ self.gating_dim = config.dim
159
+ self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
160
+ self.reset_parameters()
161
+
162
+ def reset_parameters(self) -> None:
163
+ import torch.nn.init as init
164
+ init.kaiming_uniform_(self.weight, a=math.sqrt(5))
165
+
166
+ def forward(self, hidden_states):
167
+ bsz, seq_len, h = hidden_states.shape
168
+ hidden_states = hidden_states.view(-1, h)
169
+ logits = F.linear(hidden_states, self.weight, None)
170
+ if self.scoring_func == 'softmax':
171
+ scores = logits.softmax(dim=-1)
172
+ else:
173
+ raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
174
+
175
+ topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
176
+
177
+ if self.top_k > 1 and self.norm_topk_prob:
178
+ denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
179
+ topk_weight = topk_weight / denominator
180
+
181
+ if self.training and self.alpha > 0.0:
182
+ scores_for_aux = scores
183
+ aux_topk = self.top_k
184
+ topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
185
+ if self.seq_aux:
186
+ scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
187
+ ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
188
+ ce.scatter_add_(1, topk_idx_for_aux_loss,
189
+ torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
190
+ seq_len * aux_topk / self.n_routed_experts)
191
+ aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
192
+ else:
193
+ mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
194
+ ce = mask_ce.float().mean(0)
195
+ Pi = scores_for_aux.mean(0)
196
+ fi = ce * self.n_routed_experts
197
+ aux_loss = (Pi * fi).sum() * self.alpha
198
+ else:
199
+ aux_loss = 0
200
+ return topk_idx, topk_weight, aux_loss
201
+
202
+
203
+ class MOEFeedForward(nn.Module):
204
+ def __init__(self, config: LMConfig):
205
+ super().__init__()
206
+ self.config = config
207
+ self.experts = nn.ModuleList([
208
+ FeedForward(config)
209
+ for _ in range(config.n_routed_experts)
210
+ ])
211
+ self.gate = MoEGate(config)
212
+ if config.n_shared_experts is not None:
213
+ self.shared_experts = FeedForward(config)
214
+
215
+ def forward(self, x):
216
+ identity = x
217
+ orig_shape = x.shape
218
+ bsz, seq_len, _ = x.shape
219
+ # 使用门控机制选择专家
220
+ topk_idx, topk_weight, aux_loss = self.gate(x)
221
+ x = x.view(-1, x.shape[-1])
222
+ flat_topk_idx = topk_idx.view(-1)
223
+ if self.training:
224
+ x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
225
+ y = torch.empty_like(x, dtype=torch.float16)
226
+ for i, expert in enumerate(self.experts):
227
+ y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致
228
+ y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
229
+ y = y.view(*orig_shape)
230
+ else:
231
+ y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
232
+ if self.config.n_shared_experts is not None:
233
+ y = y + self.shared_experts(identity)
234
+ self.aux_loss = aux_loss
235
+ return y
236
+
237
+ @torch.no_grad()
238
+ def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
239
+ expert_cache = torch.zeros_like(x)
240
+ idxs = flat_expert_indices.argsort()
241
+ tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
242
+ token_idxs = idxs // self.config.num_experts_per_tok
243
+ # 当tokens_per_expert = [6, 15, 20, 26],tokens_per_expert.shape[0]即为专家数量(此时为4)
244
+ # 且token_idxs = [3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...] 时
245
+ # 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token(每个token有可能被多个专家处理,这取决于num_experts_per_tok)
246
+ # 接下来9个位置token_idxs[6:15] -> [4, 5, 6, 10, 11, 12...]属于专家1处理的token...依此类推
247
+ for i, end_idx in enumerate(tokens_per_expert):
248
+ start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
249
+ if start_idx == end_idx:
250
+ continue
251
+ expert = self.experts[i]
252
+ exp_token_idx = token_idxs[start_idx:end_idx]
253
+ expert_tokens = x[exp_token_idx]
254
+ expert_out = expert(expert_tokens).to(expert_cache.dtype)
255
+ expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
256
+ expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
257
+
258
+ return expert_cache
259
+
260
+
261
+ class MiniMindBlock(nn.Module):
262
+ def __init__(self, layer_id: int, config: LMConfig):
263
+ super().__init__()
264
+ self.n_heads = config.n_heads
265
+ self.dim = config.dim
266
+ self.head_dim = config.dim // config.n_heads
267
+ self.attention = Attention(config)
268
+
269
+ self.layer_id = layer_id
270
+ self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
271
+ self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
272
+ self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
273
+
274
+ def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
275
+ h_attn, past_kv = self.attention(
276
+ self.attention_norm(x),
277
+ pos_cis,
278
+ past_key_value=past_key_value,
279
+ use_cache=use_cache
280
+ )
281
+ h = x + h_attn
282
+ out = h + self.feed_forward(self.ffn_norm(h))
283
+ return out, past_kv
284
+
285
+
286
+ class MiniMindLM(PreTrainedModel):
287
+ config_class = LMConfig
288
+
289
+ def __init__(self, params: LMConfig = None):
290
+ self.params = params or LMConfig()
291
+ super().__init__(self.params)
292
+ self.vocab_size, self.n_layers = params.vocab_size, params.n_layers
293
+ self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
294
+ self.dropout = nn.Dropout(params.dropout)
295
+ self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])
296
+ self.norm = RMSNorm(params.dim, eps=params.norm_eps)
297
+ self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
298
+ self.tok_embeddings.weight = self.output.weight
299
+ self.register_buffer("pos_cis",
300
+ precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
301
+ persistent=False)
302
+ self.OUT = CausalLMOutputWithPast()
303
+
304
+ def forward(self,
305
+ input_ids: Optional[torch.Tensor] = None,
306
+ past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
307
+ use_cache: bool = False,
308
+ logits_to_keep: Union[int, torch.Tensor] = 0,
309
+ **args):
310
+ past_key_values = past_key_values or [None] * len(self.layers)
311
+ start_pos = args.get('start_pos', 0)
312
+ h = self.dropout(self.tok_embeddings(input_ids))
313
+ pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]
314
+ past_kvs = []
315
+ for l, layer in enumerate(self.layers):
316
+ h, past_kv = layer(
317
+ h, pos_cis,
318
+ past_key_value=past_key_values[l],
319
+ use_cache=use_cache
320
+ )
321
+ past_kvs.append(past_kv)
322
+
323
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
324
+ logits = self.output(self.norm(h)[:, slice_indices, :])
325
+ aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
326
+ self.OUT.__setitem__('last_hidden_state', h)
327
+ self.OUT.__setitem__('logits', logits)
328
+ self.OUT.__setitem__('aux_loss', aux_loss)
329
+ self.OUT.__setitem__('past_key_values', past_kvs)
330
+ return self.OUT
331
+
332
+ @torch.inference_mode()
333
+ def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
334
+ stream=False, rp=1., use_cache=True, pad_token_id=0, num_return_sequences=1, **args):
335
+ # 流式生成
336
+ if stream:
337
+ return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
338
+
339
+ # 直接生成
340
+ generated = []
341
+ for i in range(input_ids.size(0)):
342
+ non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0)
343
+ for _ in range(num_return_sequences):
344
+ out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
345
+ tokens_list = [tokens[:, -1:] for tokens in out]
346
+ gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad
347
+ full_sequence = torch.cat([non_pad, gen], dim=-1)
348
+ generated.append(full_sequence)
349
+
350
+ max_length = max(seq.size(1) for seq in generated)
351
+ generated = [
352
+ torch.cat(
353
+ [seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)],
354
+ dim=-1)
355
+ for seq in generated
356
+ ]
357
+ output = torch.cat(generated, dim=0)
358
+ res = output.view(input_ids.size(0) * num_return_sequences, -1)
359
+ return res
360
+
361
+ def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):
362
+ start, first_seq, past_kvs = input_ids.shape[1], True, None
363
+ while input_ids.shape[1] < max_new_tokens - 1:
364
+ if first_seq or not use_cache:
365
+ out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False
366
+ else:
367
+ out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,
368
+ start_pos=input_ids.shape[1] - 1, **args)
369
+ logits, past_kvs = out.logits[:, -1, :], out.past_key_values
370
+ logits[:, list(set(input_ids.tolist()[0]))] /= rp
371
+ logits /= (temperature + 1e-9)
372
+ if top_p is not None and top_p < 1.0:
373
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
374
+ sorted_probs = F.softmax(sorted_logits, dim=-1)
375
+ cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
376
+ sorted_indices_to_remove = cumulative_probs > top_p
377
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
378
+ sorted_indices_to_remove[:, 0] = False
379
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
380
+ logits[indices_to_remove] = -float('Inf')
381
+ input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
382
+ input_ids = torch.cat((input_ids, input_ids_next), dim=1)
383
+ yield input_ids[:, start:]
384
+ if input_ids_next.item() == eos_token_id:
385
+ break
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4a054638e92d96fccc65e541a8ad72844a513b28910332923452d58599431159
3
+ size 103342988
special_tokens_map.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "</s>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "<unk>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "unk_token": {
24
+ "content": "<unk>",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ }
30
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_eos_token": false,
4
+ "add_prefix_space": false,
5
+ "added_tokens_decoder": {
6
+ "0": {
7
+ "content": "<unk>",
8
+ "lstrip": false,
9
+ "normalized": false,
10
+ "rstrip": false,
11
+ "single_word": false,
12
+ "special": true
13
+ },
14
+ "1": {
15
+ "content": "<s>",
16
+ "lstrip": false,
17
+ "normalized": false,
18
+ "rstrip": false,
19
+ "single_word": false,
20
+ "special": true
21
+ },
22
+ "2": {
23
+ "content": "</s>",
24
+ "lstrip": false,
25
+ "normalized": false,
26
+ "rstrip": false,
27
+ "single_word": false,
28
+ "special": true
29
+ }
30
+ },
31
+ "additional_special_tokens": [],
32
+ "bos_token": "<s>",
33
+ "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<s>system\\n' + system_message + '</s>\\n' }}{% else %}{{ '<s>system\\n你是 MiniMind,是一个有用的人工智能助手。</s>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}",
34
+ "clean_up_tokenization_spaces": false,
35
+ "eos_token": "</s>",
36
+ "extra_special_tokens": {},
37
+ "legacy": true,
38
+ "model_max_length": 32768,
39
+ "pad_token": "<unk>",
40
+ "sp_model_kwargs": {},
41
+ "spaces_between_special_tokens": false,
42
+ "tokenizer_class": "PreTrainedTokenizerFast",
43
+ "unk_token": "<unk>"
44
+ }