text: from torch import nn, torch
  from torch.autograd import Variable
  
  from config import DEVICE
  
  
  class SelfAttention(nn.Module):
      def __init__(self, attention_size,
                   batch_first=False,
                   layers=1,
                   dropout=.0,
                   non_linearity="tanh"):
          super(SelfAttention, self).__init__()
  
          self.batch_first = batch_first
  
          if non_linearity == "relu":
              activation = nn.ReLU()
          else:
              activation = nn.Tanh()
  
          modules = []
          for i in range(layers - 1):
              modules.append(nn.Linear(attention_size, attention_size))
              modules.append(activation)
              modules.append(nn.Dropout(dropout))
  
          # last attention layer must output 1
          modules.append(nn.Linear(attention_size, 1))
          modules.append(activation)
          modules.append(nn.Dropout(dropout))
  
          self.attention = nn.Sequential(*modules)
          self.softmax = nn.Softmax(dim=-1)
  
      @staticmethod
      def get_mask(attentions, lengths):
          """
          Construct mask for padded itemsteps, based on lengths
          """
          max_len = max(lengths.data)
          mask = Variable(torch.ones(attentions.size())).detach()
          mask = mask.to(DEVICE)
  
          for i, l in enumerate(lengths.data):  # skip the first sentence
              if l < max_len:
                  mask[i, l:] = 0
          return mask
  
      def forward(self, inputs, lengths):
  
          ##################################################################
          # STEP 1 - perform dot product
          # of the attention vector and each hidden state
          ##################################################################
  
          # inputs is a 3D Tensor: batch, len, hidden_size
          # scores is a 2D Tensor: batch, len
          scores = self.attention(inputs).squeeze()
          scores = self.softmax(scores)
  
          ##################################################################
          # Step 2 - Masking
          ##################################################################
  
          # construct a mask, based on sentence lengths
          mask = self.get_mask(scores, lengths)
  
          # apply the mask - zero out masked timesteps
          masked_scores = scores * mask
  
          # re-normalize the masked scores
          _sums = masked_scores.sum(-1, keepdim=True)  # sums per row
          scores = masked_scores.div(_sums)  # divide by row sum
  
          ##################################################################
          # Step 3 - Weighted sum of hidden states, by the attention scores
          ##################################################################
  
          # multiply each hidden state with the attention weights
          weighted = torch.mul(inputs, scores.unsqueeze(-1).expand_as(inputs))
  
          # sum the hidden states
          representations = weighted.sum(1).squeeze()
  
          return representations, scores
  