Coverage for hopwise/model/layers.py: 82%

1# @Time : 2020/6/27 16:40

2# @Author : Shanlei Mu

3# @Email : slmu@ruc.edu.cn

4# @File : layers.py

6# UPDATE:

7# @Time : 2022/7/16, 2020/8/24 14:58, 2020/9/16, 2020/9/21, 2020/10/9, 2021/05/01

8# @Author : Zhen Tian, Yujie Lu, Xingyu Pan, Zhichao Feng, Hui Wang, Xinyan Fan

9# @Email : chenyuwuxinn@gmail.com, yujielu1998@gmail.com, panxy@ruc.edu.cn, fzcbupt@gmail.com, hui.wang@ruc.edu.cn, xinyan.fan@ruc.edu.cn # noqa: E501

11# @Time : 2025

12# @Author : Giacomo Medda, Alessandro Soccol

13# @Email : giacomo.medda@unica.it, alessandro.soccol@unica.it

15"""hopwise.model.layers

16#############################

17Common Layers in recommender system

18"""

20import copy

21import math

23import numpy as np

24import torch

25import torch.nn.functional as fn

26from torch import nn

27from torch.nn.init import normal_

29from hopwise.utils import FeatureSource, FeatureType

32class MLPLayers(nn.Module):

33 r"""MLPLayers

35 Args:

36 - layers(list): a list contains the size of each layer in mlp layers

37 - dropout(float): probability of an element to be zeroed. Default: 0

38 - activation(str): activation function after each layer in mlp layers. Default: 'relu'.

39 candidates: 'sigmoid', 'tanh', 'relu', 'leekyrelu', 'none'

41 Shape:

43 - Input: (:math:`N`, \*, :math:`H_{in}`) where \* means any number of additional dimensions

44 :math:`H_{in}` must equal to the first value in `layers`

45 - Output: (:math:`N`, \*, :math:`H_{out}`) where :math:`H_{out}` equals to the last value in `layers`

47 Examples::

49 >>> m = MLPLayers([64, 32, 16], 0.2, 'relu')

50 >>> input = torch.randn(128, 64)

51 >>> output = m(input)

52 >>> print(output.size())

53 >>> torch.Size([128, 16])

54 """

56 def __init__(

57 self,

58 layers,

59 dropout=0.0,

60 activation="relu",

61 bn=False,

62 init_method=None,

63 last_activation=True,

64 ):

65 super().__init__()

66 self.layers = layers

67 self.dropout = dropout

68 self.activation = activation

69 self.use_bn = bn

70 self.init_method = init_method

72 mlp_modules = []

73 for idx, (input_size, output_size) in enumerate(zip(self.layers[:-1], self.layers[1:])):

74 mlp_modules.append(nn.Dropout(p=self.dropout))

75 mlp_modules.append(nn.Linear(input_size, output_size))

76 if self.use_bn:

77 mlp_modules.append(nn.BatchNorm1d(num_features=output_size))

78 activation_func = activation_layer(self.activation, output_size)

79 if activation_func is not None:

80 mlp_modules.append(activation_func)

81 if self.activation is not None and not last_activation:

82 mlp_modules.pop()

83 self.mlp_layers = nn.Sequential(*mlp_modules)

84 if self.init_method is not None:

85 self.apply(self.init_weights)

87 def init_weights(self, module):

88 # We just initialize the module with normal distribution as the paper said

89 if isinstance(module, nn.Linear):

90 if self.init_method == "norm":

91 normal_(module.weight.data, 0, 0.01)

92 if module.bias is not None:

93 module.bias.data.fill_(0.0)

95 def forward(self, input_feature):

96 return self.mlp_layers(input_feature)

99def activation_layer(activation_name="relu", emb_dim=None):

100 """Construct activation layers

101

102 Args:

103 activation_name: str, name of activation function

104 emb_dim: int, used for Dice activation

105

106 Return:

107 activation: activation layer

108 """

109 if activation_name is None:

110 activation = None

111 elif isinstance(activation_name, str):

112 if activation_name.lower() == "sigmoid":

113 activation = nn.Sigmoid()

114 elif activation_name.lower() == "tanh":

115 activation = nn.Tanh()

116 elif activation_name.lower() == "relu":

117 activation = nn.ReLU()

118 elif activation_name.lower() == "leakyrelu":

119 activation = nn.LeakyReLU()

120 elif activation_name.lower() == "dice":

121 activation = Dice(emb_dim)

122 elif activation_name.lower() == "none":

123 activation = None

124 elif issubclass(activation_name, nn.Module):

125 activation = activation_name()

126 else:

127 raise NotImplementedError(f"activation function {activation_name} is not implemented")

128

129 return activation

130

131

132class FMEmbedding(nn.Module):

133 r"""Embedding for token fields.

134

135 Args:

136 field_dims: list, the number of tokens in each token fields

137 offsets: list, the dimension offset of each token field

138 embed_dim: int, the dimension of output embedding vectors

139

140 Input:

141 input_x: tensor, A 3D tensor with shape:``(batch_size,field_size)``.

142

143 Return:

144 output: tensor, A 3D tensor with shape: ``(batch_size,field_size,embed_dim)``.

145 """

146

147 def __init__(self, field_dims, offsets, embed_dim):

148 super().__init__()

149 self.embedding = nn.Embedding(sum(field_dims), embed_dim)

150 self.offsets = offsets

151

152 def forward(self, input_x):

153 input_x = input_x + input_x.new_tensor(self.offsets).unsqueeze(0)

154 output = self.embedding(input_x)

155 return output

156

157

158class FLEmbedding(nn.Module):

159 r"""Embedding for float fields.

160

161 Args:

162 field_dims: list, the number of float in each float fields

163 offsets: list, the dimension offset of each float field

164 embed_dim: int, the dimension of output embedding vectors

165

166 Input:

167 input_x: tensor, A 3D tensor with shape:``(batch_size,field_size,2)``.

168

169 Return:

170 output: tensor, A 3D tensor with shape: ``(batch_size,field_size,embed_dim)``.

171 """

172

173 def __init__(self, field_dims, offsets, embed_dim):

174 super().__init__()

175 self.embedding = nn.Embedding(sum(field_dims), embed_dim)

176 self.offsets = offsets

177

178 def forward(self, input_x):

179 base, index = torch.split(input_x, [1, 1], dim=-1)

180 index = index.squeeze(-1).long()

181 index = index + index.new_tensor(self.offsets).unsqueeze(0)

182 output = base * self.embedding(index)

183 return output

184

185

186class BaseFactorizationMachine(nn.Module):

187 r"""Calculate FM result over the embeddings

188

189 Args:

190 reduce_sum: bool, whether to sum the result, default is True.

191

192 Input:

193 input_x: tensor, A 3D tensor with shape:``(batch_size,field_size,embed_dim)``.

194

195 Output

196 output: tensor, A 3D tensor with shape: ``(batch_size,1)`` or ``(batch_size, embed_dim)``.

197 """

198

199 def __init__(self, reduce_sum=True):

200 super().__init__()

201 self.reduce_sum = reduce_sum

202

203 def forward(self, input_x):

204 square_of_sum = torch.sum(input_x, dim=1) ** 2

205 sum_of_square = torch.sum(input_x**2, dim=1)

206 output = square_of_sum - sum_of_square

207 if self.reduce_sum:

208 output = torch.sum(output, dim=1, keepdim=True)

209 output = 0.5 * output

210 return output

211

212

213class BiGNNLayer(nn.Module):

214 r"""Propagate a layer of Bi-interaction GNN

215

216 .. math::

217 output = (L+I)EW_1 + LE \otimes EW_2

218 """

219

220 def __init__(self, in_dim, out_dim):

221 super().__init__()

222 self.in_dim = in_dim

223 self.out_dim = out_dim

224 self.linear = torch.nn.Linear(in_features=in_dim, out_features=out_dim)

225 self.interActTransform = torch.nn.Linear(in_features=in_dim, out_features=out_dim)

226

227 def forward(self, lap_matrix, eye_matrix, features):

228 # for GCF ajdMat is a (N+M) by (N+M) mat

229 # lap_matrix L = D^-1(A)D^-1 # 拉普拉斯矩阵

230 x = torch.sparse.mm(lap_matrix, features)

231

232 inter_part1 = self.linear(features + x)

233 inter_feature = torch.mul(x, features)

234 inter_part2 = self.interActTransform(inter_feature)

235

236 return inter_part1 + inter_part2

237

238

239class AttLayer(nn.Module):

240 """Calculate the attention signal(weight) according the input tensor.

241

242 Args:

243 infeatures (torch.FloatTensor): A 3D input tensor with shape of[batch_size, M, embed_dim].

244

245 Returns:

246 torch.FloatTensor: Attention weight of input. shape of [batch_size, M].

247 """

248

249 def __init__(self, in_dim, att_dim):

250 super().__init__()

251 self.in_dim = in_dim

252 self.att_dim = att_dim

253 self.w = torch.nn.Linear(in_features=in_dim, out_features=att_dim, bias=False)

254 self.h = nn.Parameter(torch.randn(att_dim), requires_grad=True)

255

256 def forward(self, infeatures):

257 att_signal = self.w(infeatures) # [batch_size, M, att_dim]

258 att_signal = fn.relu(att_signal) # [batch_size, M, att_dim]

259

260 att_signal = torch.mul(att_signal, self.h) # [batch_size, M, att_dim]

261 att_signal = torch.sum(att_signal, dim=2) # [batch_size, M]

262 att_signal = fn.softmax(att_signal, dim=1) # [batch_size, M]

263

264 return att_signal

265

266

267class Dice(nn.Module):

268 r"""Dice activation function

269

270 .. math::

271 f(s)=p(s) \cdot s+(1-p(s)) \cdot \alpha s

272

273 .. math::

274 p(s)=\frac{1} {1 + e^{-\frac{s-E[s]} {\sqrt {Var[s] + \epsilon}}}}

275 """

276

277 def __init__(self, emb_size):

278 super().__init__()

279

280 self.sigmoid = nn.Sigmoid()

281 self.alpha = torch.zeros((emb_size,))

282

283 def forward(self, score):

284 self.alpha = self.alpha.to(score.device)

285 score_p = self.sigmoid(score)

286

287 return self.alpha * (1 - score_p) * score + score_p * score

288

289

290class SequenceAttLayer(nn.Module):

291 """Attention Layer. Get the representation of each user in the batch.

292

293 Args:

294 queries (torch.Tensor): candidate ads, [B, H], H means embedding_size * feat_num

295 keys (torch.Tensor): user_hist, [B, T, H]

296 keys_length (torch.Tensor): mask, [B]

297

298 Returns:

299 torch.Tensor: result

300 """

301

302 def __init__(

303 self,

304 mask_mat,

305 att_hidden_size=(80, 40),

306 activation="sigmoid",

307 softmax_stag=False,

308 return_seq_weight=True,

309 ):

310 super().__init__()

311 self.att_hidden_size = att_hidden_size

312 self.activation = activation

313 self.softmax_stag = softmax_stag

314 self.return_seq_weight = return_seq_weight

315 self.mask_mat = mask_mat

316 self.att_mlp_layers = MLPLayers(self.att_hidden_size, activation=self.activation, bn=False)

317 self.dense = nn.Linear(self.att_hidden_size[-1], 1)

318

319 def forward(self, queries, keys, keys_length):

320 embedding_size = queries.shape[-1] # H

321 hist_len = keys.shape[1] # T

322 queries = queries.repeat(1, hist_len)

323

324 queries = queries.view(-1, hist_len, embedding_size)

325

326 # MLP Layer

327 input_tensor = torch.cat([queries, keys, queries - keys, queries * keys], dim=-1)

328 output = self.att_mlp_layers(input_tensor)

329 output = torch.transpose(self.dense(output), -1, -2)

330

331 # get mask

332 output = output.squeeze(1)

333 mask = self.mask_mat.repeat(output.size(0), 1)

334 mask = mask >= keys_length.unsqueeze(1)

335

336 # mask

337 if self.softmax_stag:

338 mask_value = -np.inf

339 else:

340 mask_value = 0.0

341

342 output = output.masked_fill(mask=mask, value=torch.tensor(mask_value))

343 output = output.unsqueeze(1)

344 output = output / (embedding_size**0.5)

345

346 # get the weight of each user's history list about the target item

347 if self.softmax_stag:

348 output = fn.softmax(output, dim=2) # [B, 1, T]

349

350 if not self.return_seq_weight:

351 output = torch.matmul(output, keys) # [B, 1, H]

352

353 return output

354

355

356class VanillaAttention(nn.Module):

357 """Vanilla attention layer is implemented by linear layer.

358

359 Args:

360 input_tensor (torch.Tensor): the input of the attention layer

361

362 Returns:

363 hidden_states (torch.Tensor): the outputs of the attention layer

364 weights (torch.Tensor): the attention weights

365

366 """

367

368 def __init__(self, hidden_dim, attn_dim):

369 super().__init__()

370 self.projection = nn.Sequential(nn.Linear(hidden_dim, attn_dim), nn.ReLU(True), nn.Linear(attn_dim, 1))

371

372 def forward(self, input_tensor):

373 # (B, Len, num, H) -> (B, Len, num, 1)

374 energy = self.projection(input_tensor)

375 weights = torch.softmax(energy.squeeze(-1), dim=-1)

376 # (B, Len, num, H) * (B, Len, num, 1) -> (B, len, H)

377 hidden_states = (input_tensor * weights.unsqueeze(-1)).sum(dim=-2)

378 return hidden_states, weights

379

380

381class MultiHeadAttention(nn.Module):

382 """Multi-head Self-attention layers, a attention score dropout layer is introduced.

383

384 Args:

385 input_tensor (torch.Tensor): the input of the multi-head self-attention layer

386 attention_mask (torch.Tensor): the attention mask for input tensor

387

388 Returns:

389 hidden_states (torch.Tensor): the output of the multi-head self-attention layer

390

391 """

392

393 def __init__(

394 self,

395 n_heads,

396 hidden_size,

397 hidden_dropout_prob,

398 attn_dropout_prob,

399 layer_norm_eps,

400 ):

401 super().__init__()

402 if hidden_size % n_heads != 0:

403 raise ValueError(

404 "The hidden size (%d) is not a multiple of the number of attention heads (%d)" % (hidden_size, n_heads)

405 )

406

407 self.num_attention_heads = n_heads

408 self.attention_head_size = int(hidden_size / n_heads)

409 self.all_head_size = self.num_attention_heads * self.attention_head_size

410 self.sqrt_attention_head_size = math.sqrt(self.attention_head_size)

411

412 self.query = nn.Linear(hidden_size, self.all_head_size)

413 self.key = nn.Linear(hidden_size, self.all_head_size)

414 self.value = nn.Linear(hidden_size, self.all_head_size)

415

416 self.softmax = nn.Softmax(dim=-1)

417 self.attn_dropout = nn.Dropout(attn_dropout_prob)

418

419 self.dense = nn.Linear(hidden_size, hidden_size)

420 self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)

421 self.out_dropout = nn.Dropout(hidden_dropout_prob)

422

423 def transpose_for_scores(self, x):

424 new_x_shape = x.size()[:-1] + (

425 self.num_attention_heads,

426 self.attention_head_size,

427 )

428 x = x.view(*new_x_shape)

429 return x

430

431 def forward(self, input_tensor, attention_mask):

432 mixed_query_layer = self.query(input_tensor)

433 mixed_key_layer = self.key(input_tensor)

434 mixed_value_layer = self.value(input_tensor)

435

436 query_layer = self.transpose_for_scores(mixed_query_layer).permute(0, 2, 1, 3)

437 key_layer = self.transpose_for_scores(mixed_key_layer).permute(0, 2, 3, 1)

438 value_layer = self.transpose_for_scores(mixed_value_layer).permute(0, 2, 1, 3)

439

440 # Take the dot product between "query" and "key" to get the raw attention scores.

441 attention_scores = torch.matmul(query_layer, key_layer)

442

443 attention_scores = attention_scores / self.sqrt_attention_head_size

444 # Apply the attention mask is (precomputed for all layers in BertModel forward() function)

445 # [batch_size heads seq_len seq_len] scores

446 # [batch_size 1 1 seq_len]

447 attention_scores = attention_scores + attention_mask

448

449 # Normalize the attention scores to probabilities.

450 attention_probs = self.softmax(attention_scores)

451 # This is actually dropping out entire tokens to attend to, which might

452 # seem a bit unusual, but is taken from the original Transformer paper.

453

454 attention_probs = self.attn_dropout(attention_probs)

455 context_layer = torch.matmul(attention_probs, value_layer)

456 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()

457 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)

458 context_layer = context_layer.view(*new_context_layer_shape)

459 hidden_states = self.dense(context_layer)

460 hidden_states = self.out_dropout(hidden_states)

461 hidden_states = self.LayerNorm(hidden_states + input_tensor)

462

463 return hidden_states

464

465

466class FeedForward(nn.Module):

467 """Point-wise feed-forward layer is implemented by two dense layers.

468

469 Args:

470 input_tensor (torch.Tensor): the input of the point-wise feed-forward layer

471

472 Returns:

473 hidden_states (torch.Tensor): the output of the point-wise feed-forward layer

474

475 """

476

477 def __init__(self, hidden_size, inner_size, hidden_dropout_prob, hidden_act, layer_norm_eps):

478 super().__init__()

479 self.dense_1 = nn.Linear(hidden_size, inner_size)

480 self.intermediate_act_fn = self.get_hidden_act(hidden_act)

481

482 self.dense_2 = nn.Linear(inner_size, hidden_size)

483 self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)

484 self.dropout = nn.Dropout(hidden_dropout_prob)

485

486 def get_hidden_act(self, act):

487 ACT2FN = {

488 "gelu": self.gelu,

489 "relu": fn.relu,

490 "swish": self.swish,

491 "tanh": torch.tanh,

492 "sigmoid": torch.sigmoid,

493 }

494 return ACT2FN[act]

495

496 def gelu(self, x):

497 """Implementation of the gelu activation function.

498

499 For information: OpenAI GPT's gelu is slightly different (and gives slightly different results)::

500

501 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))

502

503 Also see https://arxiv.org/abs/1606.08415

504 """

505 return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

506

507 def swish(self, x):

508 return x * torch.sigmoid(x)

509

510 def forward(self, input_tensor):

511 hidden_states = self.dense_1(input_tensor)

512 hidden_states = self.intermediate_act_fn(hidden_states)

513

514 hidden_states = self.dense_2(hidden_states)

515 hidden_states = self.dropout(hidden_states)

516 hidden_states = self.LayerNorm(hidden_states + input_tensor)

517

518 return hidden_states

519

520

521class TransformerLayer(nn.Module):

522 """One transformer layer consists of a multi-head self-attention layer and a point-wise feed-forward layer.

523

524 Args:

525 hidden_states (torch.Tensor): the input of the multi-head self-attention sublayer

526 attention_mask (torch.Tensor): the attention mask for the multi-head self-attention sublayer

527

528 Returns:

529 feedforward_output (torch.Tensor): The output of the point-wise feed-forward sublayer,

530 is the output of the transformer layer.

531

532 """

533

534 def __init__(

535 self,

536 n_heads,

537 hidden_size,

538 intermediate_size,

539 hidden_dropout_prob,

540 attn_dropout_prob,

541 hidden_act,

542 layer_norm_eps,

543 ):

544 super().__init__()

545 self.multi_head_attention = MultiHeadAttention(

546 n_heads, hidden_size, hidden_dropout_prob, attn_dropout_prob, layer_norm_eps

547 )

548 self.feed_forward = FeedForward(

549 hidden_size,

550 intermediate_size,

551 hidden_dropout_prob,

552 hidden_act,

553 layer_norm_eps,

554 )

555

556 def forward(self, hidden_states, attention_mask):

557 attention_output = self.multi_head_attention(hidden_states, attention_mask)

558 feedforward_output = self.feed_forward(attention_output)

559 return feedforward_output

560

561

562class TransformerEncoder(nn.Module):

563 r"""One TransformerEncoder consists of several TransformerLayers.

564

565 Args:

566 n_layers(num): num of transformer layers in transformer encoder. Default: 2

567 n_heads(num): num of attention heads for multi-head attention layer. Default: 2

568 hidden_size(num): the input and output hidden size. Default: 64

569 inner_size(num): the dimensionality in feed-forward layer. Default: 256

570 hidden_dropout_prob(float): probability of an element to be zeroed. Default: 0.5

571 attn_dropout_prob(float): probability of an attention score to be zeroed. Default: 0.5

572 hidden_act(str): activation function in feed-forward layer. Default: 'gelu'

573 candidates: 'gelu', 'relu', 'swish', 'tanh', 'sigmoid'

574 layer_norm_eps(float): a value added to the denominator for numerical stability. Default: 1e-12

575

576 """

577

578 def __init__(

579 self,

580 n_layers=2,

581 n_heads=2,

582 hidden_size=64,

583 inner_size=256,

584 hidden_dropout_prob=0.5,

585 attn_dropout_prob=0.5,

586 hidden_act="gelu",

587 layer_norm_eps=1e-12,

588 ):

589 super().__init__()

590 layer = TransformerLayer(

591 n_heads,

592 hidden_size,

593 inner_size,

594 hidden_dropout_prob,

595 attn_dropout_prob,

596 hidden_act,

597 layer_norm_eps,

598 )

599 self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(n_layers)])

600

601 def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True):

602 """Args:

603 hidden_states (torch.Tensor): the input of the TransformerEncoder

604 attention_mask (torch.Tensor): the attention mask for the input hidden_states

605 output_all_encoded_layers (Bool): whether output all transformer layers' output

606

607 Returns:

608 all_encoder_layers (list): if output_all_encoded_layers is True, return a list consists of all transformer

609 layers' output, otherwise return a list only consists of the output of last transformer layer.

610

611 """

612 all_encoder_layers = []

613 for layer_module in self.layer:

614 hidden_states = layer_module(hidden_states, attention_mask)

615 if output_all_encoded_layers:

616 all_encoder_layers.append(hidden_states)

617 if not output_all_encoded_layers:

618 all_encoder_layers.append(hidden_states)

619 return all_encoder_layers

620

621

622class ItemToInterestAggregation(nn.Module):

623 def __init__(self, seq_len, hidden_size, k_interests=5):

624 super().__init__()

625 self.k_interests = k_interests # k latent interests

626 self.theta = nn.Parameter(torch.randn([hidden_size, k_interests]))

627

628 def forward(self, input_tensor): # [B, L, d] -> [B, k, d]

629 D_matrix = torch.matmul(input_tensor, self.theta) # [B, L, k]

630 D_matrix = nn.Softmax(dim=-2)(D_matrix)

631 result = torch.einsum("nij, nik -> nkj", input_tensor, D_matrix) # #[B, k, d]

632

633 return result

634

635

636class LightMultiHeadAttention(nn.Module):

637 def __init__(

638 self,

639 n_heads,

640 k_interests,

641 hidden_size,

642 seq_len,

643 hidden_dropout_prob,

644 attn_dropout_prob,

645 layer_norm_eps,

646 ):

647 super().__init__()

648 if hidden_size % n_heads != 0:

649 raise ValueError(

650 "The hidden size (%d) is not a multiple of the number of attention heads (%d)" % (hidden_size, n_heads)

651 )

652

653 self.num_attention_heads = n_heads

654 self.attention_head_size = int(hidden_size / n_heads)

655 self.all_head_size = self.num_attention_heads * self.attention_head_size

656

657 # initialization for low-rank decomposed self-attention

658 self.query = nn.Linear(hidden_size, self.all_head_size)

659 self.key = nn.Linear(hidden_size, self.all_head_size)

660 self.value = nn.Linear(hidden_size, self.all_head_size)

661

662 self.attpooling_key = ItemToInterestAggregation(seq_len, hidden_size, k_interests)

663 self.attpooling_value = ItemToInterestAggregation(seq_len, hidden_size, k_interests)

664

665 # initialization for decoupled position encoding

666 self.attn_scale_factor = 2

667 self.pos_q_linear = nn.Linear(hidden_size, self.all_head_size)

668 self.pos_k_linear = nn.Linear(hidden_size, self.all_head_size)

669 self.pos_scaling = float(self.attention_head_size * self.attn_scale_factor) ** -0.5

670 self.pos_ln = nn.LayerNorm(hidden_size, eps=layer_norm_eps)

671

672 self.attn_dropout = nn.Dropout(attn_dropout_prob)

673

674 self.dense = nn.Linear(hidden_size, hidden_size)

675 self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)

676 self.out_dropout = nn.Dropout(hidden_dropout_prob)

677

678 def transpose_for_scores(self, x): # transfor to multihead

679 new_x_shape = x.size()[:-1] + (

680 self.num_attention_heads,

681 self.attention_head_size,

682 )

683 x = x.view(*new_x_shape)

684 return x.permute(0, 2, 1, 3)

685

686 def forward(self, input_tensor, pos_emb):

687 # linear map

688 mixed_query_layer = self.query(input_tensor)

689 mixed_key_layer = self.key(input_tensor)

690 mixed_value_layer = self.value(input_tensor)

691

692 # low-rank decomposed self-attention: relation of items

693 query_layer = self.transpose_for_scores(mixed_query_layer)

694 key_layer = self.transpose_for_scores(self.attpooling_key(mixed_key_layer))

695 value_layer = self.transpose_for_scores(self.attpooling_value(mixed_value_layer))

696

697 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

698 attention_scores = attention_scores / math.sqrt(self.attention_head_size)

699

700 # normalize the attention scores to probabilities.

701 attention_probs = nn.Softmax(dim=-2)(attention_scores)

702 attention_probs = self.attn_dropout(attention_probs)

703 context_layer_item = torch.matmul(attention_probs, value_layer)

704

705 # decoupled position encoding: relation of positions

706 value_layer_pos = self.transpose_for_scores(mixed_value_layer)

707 pos_emb = self.pos_ln(pos_emb).unsqueeze(0)

708 pos_query_layer = self.transpose_for_scores(self.pos_q_linear(pos_emb)) * self.pos_scaling

709 pos_key_layer = self.transpose_for_scores(self.pos_k_linear(pos_emb))

710

711 abs_pos_bias = torch.matmul(pos_query_layer, pos_key_layer.transpose(-1, -2))

712 abs_pos_bias = abs_pos_bias / math.sqrt(self.attention_head_size)

713 abs_pos_bias = nn.Softmax(dim=-2)(abs_pos_bias)

714

715 context_layer_pos = torch.matmul(abs_pos_bias, value_layer_pos)

716

717 context_layer = context_layer_item + context_layer_pos

718

719 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()

720 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)

721 context_layer = context_layer.view(*new_context_layer_shape)

722 hidden_states = self.dense(context_layer)

723 hidden_states = self.out_dropout(hidden_states)

724 hidden_states = self.LayerNorm(hidden_states + input_tensor)

725

726 return hidden_states

727

728

729class LightTransformerLayer(nn.Module):

730 """One transformer layer consists of a multi-head self-attention layer and a point-wise feed-forward layer.

731

732 Args:

733 hidden_states (torch.Tensor): the input of the multi-head self-attention sublayer

734 attention_mask (torch.Tensor): the attention mask for the multi-head self-attention sublayer

735

736 Returns:

737 feedforward_output (torch.Tensor): the output of the point-wise feed-forward sublayer, is the output of the transformer layer

738 """ # noqa: E501

739

740 def __init__(

741 self,

742 n_heads,

743 k_interests,

744 hidden_size,

745 seq_len,

746 intermediate_size,

747 hidden_dropout_prob,

748 attn_dropout_prob,

749 hidden_act,

750 layer_norm_eps,

751 ):

752 super().__init__()

753 self.multi_head_attention = LightMultiHeadAttention(

754 n_heads,

755 k_interests,

756 hidden_size,

757 seq_len,

758 hidden_dropout_prob,

759 attn_dropout_prob,

760 layer_norm_eps,

761 )

762 self.feed_forward = FeedForward(

763 hidden_size,

764 intermediate_size,

765 hidden_dropout_prob,

766 hidden_act,

767 layer_norm_eps,

768 )

769

770 def forward(self, hidden_states, pos_emb):

771 attention_output = self.multi_head_attention(hidden_states, pos_emb)

772 feedforward_output = self.feed_forward(attention_output)

773 return feedforward_output

774

775

776class LightTransformerEncoder(nn.Module):

777 r"""One LightTransformerEncoder consists of several LightTransformerLayers.

778

779 Args:

780 n_layers(num): num of transformer layers in transformer encoder. Default: 2

781 n_heads(num): num of attention heads for multi-head attention layer. Default: 2

782 hidden_size(num): the input and output hidden size. Default: 64

783 inner_size(num): the dimensionality in feed-forward layer. Default: 256

784 hidden_dropout_prob(float): probability of an element to be zeroed. Default: 0.5

785 attn_dropout_prob(float): probability of an attention score to be zeroed. Default: 0.5

786 hidden_act(str): activation function in feed-forward layer. Default: 'gelu'.

787 candidates: 'gelu', 'relu', 'swish', 'tanh', 'sigmoid'

788 layer_norm_eps(float): a value added to the denominator for numerical stability. Default: 1e-12

789 """

790

791 def __init__(

792 self,

793 n_layers=2,

794 n_heads=2,

795 k_interests=5,

796 hidden_size=64,

797 seq_len=50,

798 inner_size=256,

799 hidden_dropout_prob=0.5,

800 attn_dropout_prob=0.5,

801 hidden_act="gelu",

802 layer_norm_eps=1e-12,

803 ):

804 super().__init__()

805 layer = LightTransformerLayer(

806 n_heads,

807 k_interests,

808 hidden_size,

809 seq_len,

810 inner_size,

811 hidden_dropout_prob,

812 attn_dropout_prob,

813 hidden_act,

814 layer_norm_eps,

815 )

816 self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(n_layers)])

817

818 def forward(self, hidden_states, pos_emb, output_all_encoded_layers=True):

819 """Args:

820 hidden_states (torch.Tensor): the input of the TrandformerEncoder

821 attention_mask (torch.Tensor): the attention mask for the input hidden_states

822 output_all_encoded_layers (Bool): whether output all transformer layers' output

823

824 Returns:

825 all_encoder_layers (list): if output_all_encoded_layers is True, return a list consists of all transformer layers' output,

826 otherwise return a list only consists of the output of last transformer layer.

827 """ # noqa: E501

828 all_encoder_layers = []

829 for layer_module in self.layer:

830 hidden_states = layer_module(hidden_states, pos_emb)

831 if output_all_encoded_layers:

832 all_encoder_layers.append(hidden_states)

833 if not output_all_encoded_layers:

834 all_encoder_layers.append(hidden_states)

835 return all_encoder_layers

836

837

838class ContextSeqEmbAbstractLayer(nn.Module):

839 """For Deep Interest Network and feature-rich sequential recommender systems, return features embedding matrices.""" # noqa: E501

840

841 def __init__(self):

842 super().__init__()

843 self.token_field_offsets = {}

844 self.float_field_offsets = {}

845 self.token_embedding_table = nn.ModuleDict()

846 self.float_embedding_table = nn.ModuleDict()

847 self.token_seq_embedding_table = nn.ModuleDict()

848 self.float_seq_embedding_table = nn.ModuleDict()

849

850 self.token_field_names = None

851 self.token_field_dims = None

852 self.float_field_names = None

853 self.float_field_dims = None

854 self.token_seq_field_names = None

855 self.token_seq_field_dims = None

856 self.float_seq_field_names = None

857 self.float_seq_field_dims = None

858 self.num_feature_field = None

859

860 def get_fields_name_dim(self):

861 """Get user feature field and item feature field."""

862 self.token_field_names = {type: [] for type in self.types}

863 self.token_field_dims = {type: [] for type in self.types}

864 self.float_field_names = {type: [] for type in self.types}

865 self.float_field_dims = {type: [] for type in self.types}

866 self.token_seq_field_names = {type: [] for type in self.types}

867 self.token_seq_field_dims = {type: [] for type in self.types}

868 self.num_feature_field = {type: 0 for type in self.types}

869 self.float_seq_field_names = {type: [] for type in self.types}

870 self.float_seq_field_dims = {type: [] for type in self.types}

871

872 for type in self.types:

873 for field_name in self.field_names[type]:

874 if self.dataset.field2type[field_name] == FeatureType.TOKEN:

875 self.token_field_names[type].append(field_name)

876 self.token_field_dims[type].append(self.dataset.num(field_name))

877 elif self.dataset.field2type[field_name] == FeatureType.TOKEN_SEQ:

878 self.token_seq_field_names[type].append(field_name)

879 self.token_seq_field_dims[type].append(self.dataset.num(field_name))

880 elif (

881 self.dataset.field2type[field_name] == FeatureType.FLOAT

882 and field_name in self.dataset.config["numerical_features"]

883 ):

884 self.float_field_names[type].append(field_name)

885 self.float_field_dims[type].append(self.dataset.num(field_name))

886 elif (

887 self.dataset.field2type[field_name] == FeatureType.FLOAT_SEQ

888 and field_name in self.dataset.config["numerical_features"]

889 ):

890 self.float_seq_field_names[type].append(field_name)

891 self.float_seq_field_dims[type].append(self.dataset.num(field_name))

892 else:

893 continue

894 self.num_feature_field[type] += 1

895

896 def get_embedding(self):

897 """Get embedding of all features."""

898 for type in self.types:

899 if len(self.token_field_dims[type]) > 0:

900 self.token_field_offsets[type] = np.array(

901 (0, *np.cumsum(self.token_field_dims[type])[:-1]), dtype=np.long

902 )

903 self.token_embedding_table[type] = FMEmbedding(

904 self.token_field_dims[type],

905 self.token_field_offsets[type],

906 self.embedding_size,

907 ).to(self.device)

908 if len(self.float_field_dims[type]) > 0:

909 self.float_field_offsets[type] = np.array(

910 (0, *np.cumsum(self.float_field_dims[type])[:-1]), dtype=np.long

911 )

912 self.float_embedding_table[type] = FLEmbedding(

913 self.float_field_dims[type],

914 self.float_field_offsets[type],

915 self.embedding_size,

916 ).to(self.device)

917 if len(self.token_seq_field_dims) > 0:

918 self.token_seq_embedding_table[type] = nn.ModuleList()

919 for token_seq_field_dim in self.token_seq_field_dims[type]:

920 self.token_seq_embedding_table[type].append(

921 nn.Embedding(token_seq_field_dim, self.embedding_size).to(self.device)

922 )

923 if len(self.float_seq_field_dims) > 0:

924 self.float_seq_embedding_table[type] = nn.ModuleList()

925 for float_seq_field_dim in self.float_seq_field_dims[type]:

926 self.float_seq_embedding_table[type].append(

927 nn.Embedding(float_seq_field_dim, self.embedding_size).to(self.device)

928 )

929

930 def embed_float_fields(self, float_fields, type, embed=True):

931 """Get the embedding of float fields.

932 In the following three functions("embed_float_fields" "embed_token_fields" "embed_token_seq_fields")

933 when the type is user, [batch_size, max_item_length] should be recognised as [batch_size]

934

935 Args:

936 float_fields(torch.Tensor): [batch_size, max_item_length, num_float_field]

937 type(str): user or item

938 embed(bool): embed or not

939

940 Returns:

941 torch.Tensor: float fields embedding. [batch_size, max_item_length, num_float_field, embed_dim]

942

943 """

944 if float_fields is None:

945 return None

946

947 if type == "item":

948 embedding_shape = float_fields.shape[:-1] + (-1,)

949 float_fields = float_fields.reshape(-1, float_fields.shape[-2], float_fields.shape[-1])

950 float_embedding = self.float_embedding_table[type](float_fields)

951 float_embedding = float_embedding.view(embedding_shape)

952 else:

953 float_embedding = self.float_embedding_table[type](float_fields)

954

955 return float_embedding

956

957 def embed_token_fields(self, token_fields, type):

958 """Get the embedding of token fields

959

960 Args:

961 token_fields(torch.Tensor): input, [batch_size, max_item_length, num_token_field]

962 type(str): user or item

963

964 Returns:

965 torch.Tensor: token fields embedding, [batch_size, max_item_length, num_token_field, embed_dim]

966

967 """

968 if token_fields is None:

969 return None

970 # [batch_size, max_item_length, num_token_field, embed_dim]

971 if type == "item":

972 embedding_shape = token_fields.shape + (-1,)

973 token_fields = token_fields.reshape(-1, token_fields.shape[-1])

974 token_embedding = self.token_embedding_table[type](token_fields)

975 token_embedding = token_embedding.view(embedding_shape)

976 else:

977 token_embedding = self.token_embedding_table[type](token_fields)

978 return token_embedding

979

980 def embed_float_seq_fields(self, float_seq_fields, type):

981 """Embed the float sequence feature columns

982

983 Args:

984 float_seq_fields (torch.FloatTensor): The input tensor. shape of [batch_size, seq_len, 2]

985 mode (str): How to aggregate the embedding of feature in this field. default=mean

986

987 Returns:

988 torch.FloatTensor: The result embedding tensor of float sequence columns.

989 """

990 fields_result = []

991 for i, float_seq_field in enumerate(float_seq_fields):

992 embedding_table = self.float_seq_embedding_table[type][i]

993 base, index = torch.split(float_seq_field, [1, 1], dim=-1)

994 index = index.squeeze(-1)

995 mask = index != 0

996 mask = mask.float()

997 value_cnt = torch.sum(mask, dim=-1, keepdim=True)

998 float_seq_embedding = base * embedding_table(index.long())

999 mask = mask.unsqueeze(-1).expand_as(float_seq_embedding)

1000 if self.pooling_mode == "max":

1001 masked_float_seq_embedding = float_seq_embedding - (1 - mask) * 1e9

1002 result = torch.max(masked_float_seq_embedding, dim=-2, keepdim=True)

1003 result = result.values

1004 elif self.pooling_mode == "sum":

1005 masked_float_seq_embedding = float_seq_embedding * mask.float()

1006 result = torch.sum(masked_float_seq_embedding, dim=-2, keepdim=True)

1007 else:

1008 masked_float_seq_embedding = float_seq_embedding * mask.float()

1009 result = torch.sum(masked_float_seq_embedding, dim=-2)

1010 eps = torch.FloatTensor([1e-8]).to(self.device)

1011 result = torch.div(result, value_cnt + eps)

1012 result = result.unsqueeze(-2)

1013

1014 fields_result.append(result)

1015 if len(fields_result) == 0:

1016 return None

1017 else:

1018 return torch.cat(fields_result, dim=-2)

1019

1020 def embed_token_seq_fields(self, token_seq_fields, type):

1021 """Get the embedding of token_seq fields.

1022

1023 Args:

1024 token_seq_fields(torch.Tensor): input, [batch_size, max_item_length, seq_len]`

1025 type(str): user or item

1026 mode(str): mean/max/sum

1027

1028 Returns:

1029 torch.Tensor: result [batch_size, max_item_length, num_token_seq_field, embed_dim]

1030

1031 """

1032 fields_result = []

1033 for i, token_seq_field in enumerate(token_seq_fields):

1034 embedding_table = self.token_seq_embedding_table[type][i]

1035 mask = token_seq_field != 0 # [batch_size, max_item_length, seq_len]

1036 mask = mask.float()

1037 value_cnt = torch.sum(mask, dim=-1, keepdim=True) # [batch_size, max_item_length, 1]

1038 token_seq_embedding = embedding_table(token_seq_field) # [batch_size, max_item_length, seq_len, embed_dim]

1039 mask = mask.unsqueeze(-1).expand_as(token_seq_embedding)

1040 if self.pooling_mode == "max":

1041 masked_token_seq_embedding = token_seq_embedding - (1 - mask) * 1e9

1042 result = torch.max(

1043 masked_token_seq_embedding, dim=-2, keepdim=True

1044 ) # [batch_size, max_item_length, 1, embed_dim]

1045 result = result.values

1046 elif self.pooling_mode == "sum":

1047 masked_token_seq_embedding = token_seq_embedding * mask.float()

1048 result = torch.sum(

1049 masked_token_seq_embedding, dim=-2, keepdim=True

1050 ) # [batch_size, max_item_length, 1, embed_dim]

1051 else:

1052 masked_token_seq_embedding = token_seq_embedding * mask.float()

1053 result = torch.sum(masked_token_seq_embedding, dim=-2) # [batch_size, max_item_length, embed_dim]

1054 eps = torch.FloatTensor([1e-8]).to(self.device)

1055 result = torch.div(result, value_cnt + eps) # [batch_size, max_item_length, embed_dim]

1056 result = result.unsqueeze(-2) # [batch_size, max_item_length, 1, embed_dim]

1057

1058 fields_result.append(result)

1059 if len(fields_result) == 0:

1060 return None

1061 else:

1062 return torch.cat(fields_result, dim=-2) # [batch_size, max_item_length, num_token_seq_field, embed_dim]

1063

1064 def embed_input_fields(self, user_idx, item_idx):

1065 """Get the embedding of user_idx and item_idx

1066

1067 Args:

1068 user_idx(torch.Tensor): interaction['user_id']

1069 item_idx(torch.Tensor): interaction['item_id_list']

1070

1071 Returns:

1072 dict: embedding of user feature and item feature

1073

1074 """

1075 user_item_feat = {"user": self.user_feat, "item": self.item_feat}

1076 user_item_idx = {"user": user_idx, "item": item_idx}

1077 float_fields_embedding = {}

1078 float_seq_fields_embedding = {}

1079 token_fields_embedding = {}

1080 token_seq_fields_embedding = {}

1081 sparse_embedding = {}

1082 dense_embedding = {}

1083

1084 for type in self.types:

1085 float_fields = []

1086 for field_name in self.float_field_names[type]:

1087 feature = user_item_feat[type][field_name][user_item_idx[type]]

1088 float_fields.append(feature if len(feature.shape) == (3 + (type == "item")) else feature.unsqueeze(-2))

1089 if len(float_fields) > 0:

1090 float_fields = torch.cat(float_fields, dim=-1) # [batch_size, max_item_length, num_float_field]

1091 else:

1092 float_fields = None

1093 float_fields_embedding[type] = self.embed_float_fields(float_fields, type)

1094

1095 float_seq_fields = []

1096 for field_name in self.float_seq_field_names[type]:

1097 feature = user_item_feat[type][field_name][user_item_idx[type]]

1098 float_seq_fields.append(feature)

1099 # [batch_size, max_item_length, num_token_seq_field, embed_dim] or None

1100 float_seq_fields_embedding[type] = self.embed_float_seq_fields(float_seq_fields, type)

1101

1102 if float_fields_embedding[type] is None:

1103 dense_embedding[type] = float_seq_fields_embedding[type]

1104 elif float_seq_fields_embedding[type] is None:

1105 dense_embedding[type] = float_fields_embedding[type]

1106 else:

1107 dense_embedding[type] = torch.cat(

1108 [

1109 float_fields_embedding[type],

1110 float_seq_fields_embedding[type],

1111 ],

1112 dim=-2,

1113 )

1114

1115 token_fields = []

1116 for field_name in self.token_field_names[type]:

1117 feature = user_item_feat[type][field_name][user_item_idx[type]]

1118 token_fields.append(feature.unsqueeze(-1))

1119 if len(token_fields) > 0:

1120 token_fields = torch.cat(token_fields, dim=-1) # [batch_size, max_item_length, num_token_field]

1121 else:

1122 token_fields = None

1123 # [batch_size, max_item_length, num_token_field, embed_dim] or None

1124 token_fields_embedding[type] = self.embed_token_fields(token_fields, type)

1125

1126 token_seq_fields = []

1127 for field_name in self.token_seq_field_names[type]:

1128 feature = user_item_feat[type][field_name][user_item_idx[type]]

1129 token_seq_fields.append(feature)

1130 # [batch_size, max_item_length, num_token_seq_field, embed_dim] or None

1131 token_seq_fields_embedding[type] = self.embed_token_seq_fields(token_seq_fields, type)

1132

1133 if token_fields_embedding[type] is None:

1134 sparse_embedding[type] = token_seq_fields_embedding[type]

1135 elif token_seq_fields_embedding[type] is None:

1136 sparse_embedding[type] = token_fields_embedding[type]

1137 else:

1138 sparse_embedding[type] = torch.cat(

1139 [

1140 token_fields_embedding[type],

1141 token_seq_fields_embedding[type],

1142 ],

1143 dim=-2,

1144 )

1145

1146 # sparse_embedding[type]

1147 # shape: [batch_size, max_item_length, num_token_seq_field+num_token_field, embed_dim] or None

1148 # dense_embedding[type]

1149 # shape: [batch_size, max_item_length, num_float_field]

1150 # or [batch_size, max_item_length, num_float_field, embed_dim] or None

1151 return sparse_embedding, dense_embedding

1152

1153 def forward(self, user_idx, item_idx):

1154 return self.embed_input_fields(user_idx, item_idx)

1155

1156

1157class ContextSeqEmbLayer(ContextSeqEmbAbstractLayer):

1158 """For Deep Interest Network, return all features (including user features and item features) embedding matrices.""" # noqa: E501

1159

1160 def __init__(self, dataset, embedding_size, pooling_mode, device):

1161 super().__init__()

1162 self.device = device

1163 self.embedding_size = embedding_size

1164 self.dataset = dataset

1165 self.user_feat = self.dataset.get_user_feature().to(self.device)

1166 self.item_feat = self.dataset.get_item_feature().to(self.device)

1167

1168 self.field_names = {

1169 "user": list(self.user_feat.interaction.keys()),

1170 "item": list(self.item_feat.interaction.keys()),

1171 }

1172

1173 self.types = ["user", "item"]

1174 self.pooling_mode = pooling_mode

1175 try:

1176 assert self.pooling_mode in ["mean", "max", "sum"]

1177 except AssertionError:

1178 raise AssertionError("Make sure 'pooling_mode' in ['mean', 'max', 'sum']!")

1179 self.get_fields_name_dim()

1180 self.get_embedding()

1181

1182

1183class FeatureSeqEmbLayer(ContextSeqEmbAbstractLayer):

1184 """For feature-rich sequential recommenders, return item features embedding matrices according to

1185 selected features.

1186 """

1187

1188 def __init__(self, dataset, embedding_size, selected_features, pooling_mode, device):

1189 super().__init__()

1190

1191 self.device = device

1192 self.embedding_size = embedding_size

1193 self.dataset = dataset

1194 self.user_feat = None

1195 self.item_feat = self.dataset.get_item_feature().to(self.device)

1196

1197 self.field_names = {"item": selected_features}

1198

1199 self.types = ["item"]

1200 self.pooling_mode = pooling_mode

1201 try:

1202 assert self.pooling_mode in ["mean", "max", "sum"]

1203 except AssertionError:

1204 raise AssertionError("Make sure 'pooling_mode' in ['mean', 'max', 'sum']!")

1205 self.get_fields_name_dim()

1206 self.get_embedding()

1207

1208

1209class CNNLayers(nn.Module):

1210 r"""CNNLayers

1211

1212 Args:

1213 - channels(list): a list contains the channels of each layer in cnn layers

1214 - kernel(list): a list contains the kernels of each layer in cnn layers

1215 - strides(list): a list contains the channels of each layer in cnn layers

1216 - activation(str): activation function after each layer in mlp layers. Default: 'relu'

1217 candidates: 'sigmoid', 'tanh', 'relu', 'leekyrelu', 'none'

1218

1219 Shape:

1220 - Input: :math:`(N, C_{in}, H_{in}, W_{in})`

1221 - Output: :math:`(N, C_{out}, H_{out}, W_{out})` where

1222

1223 .. math::

1224 H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0]

1225 \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor

1226

1227 .. math::

1228 W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1]

1229 \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor

1230

1231 Examples::

1232

1233 >>> m = CNNLayers([1, 32, 32], [2,2], [2,2], 'relu')

1234 >>> input = torch.randn(128, 1, 64, 64)

1235 >>> output = m(input)

1236 >>> print(output.size())

1237 >>> torch.Size([128, 32, 16, 16])

1238 """

1239

1240 def __init__(self, channels, kernels, strides, activation="relu", init_method=None):

1241 super().__init__()

1242 self.channels = channels

1243 self.kernels = kernels

1244 self.strides = strides

1245 self.activation = activation

1246 self.init_method = init_method

1247 self.num_of_nets = len(self.channels) - 1

1248

1249 if len(kernels) != len(strides) or self.num_of_nets != (len(kernels)):

1250 raise RuntimeError("channels, kernels and strides don't match\n")

1251

1252 cnn_modules = []

1253

1254 for i in range(self.num_of_nets):

1255 cnn_modules.append(

1256 nn.Conv2d(

1257 self.channels[i],

1258 self.channels[i + 1],

1259 self.kernels[i],

1260 stride=self.strides[i],

1261 )

1262 )

1263 if self.activation.lower() == "sigmoid":

1264 cnn_modules.append(nn.Sigmoid())

1265 elif self.activation.lower() == "tanh":

1266 cnn_modules.append(nn.Tanh())

1267 elif self.activation.lower() == "relu":

1268 cnn_modules.append(nn.ReLU())

1269 elif self.activation.lower() == "leakyrelu":

1270 cnn_modules.append(nn.LeakyReLU())

1271 elif self.activation.lower() == "none":

1272 pass

1273

1274 self.cnn_layers = nn.Sequential(*cnn_modules)

1275

1276 if self.init_method is not None:

1277 self.apply(self.init_weights)

1278

1279 def init_weights(self, module):

1280 # We just initialize the module with normal distribution as the paper said

1281 if isinstance(module, nn.Conv2d):

1282 if self.init_method == "norm":

1283 normal_(module.weight.data, 0, 0.01)

1284 if module.bias is not None:

1285 module.bias.data.fill_(0.0)

1286

1287 def forward(self, input_feature):

1288 return self.cnn_layers(input_feature)

1289

1290

1291class FMFirstOrderLinear(nn.Module):

1292 """Calculate the first order score of the input features.

1293 This class is a member of ContextRecommender, you can call it easily when inherit ContextRecommender.

1294

1295 """

1296

1297 def __init__(self, config, dataset, output_dim=1):

1298 super().__init__()

1299 self.field_names = dataset.fields(

1300 source=[

1301 FeatureSource.INTERACTION,

1302 FeatureSource.USER,

1303 FeatureSource.USER_ID,

1304 FeatureSource.ITEM,

1305 FeatureSource.ITEM_ID,

1306 ]

1307 )

1308 self.LABEL = config["LABEL_FIELD"]

1309 self.device = config["device"]

1310 self.numerical_features = config["numerical_features"]

1311 self.token_field_names = []

1312 self.token_field_dims = []

1313 self.float_field_names = []

1314 self.float_field_dims = []

1315 self.token_seq_field_names = []

1316 self.token_seq_field_dims = []

1317 self.float_seq_field_names = []

1318 self.float_seq_field_dims = []

1319

1320 for field_name in self.field_names:

1321 if field_name == self.LABEL:

1322 continue

1323 if dataset.field2type[field_name] == FeatureType.TOKEN:

1324 self.token_field_names.append(field_name)

1325 self.token_field_dims.append(dataset.num(field_name))

1326 elif dataset.field2type[field_name] == FeatureType.TOKEN_SEQ:

1327 self.token_seq_field_names.append(field_name)

1328 self.token_seq_field_dims.append(dataset.num(field_name))

1329 elif dataset.field2type[field_name] == FeatureType.FLOAT and field_name in self.numerical_features:

1330 self.float_field_names.append(field_name)

1331 self.float_field_dims.append(dataset.num(field_name))

1332 elif dataset.field2type[field_name] == FeatureType.FLOAT_SEQ and field_name in self.numerical_features:

1333 self.float_seq_field_names.append(field_name)

1334 self.float_seq_field_dims.append(dataset.num(field_name))

1335

1336 if len(self.token_field_dims) > 0:

1337 self.token_field_offsets = np.array((0, *np.cumsum(self.token_field_dims)[:-1]), dtype=np.long)

1338 self.token_embedding_table = FMEmbedding(self.token_field_dims, self.token_field_offsets, output_dim)

1339 if len(self.float_field_dims) > 0:

1340 self.float_field_offsets = np.array((0, *np.cumsum(self.float_field_dims)[:-1]), dtype=np.long)

1341 self.float_embedding_table = FLEmbedding(self.float_field_dims, self.float_field_offsets, output_dim)

1342 if len(self.token_seq_field_dims) > 0:

1343 self.token_seq_embedding_table = nn.ModuleList()

1344 for token_seq_field_dim in self.token_seq_field_dims:

1345 self.token_seq_embedding_table.append(nn.Embedding(token_seq_field_dim, output_dim))

1346 if len(self.float_seq_field_dims) > 0:

1347 self.float_seq_embedding_table = nn.ModuleList()

1348 for float_seq_field_dim in self.float_seq_field_dims:

1349 self.float_seq_embedding_table.append(nn.Embedding(float_seq_field_dim, output_dim))

1350

1351 self.bias = nn.Parameter(torch.zeros((output_dim,)), requires_grad=True)

1352

1353 def embed_float_fields(self, float_fields):

1354 """Embed the float feature columns

1355

1356 Args:

1357 float_fields (torch.FloatTensor): The input dense tensor. shape of [batch_size, num_float_field, 2]

1358 embed (bool): Return the embedding of columns or just the columns itself. Defaults to ``True``.

1359

1360 Returns:

1361 torch.FloatTensor: The result embedding tensor of float columns.

1362 """

1363 # input Tensor shape : [batch_size, num_float_field]

1364 if float_fields is None:

1365 return None

1366 # [batch_size, num_float_field, embed_dim]

1367 float_embedding = self.float_embedding_table(float_fields)

1368

1369 # [batch_size, 1, output_dim]

1370 float_embedding = torch.sum(float_embedding, dim=1, keepdim=True)

1371 return float_embedding

1372

1373 def embed_float_seq_fields(self, float_seq_fields, mode="mean"):

1374 """Embed the float sequence feature columns

1375

1376 Args:

1377 float_seq_fields (torch.LongTensor): The input tensor. shape of [batch_size, seq_len, 2]

1378 mode (str): How to aggregate the embedding of feature in this field. default=mean

1379

1380 Returns:

1381 torch.FloatTensor: The result embedding tensor of float sequence columns.

1382 """

1383 # input is a list of Tensor shape of [batch_size, seq_len]

1384 fields_result = []

1385 for i, float_seq_field in enumerate(float_seq_fields):

1386 embedding_table = self.float_seq_embedding_table[i]

1387 base, index = torch.split(float_seq_field, [1, 1], dim=-1)

1388 index = index.squeeze(-1)

1389 mask = index != 0 # [batch_size, seq_len]

1390 mask = mask.float()

1391 value_cnt = torch.sum(mask, dim=1, keepdim=True) # [batch_size, 1]

1392

1393 float_seq_embedding = base * embedding_table(index.long()) # [batch_size, seq_len, embed_dim]

1394

1395 mask = mask.unsqueeze(2).expand_as(float_seq_embedding) # [batch_size, seq_len, embed_dim]

1396 if mode == "max":

1397 masked_float_seq_embedding = float_seq_embedding - (1 - mask) * 1e9 # [batch_size, seq_len, embed_dim]

1398 result = torch.max(masked_float_seq_embedding, dim=1, keepdim=True) # [batch_size, 1, embed_dim]

1399 elif mode == "sum":

1400 masked_float_seq_embedding = float_seq_embedding * mask.float()

1401 result = torch.sum(masked_float_seq_embedding, dim=1, keepdim=True) # [batch_size, 1, embed_dim]

1402 else:

1403 masked_float_seq_embedding = float_seq_embedding * mask.float()

1404 result = torch.sum(masked_float_seq_embedding, dim=1) # [batch_size, embed_dim]

1405 eps = torch.FloatTensor([1e-8]).to(self.device)

1406 result = torch.div(result, value_cnt + eps) # [batch_size, embed_dim]

1407 result = result.unsqueeze(1) # [batch_size, 1, embed_dim]

1408 fields_result.append(result)

1409 if len(fields_result) == 0:

1410 return None

1411 else:

1412 return torch.sum(

1413 torch.cat(fields_result, dim=1), dim=1, keepdim=True

1414 ) # [batch_size, num_token_seq_field, embed_dim]

1415

1416 def embed_token_fields(self, token_fields):

1417 """Calculate the first order score of token feature columns

1418

1419 Args:

1420 token_fields (torch.LongTensor): The input tensor. shape of [batch_size, num_token_field]

1421

1422 Returns:

1423 torch.FloatTensor: The first order score of token feature columns

1424 """

1425 # input Tensor shape : [batch_size, num_token_field]

1426 if token_fields is None:

1427 return None

1428 # [batch_size, num_token_field, embed_dim]

1429 token_embedding = self.token_embedding_table(token_fields)

1430 # [batch_size, 1, output_dim]

1431 token_embedding = torch.sum(token_embedding, dim=1, keepdim=True)

1432

1433 return token_embedding

1434

1435 def embed_token_seq_fields(self, token_seq_fields):

1436 """Calculate the first order score of token sequence feature columns

1437

1438 Args:

1439 token_seq_fields (torch.LongTensor): The input tensor. shape of [batch_size, seq_len]

1440

1441 Returns:

1442 torch.FloatTensor: The first order score of token sequence feature columns

1443 """

1444 # input is a list of Tensor shape of [batch_size, seq_len]

1445 fields_result = []

1446 for i, token_seq_field in enumerate(token_seq_fields):

1447 embedding_table = self.token_seq_embedding_table[i]

1448 mask = token_seq_field != 0 # [batch_size, seq_len]

1449 mask = mask.float()

1450

1451 token_seq_embedding = embedding_table(token_seq_field) # [batch_size, seq_len, output_dim]

1452

1453 mask = mask.unsqueeze(2).expand_as(token_seq_embedding) # [batch_size, seq_len, output_dim]

1454 masked_token_seq_embedding = token_seq_embedding * mask.float()

1455 result = torch.sum(masked_token_seq_embedding, dim=1, keepdim=True) # [batch_size, 1, output_dim]

1456

1457 fields_result.append(result)

1458 if len(fields_result) == 0:

1459 return None

1460 else:

1461 return torch.sum(torch.cat(fields_result, dim=1), dim=1, keepdim=True) # [batch_size, 1, output_dim]

1462

1463 def forward(self, interaction):

1464 total_fields_embedding = []

1465 float_fields = []

1466 for field_name in self.float_field_names:

1467 if len(interaction[field_name].shape) == 3: # noqa: PLR2004

1468 float_fields.append(interaction[field_name])

1469 else:

1470 float_fields.append(interaction[field_name].unsqueeze(1))

1471

1472 if len(float_fields) > 0:

1473 float_fields = torch.cat(float_fields, dim=1)

1474 else:

1475 float_fields = None

1476

1477 float_fields_embedding = self.embed_float_fields(float_fields)

1478

1479 if float_fields_embedding is not None:

1480 total_fields_embedding.append(float_fields_embedding)

1481

1482 float_seq_fields = []

1483 for field_name in self.float_seq_field_names:

1484 float_seq_fields.append(interaction[field_name])

1485

1486 float_seq_fields_embedding = self.embed_float_seq_fields(float_seq_fields)

1487

1488 if float_seq_fields_embedding is not None:

1489 total_fields_embedding.append(float_seq_fields_embedding)

1490

1491 token_fields = []

1492 for field_name in self.token_field_names:

1493 token_fields.append(interaction[field_name].unsqueeze(1))

1494 if len(token_fields) > 0:

1495 token_fields = torch.cat(token_fields, dim=1) # [batch_size, num_token_field]

1496 else:

1497 token_fields = None

1498 # [batch_size, 1, output_dim] or None

1499 token_fields_embedding = self.embed_token_fields(token_fields)

1500 if token_fields_embedding is not None:

1501 total_fields_embedding.append(token_fields_embedding)

1502

1503 token_seq_fields = []

1504 for field_name in self.token_seq_field_names:

1505 token_seq_fields.append(interaction[field_name])

1506 # [batch_size, 1, output_dim] or None

1507 token_seq_fields_embedding = self.embed_token_seq_fields(token_seq_fields)

1508 if token_seq_fields_embedding is not None:

1509 total_fields_embedding.append(token_seq_fields_embedding)

1510

1511 return torch.sum(torch.cat(total_fields_embedding, dim=1), dim=1) + self.bias # [batch_size, output_dim]

1512

1513

1514class SparseDropout(nn.Module):

1515 """This is a Module that execute Dropout on Pytorch sparse tensor."""

1516

1517 def __init__(self, p=0.5):

1518 super().__init__()

1519 # p is ratio of dropout

1520 # convert to keep probability

1521 self.kprob = 1 - p

1522

1523 def forward(self, x):

1524 if not self.training:

1525 return x

1526

1527 mask = ((torch.rand(x._values().size()) + self.kprob).floor()).type(torch.bool)

1528 rc = x._indices()[:, mask]

1529 val = x._values()[mask] * (1.0 / self.kprob)

1530 return torch.sparse.FloatTensor(rc, val, x.shape)

1531

1532

1533class ResidualBlock(nn.Module):

1534 """

1535 A Residual Block module. Used in RPG

1536

1537 This module performs a linear transformation followed by a SiLU activation,

1538 and then adds the result to the original input, creating a residual connection.

1539

1540 Args:

1541 hidden_size (int): The size of the hidden layers in the block.

1542 """

1543

1544 def __init__(self, hidden_size):

1545 super().__init__()

1546 self.linear = nn.Linear(hidden_size, hidden_size)

1547 # Initialize as an identity mapping

1548 torch.nn.init.zeros_(self.linear.weight)

1549 torch.nn.init.zeros_(self.linear.bias)

1550 # Use SiLU activation to keep consistent with the Llama model

1551 self.act = nn.SiLU()

1552

1553 def forward(self, x):

1554 """

1555 Forward pass of the ResBlock.

1556

1557 Args:

1558 x (torch.Tensor): Input tensor.

1559

1560 Returns:

1561 torch.Tensor: Output after the residual connection and activation.

1562 """

1563 return x + self.act(self.linear(x))