T5 Relative Position Bucket

Problem

Compute the T5-style relative position bucket for a nonnegative distance. Distances below max_exact = num_buckets / 2 map to their own bucket; larger distances go on a logarithmic grid that spreads them up to bucket num_buckets - 1.

The signature is t5_relative_bucket(distance, num_buckets, max_distance). Return an integer bucket, with distance <= 0 mapping to bucket 0 and the logarithmic value clamped to num_buckets - 1.

Hint

Small distances get exact buckets and only the long tail is log-compressed, then clamped at the last bucket.

Test cases

case 1

in: [0,8,128]

expected: 0

case 2

in: [3,8,128]

expected: 3

case 3

in: [5,8,128]

expected: 4

case 4

in: [100,8,128]

expected: 7

case 5

in: [1000,16,64]

expected: 15

Starter codePython

import math
def t5_relative_bucket(distance, num_buckets, max_distance):
    # Your code here
    pass

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