Repetition Penalty

LLMs sometimes get stuck repeating themselves. A repetition penalty $\theta > 1$ makes tokens that have already been generated less likely before sampling the next one:

• if the token's logit is positive, divide it by $\theta$
• if it's negative or zero, multiply it by $\theta$

(Dividing a negative logit would make it more likely — a classic bug.)

Implement apply_repetition_penalty(logits, generated_ids, penalty) returning a new NumPy array. Each token is penalised once, however many times it appeared. Don't modify the input array.

Example

apply_repetition_penalty(np.array([2.0, -1.0, 0.5]), [0, 1, 1], 2.0)
# array([ 1. , -2. ,  0.5])