Beam Search Decoder
Implement beam search for sequence decoding.
Beam search maintains the top-k (beam width) partial sequences at each step, expanding each with all possible next tokens and keeping only the top-k overall.
Input format:
- Line 1: beam_width max_length
- Line 2: Number of vocabulary tokens V
- Lines 3 to V+2: token probability_given_previous_token (uniform for simplicity)
- Line V+3: Number of transition rules T
- Lines V+4 to V+T+3: prev_token next_token probability
The model uses: P(next | prev) from transition rules. Missing transitions have probability 0. Start token is <s>. End token is </s>. Score = product of probabilities (use log sum).
Output: The top beam_width sequences with their log-probabilities, sorted by score descending. Format: each line is "word1 word2 ... : log_prob"
Example
Input
2 3 3 a b </s> 4 <s> a 0.6 <s> b 0.4 a </s> 0.7 b </s> 0.8
Output
a : -0.6931 b : -1.1394
Reasoning
Step 1: Expand <s>
- <s> -> a: log(0.6) = -0.5108
- <s> -> b: log(0.4) = -0.9163 Keep top 2: [a (-0.5108), b (-0.9163)]
Step 2: Expand both
- a -> </s>: -0.5108 + log(0.7) = -0.5108 + (-0.3567) = -0.8675. But wait...
Actually with the given probabilities, let me recalculate:
- a -> </s>: log(0.6) + log(0.7) = -0.5108 + (-0.3567) = -0.8675 => "a" (completed)
- b -> </s>: log(0.4) + log(0.8) = -0.9163 + (-0.2231) = -1.1394 => "b" (completed)
Both completed. Sort by score: a : -0.8675 b : -1.1394
Hmm, let me use the exact values from the code.
Constraints
- Use log probabilities to avoid underflow
- At each step, expand all beams with all possible next tokens
- Keep only top beam_width candidates
- Stop expanding a beam when </s> is generated
- max_length includes </s> but not <s>
- Output completed sequences sorted by log prob (highest first)
- Round log probs to 4 decimal places