No-Repeat N-Gram Block

Problem

Ban tokens that would repeat an n-gram. Look at the last n - 1 generated tokens as the suffix; for every earlier occurrence of that suffix, ban the token that followed it by setting its logit to -1e9. For n = 1, ban every already-generated token.

The signature is no_repeat_ngram_block(logits, generated_ids, n). Return a copy of the logits.

Hint

Compare the trailing (n - 1)-gram with every earlier window of the same length.

Test cases

case 1

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

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

case 2

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

expected: [-1000000000,-1000000000,0]

case 3

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

expected: [0,0,0,0]

case 4

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

expected: [0,-1000000000,0]

Starter codePython

def no_repeat_ngram_block(logits, generated_ids, n):
    # Your code here
    pass

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