Top-k and Top-p Sampling | Dataford Interview Questions

Write a function to simulate top-k or top-p (nucleus) sampling during text generation at Mistral AI.

Implement sample_token(logits, top_k=None, top_p=None, seed=0), returning one sampled token index. logits is a non-empty list of floats, exactly one of top_k or top_p is provided, and seed must make sampling reproducible. Apply softmax only after filtering, then sample proportionally to the remaining probabilities. For top-p, sort tokens by descending logit and retain tokens through the first cumulative probability reaching top_p.

Constraints

1 <= len(logits) <= 10000
All logits are finite floating-point values
Exactly one of top_k and top_p is provided
1 <= top_k <= len(logits)
0 < top_p <= 1
Return a token index from 0 through len(logits) - 1
The same inputs and seed must produce the same result

Function Signature

def sample_token(logits, top_k=None, top_p=None, seed=0):

Your question is Top-k and Top-p Sampling. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.

Problem

Write a function to simulate top-k or top-p (nucleus) sampling during text generation at Mistral AI.

Constraints

1 <= len(logits) <= 10000
All logits are finite floating-point values
Exactly one of top_k and top_p is provided
1 <= top_k <= len(logits)
0 < top_p <= 1
Return a token index from 0 through len(logits) - 1
The same inputs and seed must produce the same result

Function Signature

def sample_token(logits, top_k=None, top_p=None, seed=0):

Your question is Top-k and Top-p Sampling. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.