SFT Loss Masking

Problem

Compute a masked SFT loss over completion tokens: -sum(log_probs[i] * mask[i]) / sum(mask). If the mask sums to zero or the batch is empty return 0.0.

Hint

Masking lets you train only on answer tokens, ignoring the prompt.

Test cases

case 1

in: [[-1,-1],[1,1]]

expected: 1

case 2

in: [[-1,-5],[1,0]]

expected: 1

case 3

in: [[0],[0]]

expected: 0

case 4

in: [[],[]]

expected: 0

Starter codePython

def sft_masked_loss(log_probs, mask):
    # Your code here
    pass

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