Speculative Draft Accept Length

Problem

Determine how many speculative draft tokens are accepted. For each drafted token, accept it with probability min(1, target_prob(token) / draft_prob(token)); stop at the first rejection. Call random.seed(seed) and draw one random.random() per drafted token.

The signature is speculative_accept_length(draft_tokens, target_probs, draft_probs, seed). Return the accepted prefix length.

Hint

The acceptance rule keeps the speculative distribution identical to sampling from the target model.

Test cases

case 1

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

expected: 2

case 2

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

expected: 1

case 3

in: [[1],[[0.1,0.9]],[[0.9,0.1]],0]

expected: 1

case 4

in: [[0],[[0.01,0.99]],[[0.5,0.5]],0]

expected: 0

case 5

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

expected: 2

Starter codePython

import random
def speculative_accept_length(draft_tokens, target_probs, draft_probs, seed):
    # Your code here
    pass

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