Prefix LM Mask

Problem

Build a prefix-LM attention mask. The first prefix_len tokens attend to all prefix tokens (bidirectional); every other token attends to all prefix tokens plus the suffix tokens up to and including itself.

The signature is prefix_lm_mask(prefix_len, total_len). Return a total_len x total_len matrix of 1s and 0s.

Hint

Prefix-LMs combine bidirectional encoding of the prompt with causal decoding of the continuation.

Test cases

case 1

in: [2,4]

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

case 2

in: [1,3]

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

case 3

in: [3,3]

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

case 4

in: [0,2]

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

Starter codePython

def prefix_lm_mask(prefix_len, total_len):
    # Your code here
    pass

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