text: #!/usr/bin/env python3
  
  import torch
  import torch.nn.functional as F
  import numpy as np
  
  
  def create_src_lengths_mask(batch_size, src_lengths):
      """
      Generate boolean mask to prevent attention beyond the end of source
  
      Inputs:
        batch_size : int
        src_lengths : [batch_size] of sentence lengths
  
      Outputs:
        [batch_size, max_src_len]
      """
      max_srclen = src_lengths.max()
      src_indices = torch.arange(
          0,
          max_srclen,
      ).unsqueeze(0).type_as(src_lengths)
      src_indices = src_indices.expand(batch_size, max_srclen)
      src_lengths = src_lengths.unsqueeze(dim=1).expand(
          batch_size,
          max_srclen,
      )
      # returns [batch_size, max_seq_len]
      return (src_indices < src_lengths).int().detach()
  
  
  def masked_softmax(scores, src_lengths, src_length_masking=True):
      """Apply source length masking then softmax.
      Input and output have shape bsz x src_len"""
      if src_length_masking:
          bsz, max_src_len = scores.size()
          # compute masks
          src_mask = create_src_lengths_mask(bsz, src_lengths)
          # Fill pad positions with -inf
          scores = scores.masked_fill(src_mask == 0, -np.inf)
  
      return F.softmax(scores, dim=-1)
  