CTRL-Style Repetition Penalty

Problem

Apply the CTRL-style repetition penalty to logits: for each token id already generated, divide its logit by penalty when positive and multiply it by penalty when negative; other logits stay unchanged.

The signature is repetition_penalty(logits, generated_ids, penalty). Return a copy of the logits.

Hint

Positive logits shrink and negative logits grow more negative, both discouraging repeats.

Test cases

case 1

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

expected: [0.5,-2,0.5]

case 2

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

expected: [2,2]

case 3

in: [[0,-3],[0],1.5]

expected: [0,-3]

case 4

in: [[4,2,-2],[0,2],2]

expected: [2,2,-4]

Starter codePython

def repetition_penalty(logits, generated_ids, penalty):
    # Your code here
    pass

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