NanyuDeer/kronos-prediction
0
1import math
2
3from einops import rearrange, reduce
4import torch
5import torch.nn as nn
6from torch.autograd import Function
7import torch.nn.functional as F
8
9
10class DifferentiableEntropyFunction(Function):
11 @staticmethod
12 def forward(ctx, zq, basis, K, eps):
13 zb = (zq + 1) / 2
14 zi = ((zb * basis).sum(-1)).to(torch.int64)
15 cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),
16 0,
17 zi.flatten(),
18 torch.ones_like(zi.flatten()).to(zq.dtype),
19 'sum')
20 prob = (cnt + eps) / (cnt + eps).sum()
21 H = -(prob * torch.log(prob)).sum()
22 ctx.save_for_backward(zq, zi, prob)
23 ctx.K = K
24 return H
25
26 @staticmethod
27 def backward(ctx, grad_output):
28 zq, zi, prob = ctx.saved_tensors
29 grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K
30 reord_grad = grad_array[zi.flatten()].reshape(zi.shape)
31 grad_input = reord_grad.unsqueeze(-1) * zq
32 return grad_input, None, None, None, None
33
34
35def codebook_entropy(zq, basis, K, eps=1e-4):
36 return DifferentiableEntropyFunction.apply(zq, basis, K, eps)
37
38
39class BinarySphericalQuantizer(nn.Module):
40 def __init__(self, embed_dim, beta, gamma0, gamma, zeta,
41 input_format='bchw',
42 soft_entropy=True, group_size=9,
43 persample_entropy_compute='analytical',
44 cb_entropy_compute='group',
45 l2_norm=True,
46 inv_temperature=1):
47 """
48 Paper link: https://arxiv.org/pdf/2406.07548.pdf
49 Here we use the official implementation of the BinarySphericalQuantizer.
50 """
51 super().__init__()
52 self.embed_dim = embed_dim
53 self.beta = beta # loss weight for commit loss
54 self.gamma0 = gamma0 # loss weight for entropy penalty
55 self.gamma = gamma # loss weight for entropy penalty
56 self.zeta = zeta # loss weight for entire entropy penalty
57 self.input_format = input_format
58 assert self.embed_dim % group_size == 0, "embed_dim must be divisible by group_size"
59 self.num_groups = self.embed_dim // group_size
60 self.group_size = group_size
61 assert persample_entropy_compute in ['group', 'analytical'], "persample_entropy_compute must be either 'group' or 'analytical'"
62 assert cb_entropy_compute in ['group', 'nce'], "cb_entropy_compute must be either 'group' or 'nce'"
63 self.persample_entropy_compute = persample_entropy_compute
64 self.cb_entropy_compute = cb_entropy_compute
65 self.l2_norm = l2_norm
66 self.inv_temperature = inv_temperature
67
68 self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))
69 self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))
70
71 self.num_dimensions = 2 ** embed_dim
72 self.bits_per_index = embed_dim
73
74 # we only need to keep the codebook portion up to the group size
75 # because we approximate the H loss with this subcode
76 group_codes = torch.arange(2 ** self.group_size)
77 group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]
78 self.register_buffer('group_codebook', group_codebook, persistent=False)
79
80 self.soft_entropy = soft_entropy # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf
81
82 def quantize(self, z):
83 assert z.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {z.shape[-1]}"
84
85 zhat = torch.where(z > 0,
86 torch.tensor(1, dtype=z.dtype, device=z.device),
87 torch.tensor(-1, dtype=z.dtype, device=z.device))
88 return z + (zhat - z).detach()
89
90 def forward(self, z, collect_metrics=True):
91 # if self.input_format == 'bchw':
92 # z = rearrange(z, 'b c h w -> b h w c')
93 zq = self.quantize(z)
94
95 q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
96
97 zq = zq * q_scale
98
99 if not collect_metrics:
100 return zq, zq.new_zeros(()), {}
101
102 indices = self.codes_to_indexes(zq.detach())
103 group_indices = self.codes_to_group_indexes(zq.detach())
104 if not self.training:
105 used_codes = torch.unique(indices, return_counts=False)
106 else:
107 used_codes = None
108
109 if self.soft_entropy:
110 persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)
111 entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
112 else:
113 zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)
114 persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)
115 cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)
116 entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy
117
118 # commit loss
119 commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))
120
121 # if self.input_format == 'bchw':
122 # zq = rearrange(zq, 'b h w c -> b c h w')
123
124 return (
125 zq,
126 commit_loss + self.zeta * entropy_penalty / self.inv_temperature,
127 {"H": cb_entropy, "used_codes": used_codes, "indices": indices, "group_indices": group_indices,
128 "avg_prob": avg_prob}
129 )
130
131 def soft_entropy_loss(self, z):
132 # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size
133 # the sub-code is the last group_size bits of the full code
134 group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)
135 divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)
136
137 # we calculate the distance between the divided_z and the codebook for each subgroup
138 distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)
139 prob = (-distance * self.inv_temperature).softmax(dim=-1)
140 if self.persample_entropy_compute == 'analytical':
141 if self.l2_norm:
142 p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)
143 else:
144 p = torch.sigmoid(-4 * z * self.inv_temperature)
145 prob = torch.stack([p, 1 - p], dim=-1)
146 per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
147 else:
148 per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()
149
150 # macro average of the probability of each subgroup
151 avg_prob = reduce(prob, '... g d ->g d', 'mean')
152 codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)
153
154 # the approximation of the entropy is the sum of the entropy of each subgroup
155 return per_sample_entropy, codebook_entropy.sum(), avg_prob
156
157 def get_hard_per_sample_entropy(self, zb_by_sample):
158 probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]
159 persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)
160 persample_entropy = persample_entropy.sum(-1)
161 return persample_entropy.mean()
162
163 def codes_to_indexes(self, zhat):
164 """Converts a `code` to an index in the codebook.
165 Args:
166 zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
167 """
168 assert zhat.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}"
169 return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)
170
171 def codes_to_group_indexes(self, zhat):
172 """Converts a `code` to a list of indexes (in groups) in the codebook.
173 Args:
174 zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}
175 """
176 zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)
177 return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)
178
179 def indexes_to_codes(self, indices):
180 """Inverse of `indexes_to_codes`."""
181 indices = indices.unsqueeze(-1)
182 codes_non_centered = torch.remainder(
183 torch.floor_divide(indices, self.basis), 2
184 )
185 return codes_non_centered * 2 - 1
186
187 def group_indexes_to_codes(self, group_indices):
188 """Inverse of `group_indexes_to_codes`."""
189 group_indices = group_indices.unsqueeze(-1)
190 codes_non_centered = torch.remainder(
191 torch.floor_divide(group_indices, self.group_basis), 2
192 )
193 codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')
194 return codes_non_centered * 2 - 1
195
196 def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):
197 if normalize:
198 probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)
199 else:
200 probs = count
201 H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)
202 return H
203
204 def get_group_codebook_entry(self, group_indices):
205 z_q = self.group_indexes_to_codes(group_indices)
206 q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
207 z_q = z_q * q_scale
208 if self.input_format == 'bchw':
209 h, w = int(z_q.shape[1] ** 0.5)
210 assert h * w == z_q.shape[1], 'Invalid sequence length'
211 z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
212 return z_q
213
214 def get_codebook_entry(self, indices):
215 z_q = self.indexes_to_codes(indices)
216 q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.
217 z_q = z_q * q_scale
218 if self.input_format == 'bchw':
219 h, w = int(z_q.shape[1] ** 0.5)
220 assert h * w == z_q.shape[1], 'Invalid sequence length'
221 z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)
222 return z_q
223
224
225class BSQuantizer(nn.Module):
226
227 def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):
228 super().__init__()
229 self.codebook_dim = s1_bits + s2_bits
230 self.s1_bits = s1_bits
231 self.s2_bits = s2_bits
232 self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)
233
234 def bits_to_indices(self, bits):
235 bits = (bits >= 0).to(torch.long)
236 indices = 2 ** torch.arange(
237 0,
238 bits.shape[-1],
239 1,
240 dtype=torch.long,
241 device=bits.device,
242 )
243 return (bits * indices).sum(-1)
244
245 def forward(self, z, half=False, collect_metrics=True):
246 z = F.normalize(z, dim=-1)
247 quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)
248 if half:
249 q_pre = quantized[:, :, :self.s1_bits]
250 q_post = quantized[:, :, self.s1_bits:]
251 z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]
252 else:
253 z_indices = self.bits_to_indices(quantized)
254 return bsq_loss, quantized, z_indices
255
256
257class RMSNorm(torch.nn.Module):
258 def __init__(self, dim: int, eps: float = 1e-5):
259 super().__init__()
260 self.eps = eps
261 self.weight = nn.Parameter(torch.ones(dim))
262
263 def _norm(self, x):
264 return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
265
266 def forward(self, x):
267 output = self._norm(x.float()).type_as(x)
268 return output * self.weight
269
270
271class FeedForward(nn.Module):
272 def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):
273 super().__init__()
274
275 self.w1 = nn.Linear(d_model, ff_dim, bias=False)
276 self.w3 = nn.Linear(d_model, ff_dim, bias=False)
277 self.w2 = nn.Linear(ff_dim, d_model, bias=False)
278 self.ffn_dropout = nn.Dropout(ffn_dropout_p)
279
280 def forward(self, x):
281 return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
282
283
284class RotaryPositionalEmbedding(nn.Module):
285 def __init__(self, dim):
286 super().__init__()
287 inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
288 self.register_buffer("inv_freq", inv_freq)
289 self.seq_len_cached = None
290 self.cos_cached = None
291 self.sin_cached = None
292
293 def _update_cos_sin_cache(self, x, seq_len):
294 if seq_len != self.seq_len_cached:
295 self.seq_len_cached = seq_len
296 t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
297 freqs = torch.einsum('i,j->ij', t, self.inv_freq)
298 emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
299 self.cos_cached = emb.cos()[None, None, :, :]
300 self.sin_cached = emb.sin()[None, None, :, :]
301 return self.cos_cached, self.sin_cached
302
303 def forward(self, q, k):
304 cos, sin = self._update_cos_sin_cache(q, q.shape[-2])
305 return (
306 (q * cos) + (self._rotate_half(q) * sin),
307 (k * cos) + (self._rotate_half(k) * sin),
308 )
309
310 def _rotate_half(self, x):
311 x1, x2 = x.chunk(2, dim=-1)
312 return torch.cat((-x2, x1), dim=-1)
313
314
315class MultiHeadAttentionWithRoPE(nn.Module):
316 def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):
317 super().__init__()
318 self.d_model = d_model
319 self.n_heads = n_heads
320 self.head_dim = d_model // n_heads
321
322 self.q_proj = nn.Linear(d_model, d_model)
323 self.k_proj = nn.Linear(d_model, d_model)
324 self.v_proj = nn.Linear(d_model, d_model)
325 self.out_proj = nn.Linear(d_model, d_model)
326 self.rotary = RotaryPositionalEmbedding(self.head_dim)
327 self.attn_dropout_p = attn_dropout_p
328 self.resid_dropout = nn.Dropout(resid_dropout_p)
329
330 def forward(self, x, key_padding_mask=None):
331 batch_size, seq_len, _ = x.shape
332
333 q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
334 k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
335 v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
336
337 q, k = self.rotary(q, k)
338
339 if key_padding_mask is not None:
340 attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]
341 attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1) # [batch, n_heads, q_len, k_len]
342 else:
343 attn_mask = None
344
345 attn_output = F.scaled_dot_product_attention(
346 q, k, v,
347 attn_mask=attn_mask,
348 dropout_p=self.attn_dropout_p if self.training else 0.0,
349 is_causal=True
350 )
351
352 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
353 return self.resid_dropout(self.out_proj(attn_output))
354
355
356class MultiHeadCrossAttentionWithRoPE(nn.Module):
357 def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):
358 super().__init__()
359 self.d_model = d_model
360 self.n_heads = n_heads
361 self.head_dim = d_model // n_heads
362
363 self.q_proj = nn.Linear(d_model, d_model)
364 self.k_proj = nn.Linear(d_model, d_model)
365 self.v_proj = nn.Linear(d_model, d_model)
366 self.out_proj = nn.Linear(d_model, d_model)
367 self.rotary = RotaryPositionalEmbedding(self.head_dim)
368 self.attn_dropout_p = attn_dropout_p
369 self.resid_dropout = nn.Dropout(resid_dropout)
370
371 def forward(self, query, key, value, key_padding_mask=None):
372 batch_size, q_len, _ = query.shape
373 _, seq_len, _ = key.shape
374
375 q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)
376 k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
377 v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
378
379 q, k = self.rotary(q, k)
380
381 if key_padding_mask is not None:
382 attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
383 attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)
384 else:
385 attn_mask = None
386
387 is_causal_flag = self.training
388
389 attn_output = F.scaled_dot_product_attention(
390 q, k, v,
391 attn_mask=attn_mask,
392 dropout_p=self.attn_dropout_p if self.training else 0.0,
393 is_causal=is_causal_flag
394 )
395
396 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)
397 return self.resid_dropout(self.out_proj(attn_output))
398
399
400class HierarchicalEmbedding(nn.Module):
401 def __init__(self, s1_bits, s2_bits, d_model=256):
402 super().__init__()
403 self.s1_bits = s1_bits
404 self.s2_bits = s2_bits
405
406 vocab_s1 = 2 ** s1_bits
407 vocab_s2 = 2 ** s2_bits
408
409 self.emb_s1 = nn.Embedding(vocab_s1, d_model)
410 self.emb_s2 = nn.Embedding(vocab_s2, d_model)
411 self.d_model = d_model
412 self.fusion_proj = nn.Linear(d_model * 2, d_model)
413
414 nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)
415 nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)
416
417 def split_token(self, token_ids: torch.Tensor, s2_bits: int):
418 """Inputs:
419 token_ids (torch.Tensor): Composite token IDs of shape [batch_size, seq_len] or [N], each in range [0, 2^(s1_bits + s2_bits) - 1].
420 s2_bits (int): Number of low bits used for the fine token (s2).
421 """
422 assert isinstance(s2_bits, int) and s2_bits > 0, "s2_bits must be a positive integer"
423
424 t = token_ids.long()
425 mask = (1 << s2_bits) - 1
426 s2_ids = t & mask # extract low bits
427 s1_ids = t >> s2_bits # extract high bits
428 return s1_ids, s2_ids
429
430 def forward(self, token_ids):
431 """Inputs:
432 token_ids:
433 - tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or
434 - torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.
435 Output: [batch_size, seq_len, d_model]
436 """
437 if isinstance(token_ids, tuple) or isinstance(token_ids, list):
438 s1_ids, s2_ids = token_ids
439 else:
440 s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)
441 s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)
442 s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)
443 return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))
444
445
446class DependencyAwareLayer(nn.Module):
447 def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):
448 super().__init__()
449 self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)
450 self.norm = RMSNorm(d_model)
451
452 def forward(self, hidden_states, sibling_embed, key_padding_mask=None):
453 """hidden_states: [batch, seq_len, d_model]
454 sibling_embed: Embedding from another subtoken
455 """
456 attn_out = self.cross_attn(
457 query=sibling_embed,
458 key=hidden_states,
459 value=hidden_states,
460 key_padding_mask=key_padding_mask
461 )
462 return self.norm(hidden_states + attn_out)
463
464
465class TransformerBlock(nn.Module):
466 def __init__(self, d_model, n_heads, ff_dim=1024, ffn_dropout_p=0.0, attn_dropout_p=0.0, resid_dropout_p=0.0):
467 super().__init__()
468 self.norm1 = RMSNorm(d_model)
469 self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)
470 self.norm2 = RMSNorm(d_model)
471 self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)
472
473 def forward(self, x, key_padding_mask=None):
474 residual = x
475 x = self.norm1(x)
476 attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)
477 x = residual + attn_out
478
479 residual = x
480 x = self.norm2(x)
481 ffn_out = self.ffn(x)
482 x = residual + ffn_out
483 return x
484
485
486class DualHead(nn.Module):
487 def __init__(self, s1_bits, s2_bits, d_model):
488 super().__init__()
489 self.vocab_s1 = 2 ** s1_bits
490 self.vocab_s2 = 2 ** s2_bits
491 self.proj_s1 = nn.Linear(d_model, self.vocab_s1)
492 self.proj_s2 = nn.Linear(d_model, self.vocab_s2)
493
494 def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):
495 if padding_mask is not None:
496 valid_mask = (padding_mask == 0)
497 s1_logits = s1_logits[valid_mask]
498 s2_logits = s2_logits[valid_mask]
499 s1_targets = s1_targets[valid_mask]
500 s2_targets = s2_targets[valid_mask]
501 ce_s1 = F.cross_entropy(s1_logits, s1_targets)
502 ce_s2 = F.cross_entropy(s2_logits, s2_targets)
503 else:
504 ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))
505 ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))
506 ce_loss = (ce_s1 + ce_s2) / 2
507 return ce_loss, ce_s1, ce_s2
508
509 def forward(self, x):
510 return self.proj_s1(x)
511
512 def cond_forward(self, x2):
513 return self.proj_s2(x2)
514
515
516class FixedEmbedding(nn.Module):
517 def __init__(self, c_in, d_model):
518 super(FixedEmbedding, self).__init__()
519
520 w = torch.zeros(c_in, d_model).float()
521 w.require_grad = False
522
523 position = torch.arange(0, c_in).float().unsqueeze(1)
524 div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
525
526 w[:, 0::2] = torch.sin(position * div_term)
527 w[:, 1::2] = torch.cos(position * div_term)
528
529 self.emb = nn.Embedding(c_in, d_model)
530 self.emb.weight = nn.Parameter(w, requires_grad=False)
531
532 def forward(self, x):
533 return self.emb(x).detach()
534
535
536class TemporalEmbedding(nn.Module):
537 def __init__(self, d_model, learn_pe):
538 super(TemporalEmbedding, self).__init__()
539
540 minute_size = 60
541 hour_size = 24
542 weekday_size = 7
543 day_size = 32
544 month_size = 13
545
546 Embed = FixedEmbedding if not learn_pe else nn.Embedding
547 self.minute_embed = Embed(minute_size, d_model)
548 self.hour_embed = Embed(hour_size, d_model)
549 self.weekday_embed = Embed(weekday_size, d_model)
550 self.day_embed = Embed(day_size, d_model)
551 self.month_embed = Embed(month_size, d_model)
552
553 def forward(self, x):
554 x = x.long()
555
556 minute_x = self.minute_embed(x[:, :, 0])
557 hour_x = self.hour_embed(x[:, :, 1])
558 weekday_x = self.weekday_embed(x[:, :, 2])
559 day_x = self.day_embed(x[:, :, 3])
560 month_x = self.month_embed(x[:, :, 4])
561
562 return hour_x + weekday_x + day_x + month_x + minute_x
563
564
565
566
567
568
569
570
571