The Full Sampling Pipeline

medium
Unit 20 · RLHF & Embeddings0/2 to milestone

Real samplers chain several steps, and order matters. Implement the order Hugging Face's generate uses:

  1. Repetition penalty on previously generated tokens (positive logits ÷ θ, others × θ)
  2. Temperature: divide logits by temperature
  3. Top-k: keep the top_k highest logits (ties at the boundary all kept), and set others to -inf
  4. Top-p: softmax, sort descending, keep the smallest prefix whose cumulative probability reaches top_p (always at least one token)
  5. Renormalise

Implement next_token_probs(logits, generated=(), repetition_penalty=1.0, temperature=1.0, top_k=None, top_p=1.0) returning the final probability array. Any step at its default value is skipped.

#sampling#decoding#llm2 visible + 3 hidden tests
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
The Full Sampling Pipeline · TechGradify