def beam_search(
input_seq: List[int],
next_token_fn: Callable[[List[int]], List[float]],
max_token: int,
beam_size: int,
stop_word_id: int
) -> List[List[int]]:Part 1 — Beam search on the given signature
Perplexity phone screen with a Chinese interviewer, from one first-hand report (August 2025). The candidate was asked to write beam search: the signature was given, the job was to complete the code, and the unit tests had to pass. The report gives the signature and nothing else: no tests, no examples, no follow-ups, no outcome. Everything below beyond the signature is standard beam search and our setup.
Given — the signature (as reported)
def beam_search(
input_seq: List[int],
next_token_fn: Callable[[List[int]], List[float]],
max_token: int,
beam_size: int,
stop_word_id: int
) -> List[List[int]]:
...
Implement beam_search(input_seq, next_token_fn, max_token, beam_size, stop_word_id).
Our setup (not from the report) — ask your interviewer
The report states none of the rules below. They are standard beam search, and they are what the tests grade:
input_seq. next_token_fn, called on a sequence, gives a score for every possible next token, indexed by token id.beam_size best by cumulative score.stop_word_id is finished and is never extended again. It keeps competing on its score.max_token is reached.beam_size of them. beam_size is at least 1.Conventions the signature leaves open:
next_token_fn returns probabilities. The signature says only List[float], which could also be log-probabilities. A sequence's score is the product of its probabilities (or the sum of their logs, the same order).input_seq, whether they end with the stop token, their order, whether next_token_fn is called with the whole sequence or only the generated tokens, whether max_token counts generated tokens or the whole length, whether a finished beam keeps its slot or moves to a separate list, and whether scores are length-normalised. The tests strip the input prefix and a trailing stop token and ignore order; their inputs make the other readings give the same answer, with no ties, except where the max_token reading changes the result, and there each reading's answer is accepted.max_token ends the search before every beam has stopped, return the beams you have, or only the ones that stopped. Both are accepted.Standard explanation — no candidate report records what the interviewer accepted here.
A product of many probabilities can underflow to 0.0, so distinct beams tie at zero. log is increasing, so summing logs ranks sequences exactly as the product does, without underflow.
Standard explanation — no candidate report records what the interviewer accepted here.
Raw cumulative scores favour short sequences, because every extra token multiplies by a probability below 1. Dividing the log score by the generated length (or by a length penalty such as length to a power alpha) lets longer finished sequences compete. It only matters when finished and live beams of different lengths are compared.
Standard explanation — no candidate report records what the interviewer accepted here.
Each step builds beam_size x V candidates; a full sort is O(B V log(B V)). Taking each beam's top beam_size tokens first, then the best beam_size of those B x B, or using heapq.nlargest, gives the same result in O(B V log B). Storing back-pointers instead of copied sequences avoids copying a sequence per candidate.
One first-hand report (August 2025): a Perplexity phone screen with a Chinese interviewer. The candidate was asked to write beam search. The signature was given, the task was to complete the code, and the unit tests had to pass. The post gives the signature and nothing more: no tests, no examples, no follow-ups, no outcome. So this page rests on that one report; the algorithm is standard beam search, and every convention the signature leaves open is ours, accepted in every reasonable reading by the tests.
Part 1 — Beam search on the given signature
Perplexity phone screen with a Chinese interviewer, from one first-hand report (August 2025). The candidate was asked to write beam search: the signature was given, the job was to complete the code, and the unit tests had to pass. The report gives the signature and nothing else: no tests, no examples, no follow-ups, no outcome. Everything below beyond the signature is standard beam search and our setup.
Given — the signature (as reported)
def beam_search(
input_seq: List[int],
next_token_fn: Callable[[List[int]], List[float]],
max_token: int,
beam_size: int,
stop_word_id: int
) -> List[List[int]]:
...
Implement beam_search(input_seq, next_token_fn, max_token, beam_size, stop_word_id).
Our setup (not from the report) — ask your interviewer
The report states none of the rules below. They are standard beam search, and they are what the tests grade:
input_seq. next_token_fn, called on a sequence, gives a score for every possible next token, indexed by token id.beam_size best by cumulative score.stop_word_id is finished and is never extended again. It keeps competing on its score.max_token is reached.beam_size of them. beam_size is at least 1.Conventions the signature leaves open:
next_token_fn returns probabilities. The signature says only List[float], which could also be log-probabilities. A sequence's score is the product of its probabilities (or the sum of their logs, the same order).input_seq, whether they end with the stop token, their order, whether next_token_fn is called with the whole sequence or only the generated tokens, whether max_token counts generated tokens or the whole length, whether a finished beam keeps its slot or moves to a separate list, and whether scores are length-normalised. The tests strip the input prefix and a trailing stop token and ignore order; their inputs make the other readings give the same answer, with no ties, except where the max_token reading changes the result, and there each reading's answer is accepted.max_token ends the search before every beam has stopped, return the beams you have, or only the ones that stopped. Both are accepted.Our practice question: why sum log-probabilities instead of multiplying probabilities on long sequences?
Our practice question: how would you add length normalisation, and how does it change which finished beams win?
Our practice question: what does one step cost, and how would you avoid sorting all beam_size x vocabulary candidates?
What the report says: a phone screen with a Chinese interviewer. The signature was given, the task was to complete the code, and the unit tests had to pass. Nothing else was posted: no tests, examples, follow-ups or outcome, so this page rests on one report.
Our advice (not from the report): before coding, say which conventions you assume: probabilities or log-probabilities, what max_token counts, and whether the output keeps the input prefix and the stop token. Pool every beam's expansions and keep the global best beam_size; choosing per beam is greedy decoding. Keep a beam that emits the stop token as a result and never extend it, and make sure max_token always ends the loop.