~/problems / AI infrastructure

Next-token sampler

On a phone? Coding is easier on a laptop: email this problem to yourself . Meanwhile: fight a boss.

medium assessment 4 levels ~50 min

Level 1 Greedy decoding

A language model scores every token in its vocabulary with a number called a logit: the higher the logit, the more likely the token. A sampler turns those scores into the one token the model says next. Over four levels you'll build Sampler, starting with the simplest rule of all.

  • Sampler(seed): seed is an integer in 0..2^31 - 1. You won't need it until level 2.
  • sample(logits) -> int: logits is a non-empty list of floats, where logits[i] is token i's score. Return the index of the largest logit. If several tokens share the largest logit, return the lowest index.
s = Sampler(42)
s.sample([1.5, 3.0, -2.0, 3.0])   # 1   (tokens 1 and 3 tie; the lower index wins)
s.sample([-0.25])                 # 0
s.sample([-4.0, -1.5, -3.0])      # 1

Constraints: up to 100,000 logits per call, each in [-100, 100].

Show hint

read all four levels before you start. Later levels add temperature, top-k and top-p settings that change how sample works, and the greedy rule stays as the special case "temperature 0".

Level 2 unlocks when level 1 passes.

Level 3 unlocks when level 2 passes.

Level 4 unlocks when level 3 passes.

Topic: AI infrastructure. The plumbing around models: request batching, streaming responses, prompt caches, sampling, token limits and eval harnesses.

0:00
Ctrl ' run · Ctrl ↵ submit
esc