/
recsys_dev
/
SplitLight
Обзор
Документация
Войти
/
recsys_dev
/
SplitLight
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
runs/rs_src/models.py
230 строк
9 KB
monkey0head
rs model, small fixes
04 фев 2026, 16:13
04 фев 2026, 16:13
64d4840
Код
Авторство
О чём код?
""" Models. """ import numpy as np import torch from torch import nn # from transformers import BertConfig, BertModel from typing import Optional, List # class BERT4Rec(nn.Module): # def __init__(self, vocab_size, bert_config, add_head=True, # tie_weights=True, padding_idx=0, init_std=0.02): # super().__init__() # self.vocab_size = vocab_size # self.bert_config = bert_config # self.add_head = add_head # self.tie_weights = tie_weights # self.padding_idx = padding_idx # self.init_std = init_std # self.embed_layer = nn.Embedding(num_embeddings=vocab_size, # embedding_dim=bert_config['hidden_size'], # padding_idx=padding_idx) # self.transformer_model = BertModel(BertConfig(**bert_config)) # if self.add_head: # self.head = nn.Linear(bert_config['hidden_size'], vocab_size, bias=False) # if self.tie_weights: # self.head.weight = self.embed_layer.weight # self.init_weights() # def init_weights(self): # # initialization in huggingface transformers # # https://github.com/huggingface/transformers/blob/v4.25.1/src/transformers/models/gpt2/modeling_gpt2.py#L462 # # initialization in pytorch Embeddings # # https://github.com/pytorch/pytorch/blob/1.7/torch/nn/modules/sparse.py#L117 # self.embed_layer.weight.data.normal_(mean=0.0, std=self.init_std) # if self.padding_idx is not None: # self.embed_layer.weight.data[self.padding_idx].zero_() # def forward(self, input_ids, attention_mask): # embeds = self.embed_layer(input_ids) # transformer_outputs = self.transformer_model( # inputs_embeds=embeds, attention_mask=attention_mask) # outputs = transformer_outputs.last_hidden_state # if self.add_head: # outputs = self.head(outputs) # return outputs class RNN(nn.Module): def __init__(self, vocab_size, rnn_config, add_head=True, tie_weights=True, padding_idx=0, init_std=0.02): super().__init__() self.vocab_size = vocab_size self.rnn_config = rnn_config self.add_head = add_head self.tie_weights = tie_weights self.padding_idx = padding_idx self.init_std = init_std self.embed_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=rnn_config['input_size'], padding_idx=padding_idx) self.rnn = nn.GRU(batch_first=True, bidirectional=False, **rnn_config) if self.add_head: self.head = nn.Linear(rnn_config['hidden_size'], vocab_size, bias=False) if self.tie_weights: self.head.weight = self.embed_layer.weight self.init_weights() def init_weights(self): self.embed_layer.weight.data.normal_(mean=0.0, std=self.init_std) if self.padding_idx is not None: self.embed_layer.weight.data[self.padding_idx].zero_() # parameter attention mask added for compatibility with Lightning module, not used def forward(self, input_ids, attention_mask): embeds = self.embed_layer(input_ids) outputs, _ = self.rnn(embeds) if self.add_head: outputs = self.head(outputs) return outputs class SASRec(torch.nn.Module): """Adaptation of code from https://github.com/pmixer/SASRec.pytorch. """ def __init__(self, item_num, maxlen=128, hidden_units=64, num_blocks=1, num_heads=1, dropout_rate=0.1, initializer_range=0.02, add_head=True, padding_idx=0): super(SASRec, self).__init__() self.item_num = item_num self.maxlen = maxlen self.hidden_units = hidden_units self.num_blocks = num_blocks self.num_heads = num_heads self.dropout_rate = dropout_rate self.initializer_range = initializer_range self.add_head = add_head self.padding_idx=padding_idx self.item_emb = nn.Embedding(item_num, hidden_units, padding_idx=self.padding_idx) self.pos_emb = nn.Embedding(maxlen, hidden_units) self.emb_dropout = nn.Dropout(dropout_rate) self.attention_layernorms = nn.ModuleList() # to be Q for self-attention self.attention_layers = nn.ModuleList() self.forward_layernorms = nn.ModuleList() self.forward_layers = nn.ModuleList() self.last_layernorm = nn.LayerNorm(hidden_units, eps=1e-8) for _ in range(num_blocks): new_attn_layernorm = nn.LayerNorm(hidden_units, eps=1e-8) self.attention_layernorms.append(new_attn_layernorm) new_attn_layer = nn.MultiheadAttention(hidden_units, num_heads, dropout_rate) self.attention_layers.append(new_attn_layer) new_fwd_layernorm = nn.LayerNorm(hidden_units, eps=1e-8) self.forward_layernorms.append(new_fwd_layernorm) new_fwd_layer = PointWiseFeedForward(hidden_units, dropout_rate) self.forward_layers.append(new_fwd_layer) # parameters initialization self.apply(self._init_weights) def _init_weights(self, module): """Initialize weights. Examples: https://github.com/huggingface/transformers/blob/v4.25.1/src/transformers/models/gpt2/modeling_gpt2.py#L454 https://recbole.io/docs/_modules/recbole/model/sequential_recommender/sasrec.html#SASRec """ if isinstance(module, (nn.Linear, nn.Conv1d)): module.weight.data.normal_(mean=0.0, std=self.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=self.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) # parameter attention mask added for compatibility with GPT Lightning module, not used def forward(self, input_ids, attention_mask): seqs = self.item_emb(input_ids) seqs *= self.item_emb.embedding_dim ** 0.5 positions = np.tile(np.array(range(input_ids.shape[1])), [input_ids.shape[0], 1]) # need to be on the same device seqs += self.pos_emb(torch.LongTensor(positions).to(seqs.device)) seqs = self.emb_dropout(seqs) timeline_mask = torch.Tensor(input_ids == self.padding_idx) seqs *= ~timeline_mask.unsqueeze(-1) # broadcast in last dim tl = seqs.shape[1] # time dim len for enforce causality # need to be on the same device attention_mask = ~torch.tril(torch.ones((tl, tl), dtype=torch.bool).to(seqs.device)) for i in range(len(self.attention_layers)): seqs = torch.transpose(seqs, 0, 1) Q = self.attention_layernorms[i](seqs) mha_outputs, _ = self.attention_layers[i](Q, seqs, seqs, attn_mask=attention_mask) # key_padding_mask=timeline_mask # need_weights=False) this arg do not work? seqs = Q + mha_outputs seqs = torch.transpose(seqs, 0, 1) seqs = self.forward_layernorms[i](seqs) seqs = self.forward_layers[i](seqs) seqs *= ~timeline_mask.unsqueeze(-1) outputs = self.last_layernorm(seqs) # (U, T, C) -> (U, -1, C) if self.add_head: outputs = torch.matmul(outputs, self.item_emb.weight.transpose(0, 1)) return outputs class PointWiseFeedForward(torch.nn.Module): """Code from https://github.com/pmixer/SASRec.pytorch.""" def __init__(self, hidden_units, dropout_rate): super(PointWiseFeedForward, self).__init__() self.conv1 = torch.nn.Conv1d(hidden_units, hidden_units, kernel_size=1) self.dropout1 = torch.nn.Dropout(p=dropout_rate) self.relu = torch.nn.ReLU() self.conv2 = torch.nn.Conv1d(hidden_units, hidden_units, kernel_size=1) self.dropout2 = torch.nn.Dropout(p=dropout_rate) def forward(self, inputs): outputs = self.dropout2( self.conv2(self.relu(self.dropout1(self.conv1(inputs.transpose(-1, -2)))))) outputs = outputs.transpose(-1, -2) # as Conv1D requires (N, C, Length) outputs += inputs return outputs