Neural Mastery
← Practice
Beam Search Step with Pruningmedium

Beam Search Step with Pruning

medium Inside an LLM 20–25 min

NumPy allowed · plain Python earns 10 bonus points

LLM step 1/4

Logits become a next-token distribution

Softmax turns a model score vector into probabilities.
Open lesson →
01prompt
02tokens
03logits
04probabilities
attention→next token
mechanism65%
layer20%
pattern10%
model5%

Softmax(logits) → probability distribution

Explore multiple high-probability generation hypotheses simultaneously using Beam Search.

Task

Implement beam_search_step(current_beams, candidate_log_probs, beam_width, eos_token_id). For each beam in current_beams: if is_done is True, carry it forward unchanged; otherwise expand it across all vocabulary tokens $v$, updating score = round(beam['score'] + log_prob[v], 4) and setting is_done = True if $v == ext{eos_token_id}$. Sort all candidate hypotheses descending by score (breaking ties lexicographically by sequence). Prune to beam_width entries. Return dict with beams and all_done (True if every pruned beam is done). Raise ValueError if beam_width <= 0 or if input list lengths mismatch.

Function Signature

beam_search_step(current_beams: list[dict], candidate_log_probs: list[list[float]], beam_width: int, eos_token_id: int) -> dict

Examples

Example 1: Single step beam expansion
Input: current_beams = [{'sequence': [10], 'score': -0.5, 'is_done': False}] candidate_log_probs = [[-0.2, -1.5, -0.8]] beam_width = 2 eos_token_id = 99
Output: {'beams': [{'sequence': [10, 0], 'score': -0.7, 'is_done': False}, {'sequence': [10, 2], 'score': -1.3, 'is_done': False}], 'all_done': False}
Example 2: Completed beam remains preserved
Input: current_beams = [{'sequence': [10, 99], 'score': -0.7, 'is_done': True}, {'sequence': [10, 2], 'score': -1.3, 'is_done': False}] candidate_log_probs = [[], [-0.1, -2]] beam_width = 2 eos_token_id = 99
Output: {'beams': [{'sequence': [10, 99], 'score': -0.7, 'is_done': True}, {'sequence': [10, 2, 0], 'score': -1.4, 'is_done': False}], 'all_done': False}
Example 3: Invalid beam_width <= 0 error
Input: current_beams = [] candidate_log_probs = [] beam_width = 0 eos_token_id = 1
Output: Raises ValueError

Constraints

  • beam_width <= 0 must raise ValueError.
  • Scores rounded to 4 decimal places.
  • Completed beams must not be re-expanded.
  • A plain-Python solution earns 10 bonus points.
Python3Saved ✓
def beam_search_step(current_beams: list[dict], candidate_log_probs: list[list[float]], beam_width: int, eos_token_id: int) -> dict:
"""Expand active beams by candidate log probabilities and prune to top-B.
Args:
current_beams: List of dicts with 'sequence', 'score', and 'is_done'.
candidate_log_probs: Vocabulary log probabilities for each beam.
beam_width: Number of active hypotheses to retain (B).
eos_token_id: End of sequence token ID.
Returns:
dict with 'beams' (top beam_width dicts) and 'all_done' (bool).
"""
# Your implementation here
pass

Case 1: Single step beam expansion
Input: current_beams = [{'sequence': [10], 'score': -0.5, 'is_done': False}] candidate_log_probs = [[-0.2, -1.5, -0.8]] beam_width = 2 eos_token_id = 99
Expected: {'beams': [{'sequence': [10, 0], 'score': -0.7, 'is_done': False}, {'sequence': [10, 2], 'score': -1.3, 'is_done': False}], 'all_done': False}
Case 2: Completed beam remains preserved
Input: current_beams = [{'sequence': [10, 99], 'score': -0.7, 'is_done': True}, {'sequence': [10, 2], 'score': -1.3, 'is_done': False}] candidate_log_probs = [[], [-0.1, -2]] beam_width = 2 eos_token_id = 99
Expected: {'beams': [{'sequence': [10, 99], 'score': -0.7, 'is_done': True}, {'sequence': [10, 2, 0], 'score': -1.4, 'is_done': False}], 'all_done': False}
Case 3: Invalid beam_width <= 0 error
Input: current_beams = [] candidate_log_probs = [] beam_width = 0 eos_token_id = 1
Expected: Raises ValueError