Implement greedy decoding and beam search
Company: Microsoft
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Technical Screen
You are given access to an auto-regressive language model that can score the next token based on a prefix.
Assume:
- `next_logprobs(prefix_tokens) -> Map[token, logp]` returns log-probabilities (natural log) for the next token given the current prefix.
- `EOS` is a special end-of-sequence token.
- You are also given `max_len` and an initial prefix `prompt_tokens`.
Implement two decoding methods:
1) **Greedy decoding (greedy sampling)**
- Repeatedly select the single token with the highest log-probability at each step.
- Stop when you generate `EOS` or when the total length reaches `max_len`.
- Return the generated token sequence (including the prompt, or clearly specify whether you return only newly generated tokens).
2) **Beam search (keep top-k paths)**
- Given an integer `k` (the number of beams to keep), implement beam search that maintains the **top-k candidate sequences** by total sequence score.
- Use **cumulative log-probability** (sum of token log-probs) as the sequence score (no need for length normalization unless you choose to add it and explain).
- At each decoding step, expand each current beam by one token using `next_logprobs`, then keep only the top-k resulting sequences overall.
- Stop when all beams have generated `EOS` or when reaching `max_len`.
- Return the best sequence (and optionally the top-k sequences if requested).
Clarify any edge cases you handle (e.g., early-finished beams, ties, empty vocabulary, whether finished beams remain in the beam set without further expansion).
Quick Answer: This question evaluates understanding of autoregressive decoding strategies—greedy decoding and beam search—including cumulative log-probability scoring, beam maintenance, EOS handling, and sequence expansion logic; it is categorized under Coding & Algorithms for Software Engineer roles.
Part 1: Greedy Decoding from Prefix Log-Probabilities
You are given a simplified auto-regressive language model as a dictionary called `scores`. For any current token list `prefix_tokens`, the available next-token log-probabilities are stored in `scores[' '.join(prefix_tokens)]`. The empty prefix uses the key `''`.
Implement greedy decoding. Starting from `prompt_tokens`, repeatedly append the token with the highest log-probability for the current prefix. If multiple tokens have the same log-probability, choose the lexicographically smallest token so the result is deterministic.
Stop decoding when one of the following happens:
- you append `eos`
- the current prefix is missing from `scores` or maps to an empty dictionary
- the total sequence length reaches `max_len`
Return the full decoded sequence, including the original `prompt_tokens`.
Assume all tokens are strings with no spaces.
Constraints
- 0 <= max_len <= 200
- 0 <= len(prompt_tokens) <= 200
- 0 <= total number of next-token scores across all prefixes <= 10^4
- All tokens, including `eos`, are non-empty strings without spaces
- Larger log-probability means a more likely next token
Examples
Input: ({'': {'I': -0.1, 'You': -0.4}, 'I': {'like': -0.2, 'am': -0.3}, 'I like': {'pizza': -0.1, 'music': -0.5}, 'I like pizza': {'<EOS>': -0.05}}, [], 5, '<EOS>')
Expected Output: ['I', 'like', 'pizza', '<EOS>']
Explanation: Greedy decoding always picks the highest-scoring next token until EOS is generated.
Input: ({'Hello': {'world': -0.1}, 'Hello world': {'again': -0.2}}, ['Hello'], 2, '<EOS>')
Expected Output: ['Hello', 'world']
Explanation: The sequence stops once total length reaches `max_len`, even though more transitions exist.
Hints
- At each step, build the current prefix key with `' '.join(current_tokens)`.
- To make ties deterministic, compare by highest log-probability first and token string second.
Part 2: Beam Search with Top-k Prefixes
You are given the same simplified language model as a dictionary `scores`, where `scores[' '.join(prefix_tokens)]` returns a map of next-token log-probabilities for that exact prefix. The empty prefix uses the key `''`.
Implement beam search with beam width `k`.
Rules:
- Start with one beam: the initial `prompt_tokens` with cumulative score `0.0`
- At each step, expand every unfinished beam by one token using the map for its current prefix
- The score of a beam is the sum of generated token log-probabilities along that beam
- After all expansions for the step, keep only the top `k` beams by cumulative score
- If scores tie, keep the lexicographically smaller full sequence first
- A beam is considered finished if it ends with `eos`, reaches `max_len`, or has no outgoing transitions; finished beams remain in the beam set but are not expanded further
- Stop when all remaining beams are finished
Return the single best full sequence, including the original `prompt_tokens`.
Do not use length normalization. Assume all tokens are strings with no spaces.
Constraints
- 1 <= k <= 50
- 0 <= max_len <= 200
- 0 <= len(prompt_tokens) <= 200
- 0 <= total number of next-token scores across all prefixes <= 10^4
- All tokens, including `eos`, are non-empty strings without spaces
- Larger cumulative log-probability means a better beam
Examples
Input: ({'': {'A': -0.1, 'B': -0.2}, 'A': {'x': -2.0, '<EOS>': -10.0}, 'A x': {'<EOS>': -0.1}, 'B': {'y': -0.1}, 'B y': {'<EOS>': -0.1}}, [], 4, '<EOS>', 2)
Expected Output: ['B', 'y', '<EOS>']
Explanation: Keeping two beams allows the search to recover from the locally best first token and choose the better total sequence.
Input: ({'': {'a': -0.1, 'b': -0.2}, 'a': {'<EOS>': -0.1}, 'b': {'c': -0.05}, 'b c': {'<EOS>': -0.05}}, [], 3, '<EOS>', 2)
Expected Output: ['a', '<EOS>']
Explanation: The finished beam `['a', '<EOS>']` remains in the beam set and still wins overall.
Hints
- Each beam needs both its token list and its cumulative score.
- Finished beams should stay in the candidate set unchanged; only unfinished beams are expanded.