random_06.py

def sample_next_token(rng, logits, temperature=1.0):
    """Sample next token from language model logits"""
    # Apply temperature
    scaled_logits = [l/temperature for l in logits]
    
    # Convert to probabilities with softmax
    max_logit = max(scaled_logits)
    exp_logits = [math.exp(l - max_logit) for l in scaled_logits]
    total = sum(exp_logits)
    probs = [e/total for e in exp_logits]
    
    # Sample using our weighted random choice
    return random_choice(rng, range(len(logits)), probs)