encoders.py

import math
import copy

import torch
from torch import nn
from torch.nn import Module
import torch.nn.functional as F
from torch.nn.init import xavier_uniform_, xavier_normal_


class GRU(Module):
    def __init__(self, hidden_size, dropout_rate, device):
        super(GRU, self).__init__()
        self.hidden_size = hidden_size
        self.emb_size = hidden_size
        self.dropout_prob = dropout_rate
        self.num_layers = 1

        self.emb_dropout = nn.Dropout(self.dropout_prob)
        self.gru_layers = nn.GRU(input_size=self.emb_size,
                                 hidden_size=self.hidden_size,
                                 num_layers=self.num_layers,
                                 bias=False,
                                 batch_first=True)
        self.dense = nn.Linear(self.hidden_size, self.emb_size)
        # self.apply(self._init_weights)
        
    
    def _init_weights(self, module):
        xavier_uniform_(module.weight_hh_l0)
        xavier_uniform_(module.weight_ih_l0)
    

    def forward(self, item_embedding, pos_embeddings, item_seq, item_seq_lens):
        #item_seq_emb = item_embedding(item_seq)
        item_seq_emb = item_embedding[item_seq]
        item_seq_emb_dropout = self.emb_dropout(item_seq_emb)
        gru_output, states_hidden = self.gru_layers(item_seq_emb_dropout)

        return states_hidden


class MultiHeadAttention(Module):
    def __init__(self, n_heads, hidden_size, hidden_dropout_prob, attn_dropout_prob, layer_norm_eps):
        super(MultiHeadAttention, self).__init__()
        if hidden_size % n_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (hidden_size, n_heads)
            )

        self.num_attention_heads = n_heads
        self.attention_head_size = int(hidden_size / n_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.sqrt_attention_head_size = math.sqrt(self.attention_head_size)

        self.query = nn.Linear(hidden_size, self.all_head_size)
        self.key = nn.Linear(hidden_size, self.all_head_size)
        self.value = nn.Linear(hidden_size, self.all_head_size)

        self.softmax = nn.Softmax(dim=-1)
        self.attn_dropout = nn.Dropout(attn_dropout_prob)

        self.dense = nn.Linear(hidden_size, hidden_size)
        self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
        self.out_dropout = nn.Dropout(hidden_dropout_prob)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x

    def forward(self, input_tensor, attention_mask):
        mixed_query_layer = self.query(input_tensor)
        mixed_key_layer = self.key(input_tensor)
        mixed_value_layer = self.value(input_tensor)

        query_layer = self.transpose_for_scores(mixed_query_layer).permute(0, 2, 1, 3)
        key_layer = self.transpose_for_scores(mixed_key_layer).permute(0, 2, 3, 1)
        value_layer = self.transpose_for_scores(mixed_value_layer).permute(0, 2, 1, 3)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer)

        attention_scores = attention_scores / self.sqrt_attention_head_size
        # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
        # [batch_size heads seq_len seq_len] scores
        # [batch_size 1 1 seq_len]
        attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = self.softmax(attention_scores)
        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.

        attention_probs = self.attn_dropout(attention_probs)
        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)
        hidden_states = self.dense(context_layer)
        hidden_states = self.out_dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)

        return hidden_states


class FeedForward(Module):
    def __init__(self, hidden_size, inner_size, hidden_dropout_prob, hidden_act, layer_norm_eps):
        super(FeedForward, self).__init__()
        self.dense_1 = nn.Linear(hidden_size, inner_size)
        self.intermediate_act_fn = self.get_hidden_act(hidden_act)

        self.dense_2 = nn.Linear(inner_size, hidden_size)
        self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
        self.dropout = nn.Dropout(hidden_dropout_prob)

    def get_hidden_act(self, act):
        ACT2FN = {
            "gelu": self.gelu,
            "relu": nn.functional.relu,
            "swish": self.swish,
            "tanh": torch.tanh,
            "sigmoid": torch.sigmoid,
        }
        return ACT2FN[act]

    def gelu(self, x):
        """Implementation of the gelu activation function.
        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results)::
            0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
        Also see https://arxiv.org/abs/1606.08415
        """
        return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

    def swish(self, x):
        return x * torch.sigmoid(x)

    def forward(self, input_tensor):
        hidden_states = self.dense_1(input_tensor)
        hidden_states = self.intermediate_act_fn(hidden_states)

        hidden_states = self.dense_2(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)

        return hidden_states





class TransformerLayer(Module):
    def __init__(self, n_heads, hidden_size, inner_size, dropout, hidden_act='gelu', eps=1e-12):
        super(TransformerLayer, self).__init__()
        self.multi_head_attn = MultiHeadAttention(n_heads, hidden_size, dropout, dropout, eps)
        self.feed_forward = FeedForward(hidden_size, inner_size, dropout, hidden_act, eps)

    def forward(self, hidden_state, attn_mask):
        attn_out = self.multi_head_attn(hidden_state, attn_mask)
        feed_forward_out = self.feed_forward(attn_out)
        return feed_forward_out



class TransformerEncoder(Module):
    def __init__(self, n_layers=2, n_heads=2, hidden_size=64, inner_size=256, dropout=0.5):
        super(TransformerEncoder, self).__init__()
        layer = TransformerLayer(n_heads, hidden_size, inner_size, dropout)
        self.layer_module = nn.ModuleList([copy.deepcopy(layer) for _ in range(n_layers)])
    
    def forward(self, hidden_state, attn_mask, output_all_encode_layers=True):
        all_encode_layers = []
        for layer in self.layer_module:
            hidden_state = layer(hidden_state, attn_mask)
            if output_all_encode_layers:
                all_encode_layers.append(hidden_state)
        if not output_all_encode_layers:
            all_encode_layers.append(hidden_state)
        return all_encode_layers
    


class SASRec(Module):
    def __init__(self, hidden_size, dropout_rate, device, layer_norm_eps=1e-12):
        super(SASRec, self).__init__()
        self.hidden_size = hidden_size
        self.inner_size = hidden_size
        self.dropout = dropout_rate
        self.eps = layer_norm_eps
        self.device = device

        self.n_layers = 2
        self.n_heads = 1

        self.trm_encoder = TransformerEncoder(self.n_layers, self.n_heads, self.hidden_size, self.inner_size, self.dropout)
        self.LayerNorm = nn.LayerNorm(self.hidden_size, self.eps)
        self.dropout = nn.Dropout(self.dropout)

        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    
    def forward(self, item_embeddings, pos_embeddings, item_seq, item_seq_lens):
        # import pdb
        # pdb.set_trace()
        pos_idx = torch.arange(item_seq.size(1), dtype=torch.long).to(self.device)
        pos_idx = pos_idx.unsqueeze(0).expand_as(item_seq)
        pos_emb = pos_embeddings(pos_idx)

        item_emb = item_embeddings[item_seq]
        #item_emb = item_embeddings(item_seq)
        input_emb = item_emb + pos_emb
        input_emb = self.LayerNorm(input_emb)
        input_emb = self.dropout(input_emb)

        # get attention mask 
        attn_mask = (item_seq != 0)
        ext_attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
        ext_attn_mask = torch.where(ext_attn_mask, 0., -10000.)

        trm_output = self.trm_encoder(input_emb, ext_attn_mask, output_all_encode_layers=True)
        output = trm_output[-1]
        gather_idxs = item_seq_lens - 1
        gather_idxs = gather_idxs.view(-1, 1, 1).expand(-1, -1, output.shape[-1])
        output_tensor = output.gather(dim=1, index=gather_idxs)
        output = output_tensor.squeeze(1)
        return output