Attention Mask With Padding

Problem

Build a 2D padding mask for one sequence of length total_len where only length positions are real. Mask[i][j] is 1 when both positions i and j are real. Right padding puts real tokens first, left padding puts them last.

The signature is padding_attention_mask(length, total_len, pad_side).

Hint

Padding tokens must never be attended to, otherwise they leak garbage into the representation.

Test cases

case 1

in: [2,4,"right"]

expected: [[1,1,0,0],[1,1,0,0],[0,0,0,0],[0,0,0,0]]

case 2

in: [2,4,"left"]

expected: [[0,0,0,0],[0,0,0,0],[0,0,1,1],[0,0,1,1]]

case 3

in: [4,4,"right"]

expected: [[1,1,1,1],[1,1,1,1],[1,1,1,1],[1,1,1,1]]

case 4

in: [0,3,"right"]

expected: [[0,0,0],[0,0,0],[0,0,0]]

Starter codePython

def padding_attention_mask(length, total_len, pad_side):
    # Your code here
    pass

4 test cases · run your code against them in the editor