PIXELBANKv9.1.0
Menu

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
🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Test Results

0/0
Run code to see test results.