"""Run with Python 3: python3 t5_buckets.py (no dependencies).
Scalar equivalent of T5Attention._relative_position_bucket, default settings.
Source: https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py
The attention mask is a separate operation.
"""
from math import floor, log


def bucket(offset, bidirectional=True):
    """offset is key position minus query position, in integer tokens."""
    if not isinstance(offset, int):
        raise TypeError("offset must be an integer")
    slots = 16 if bidirectional else 32
    exact = slots // 2
    distance = abs(offset) if bidirectional else max(0, -offset)
    direction = 16 if bidirectional and offset > 0 else 0
    if distance < exact:
        return direction + distance
    growth = floor(log(distance / exact) / log(128 / exact) * exact)
    return direction + min(slots - 1, exact + growth)


if __name__ == "__main__":
    offsets = [-128, -17, -16, -8, -7, 0, 7, 17, 128]
    encoder = [bucket(x) for x in offsets]
    decoder = [bucket(x, False) for x in offsets]
    assert encoder == [15, 10, 10, 8, 7, 0, 23, 26, 31]
    assert decoder == [31, 16, 16, 8, 7, 0, 0, 0, 0]
    assert bucket(-4, False) == 4  # key 124, query 128
    print("Offsets:", offsets)
    print("Encoder:", encoder)
    print("Decoder:", decoder)
    print("Checks passed. Positive decoder offsets still require a causal mask.")