Beam Search Decoder

Complete a beam-search implementation for sequence generation using a provided `next_token_fn`. The function returns the best generated token sequences under beam size, max-token, and stop-token constraints. Problem Requirements The goal is to write a beam search algorithm. This is a method used by AI models (like translators or chat bots) to choose the best sentence from many possibilities. You are given a starting list of numbers (tokens), a function that predicts the probability of the next number, and some rules for the search. You need to return the best finished sequences. The interviewer gives you this empty function to fill in: from typing import List, Callable 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]]: pass Parameters: input_seq: The starting list of numbers (the prompt). We add new numbers to the end of this list. next_token_fn: A helper function. You give it a list of numbers, and it gives you a list of probabilities for what number comes next. max_token: The maximum number of new numbers to add. beam_size: How many "best guesses" (beams) to keep alive at each step. stop_word_id: The specific number that means "stop." If a sequence ends with this number, it is finished. Returns: A list of number sequences. They should be sorted by their total score (best score first). Finished sequences are preferred. If no sequences finish, return the best ones you have. Example: # Vocabulary: {0: 'hello', 1: 'world', 2: '<EOS>'} def simple_next_token(seq): return [0.2, 0.5, 0.3] result = beam_search( input_seq=[0], next_token_fn=simple_next_token, max_token=2, beam_size=2, stop_word_id=2 ) # result: [[0, 2], [0, 1, 2], [0, 0, 2]] # [0, 2]: score = log(0.3) ≈ -1.20 # [0, 1, 2]: score = log(0.5) + log(0.3) ≈ -1.90 # [0, 0, 2]: score = log(0.2) + log(0.3) ≈ -2.81 Goal: Your code must pass all the tests provided later. Part 1: How the Algorithm Works Comparison: Greedy vs. Beam vs. All Paths Strategy Paths Kept Accuracy Cost (Speed) Greedy 1 (Only the best one right now) Can miss the best overall answer O(T×V)O(T \times V)O(T×V) Beam Search BBB (The top B options) Good balance of speed and accuracy O(T×B×V)O(T \times B \times V)O(T×B×V) Exhaustive All possible paths Perfect accuracy O(VT)O(V^T)O(VT) (Impossible to run) Why Picking the Best One Now Fails Later Imagine a situation like this: Path A: Start with a score of 0.4. The next step gives a huge score of 0.9. Path B: Start with a score of 0.5. The next step gives a bad score of 0.3. A Greedy search picks Path B immediately because 0.5 is bigger than 0.4. It never sees that Path A becomes much better later. Beam Search keeps more than one option open (if beam_size ≥ 2). It keeps both A and B alive, eventually seeing that A is the winner. Step-by-Step Logic Start: Create a list containing just your starting sequence (input_seq) with a score of 0.0. Loop: Do this max_token times: Expand: Take every active sequence you have. Ask next_token_fn for probabilities. Create a new candidate for every possible next number. Score: Calculate the new score. New Score = Old Score + log(probability of new number). Filter: If a candidate ends with stop_word_id, move it to a "completed" list. Prune: Sort the remaining active candidates by score. Keep only the top beam_size candidates. Throw the rest away. Finish: Return the completed lists sorted by score. If you have none, return the active lists you have left. Why Use Log Math? We add log-probabilities instead of multiplying raw probabilities. Computer Errors: Multiplying many small numbers (like 0.1 * 0.1 * 0.1...) results in a number so small the computer treats it as zero (underflow). Math Rule: log(A×B)=log(A)+log(B)\log(A \times B) = \log(A) + \log(B)log(A×B)=log(A)+log(B). Adding is safer than multiplying. Better Score: A higher number (closer to 0, less negative) is better. Part 2: The Code Solution import math from typing import List, Callable 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]]: """ Beam search decoding for sequence generation. Args: input_seq: Initial token sequence (prompt) next_token_fn: Returns probability distribution over next tokens max_token: Maximum number of new tokens to generate beam_size: Number of beams to maintain stop_word_id: Token ID that signals end of generation Returns: Completed sequences sorted by cumulative log-probability (highest first) """ # Each beam stores: (sequence list, total_score) beams = [(input_seq[:], 0.0)] completed = [] for _ in range(max_token): all_candidates = [] for seq, score in beams: # Get probabilities for the next number probs = next_token_fn(seq) # Create a new path for every possible next number for token_id, prob in enumerate(probs): if prob <= 0: continue new_seq = seq + [token_id] new_score = score + math.log(prob) # Check if this path is finished if token_id == stop_word_id: completed.append((new_seq, new_score)) else: all_candidates.append((new_seq, new_score)) # If we have no active paths left, stop early if not all_candidates: break # Sort all new paths by score (best first) and keep only beam_size all_candidates.sort(key=lambda x: x[1], reverse=True) beams = all_candidates[:beam_size] # If nothing finished properly, return the active beams if not completed: completed = beams # Sort final results by score and remove the score number completed.sort(key=lambda x: x[1], reverse=True) return [seq for seq, _ in completed] Example Run-Through Let's trace this with input_seq=[0], beam_size=2, and stop_word_id=2. Step 1: Expand [0]. Candidate 1: [0, 1]. Score: -0.69. (Active) Candidate 2: [0, 0]. Score: -1.20. (Active) Candidate 3: [0, 2]. Score: -1.61. (Finished! Move to completed list). We keep the top 2 active beams: [0, 1] and [0, 0]. Step 2: Expand those two. From [0, 1], we get [0, 1, 2]. This ends in 2, so it is Finished. From [0, 0], we get [0, 0, 2]. This ends in 2, so it is Finished. No active beams are left. We stop. Final Result: We have three finished sequences. We sort them by score: [0, 1, 2] [0, 0, 2] [0, 2] Part 3: Testing the Code Your code must pass these tests. import math def test_immediate_stop(): """When stop word has highest probability, generate it immediately.""" def next_token_fn(seq): return [0.1, 0.1, 0.8] result = beam_search([0], next_token_fn, max_token=5, beam_size=3, stop_word_id=2) assert result[0] == [0, 2], f"Expected [0, 2], got {result[0]}" def test_greedy_search(): """beam_size=1 should behave like greedy search.""" def next_token_fn(seq): if len(seq) >= 3: return [0.1, 0.1, 0.8] # generate stop after 2 tokens return [0.1, 0.7, 0.2] # token 1 is best result = beam_search([0], next_token_fn, max_token=5, beam_size=1, stop_word_id=2) assert result[0] == [0, 1, 1, 2], f"Expected [0, 1, 1, 2], got {result[0]}" def test_max_token_limit(): """Generation stops at max_token even without stop word.""" def next_token_fn(seq): return [0.8, 0.2, 0.0] # stop word never generated result = beam_search([5], next_token_fn, max_token=3, beam_size=1, stop_word_id=2) assert len(result[0]) == 4, f"Expected length 4 (1 input + 3 generated), got {len(result[0])}" assert result[0] == [5, 0, 0, 0], f"Expected [5, 0, 0, 0], got {result[0]}" def test_multiple_completed_sequences(): """Should return multiple completed sequences when beam_size > 1.""" def next_token_fn(seq): if len(seq) >= 2: return [0.0, 0.0, 1.0] # force stop at step 2 return [0.3, 0.5, 0.2] result = beam_search([0], next_token_fn, max_token=3, beam_size=2, stop_word_id=2) assert len(result) >= 2, f"Expected at least 2 sequences, got {len(result)}" assert result[0] == [0, 1, 2], f"Expected [0, 1, 2], got {result[0]}" assert result[1] == [0, 0, 2], f"Expected [0, 0, 2], got {result[1]}" def test_beam_search_advantage(): """ Beam search finds a better completed sequence than greedy. Path A (token 0): 0.4 → STOP 0.9 → score = log(0.4)+log(0.9) ≈ -1.02 Path B (token 1): 0.5 → STOP 0.3 → score = log(0.5)+log(0.3) ≈ -1.90 Greedy picks B first, missing the globally better A→STOP path. """ def next_token_fn(seq): last = seq[-1] if last == 5: return [0.4, 0.5, 0.1] elif last == 0: return [0.05, 0.05, 0.9] elif last == 1: return [0.05, 0.65, 0.3] return [0.0, 0.0, 1.0] beam = beam_search([5], next_token_fn, max_token=2, beam_size=2, stop_word_id=2) beam_completed = [s for s in beam if s[-1] == 2] assert [5, 0, 2] in beam_completed, f"Expected [5, 0, 2] in results, got {beam_completed}" # [5, 0, 2] should be ranked higher than [5, 1, 2] if [5, 1, 2] in beam_completed: assert beam_completed.index([5, 0, 2]) < beam_completed.index([5, 1, 2]) def test_sorted_by_score(): """Results should be sorted by cumulative log-probability, highest first.""" def next_token_fn(seq): if len(seq) == 1: return [0.3, 0.5, 0.2] return [0.0, 0.0, 1.0] result = beam_search([0], next_token_fn, max_token=2, beam_size=3, stop_word_id=2) # [0, 1, 2]: log(0.5) + log(1.0) = -0.69 # [0, 0, 2]: log(0.3) + log(1.0) = -1.20 # [0, 2]: log(0.2) = -1.61 assert result == [[0, 1, 2], [0, 0, 2], [0, 2]], f"Unexpected order: {result}" # Run all tests if __name__ == "__main__": test_immediate_stop() test_greedy_search() test_max_token_limit() test_multiple_completed_sequences() test_beam_search_advantage() test_sorted_by_score() print("All tests passed!") Time and Space Performance Time Complexity: O(T×B×V×log⁡(B×V))O(T \times B \times V \times \log(B \times V))O(T×B×V×log(B×V)) TTT = max_token (How many steps we take). BBB = beam_size (How many beams we hold). VVV = Vocabulary size (How many total words exist). The log part comes from sorting the list of candidates. Space Complexity: O(B×T+C×T)O(B \times T + C \times T)O(B×T+C×T) B×TB \times TB×T to store the active beams. C×TC \times TC×T to store the completed beams (CCC is how many finished). We briefly need O(B×V)O(B \times V)O(B×V) memory to check all next possible words. Common Interview Questions 1. Multiplication vs. Addition Question: "Why not just multiply the probabilities?" Answer: Floating-point underflow. If you have a sentence with 100 words, and each has a probability of 0.1, the math is 0.11000.1^{100}0.1100. This number is so tiny that the computer rounds it to 0. Using log turns this into addition (about -230), which computers handle easily. 2. Fixing Bias Against Long Sentences Question: "Beam search tends to pick short sentences. Why, and how do you fix it?" Answer: Since we add negative numbers (log probabilities), the score always goes down. Longer sentences have more negative numbers added, so their score is worse. To fix this, we use Length Normalization. We divide the total score by the number of words in the sentence. This makes the comparison fair. 3. Speeding Up with a Heap Question: "Sorting all candidates is slow. How can we make it faster?" Answer: Instead of sorting the huge list of candidates (size B×VB \times VB×V), we can use a min-heap of size BBB. This lets us find the top items without sorting the whole list. It is much faster when the vocabulary is large. 4. Making Answers Different Question: "The beams all look the same. How do we get diverse answers?" Answer: Dissimilarity Penalty: Penalize a beam if it looks too much like the others. Top-k Sampling: Before beam search, only allow the top K words to be considered. Hamming Penalty: Punish beams that use the exact same word in the same spot as another beam. 5. Real-World Use with LLMs Question: "How does beam search work with memory in Transformers?" Answer: Transformers use a KV cache to remember past tokens. In beam search, every beam needs its own copy of this cache. This uses a lot of memory (O(B×L×H)O(B \times L \times H)O(B×L×H)), which is why using a large beam_size is very expensive for models like GPT.