"""Reproduce the article's toy example: python3 softmax.py. No model calls."""

import math


def softmax(scores, temperature=1.0):
    """Stable softmax for finite, nonempty scores and positive temperature."""
    if not scores or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Provide nonempty scores and a finite positive temperature")
    if not all(math.isfinite(score) for score in scores):
        raise ValueError("Scores must be finite")
    scaled = [score / temperature for score in scores]
    if not all(math.isfinite(score) for score in scaled):
        raise ValueError("Scaled scores exceed the supported numerical range")
    maximum = max(scaled)
    masses = [math.exp(score - maximum) for score in scaled]
    total = sum(masses)
    return [mass / total for mass in masses]


if __name__ == "__main__":
    probabilities = softmax([2, 1, 0])
    for token, probability in zip("ABC", probabilities):
        print(f"{token}: {probability:.2%}")
    assert math.isclose(sum(probabilities), 1.0)
    assert all(math.isclose(a, b) for a, b in zip(probabilities, softmax([7, 6, 5])))
    assert all(math.isclose(p, 1 / 3) for p in softmax([0, 0, 0]))
    assert softmax([2, 1, 0], temperature=0.5)[0] > probabilities[0]
    print("Checks passed: normalization, equal scores, shift invariance, temperature.")