Constrained Decoding Mask

Problem

Apply a hard constraint to decoding: keep logits for allowed token ids and set every other logit to -1e9 so it can never be sampled.

The signature is constrained_mask(logits, allowed_ids). Return a copy of the logits.

Hint

This is how grammars, JSON schemas and forced formats are enforced during generation.

case 1

in: [[0,0,0,0],[1,3]]

expected: [-1000000000,0,-1000000000,0]

case 2

in: [[1,2],[0,1]]

expected: [1,2]

case 3

in: [[5,5],[]]

expected: [-1000000000,-1000000000]

case 4

in: [[0.5,0.5,0.5],[2]]

expected: [-1000000000,-1000000000,0.5]

Starter codePython

def constrained_mask(logits, allowed_ids):
    # Your code here
    pass

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