Given logits from a toy language model, implement greedy, top-k and top-p decoding in NumPy.

Greedy takes the argmax. Both samplers sort the probabilities, keep a prefix, renormalise it and sample.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

def pick(logits, rng, mode="greedy", k=5, p=0.9,
         temp=1.0):
    if mode == "greedy":
        return int(np.argmax(logits))
    probs = softmax(logits / temp)
    order = np.argsort(probs)[::-1]   # high to low
    if mode == "top_k":
        keep = order[:k]
    else:                             # top_p
        cum = np.cumsum(probs[order])
        keep = order[:np.searchsorted(cum, p) + 1]
    q = probs[keep] / probs[keep].sum()
    return int(rng.choice(keep, p=q))

Interviewers check the max subtracted in softmax, that top-p keeps the token that crosses p, and the renormalisation before sampling.