Beam Search with Finished Hypotheses

Problem

Run beam search on a tiny Markov language model: transition_logits[last_token] supplies next-token logits. Start with one empty generated sequence at score 0; use start_token as its previous token. Every step expands each unfinished hypothesis with every vocabulary token. Finished hypotheses are carried forward unchanged and never expanded.

Rank the combined candidates by decreasing summed log probability, then lexicographically increasing generated token sequence. Keep at most beam_width. Stop after max_steps or when all retained hypotheses are finished. Return the best retained hypothesis as {"tokens": generated IDs including EOS if present, "score": summed log probability, "finished": boolean}. No length normalization or preference for finished hypotheses is applied; the start token is excluded.

def beam_search(transition_logits, start_token, eos_id, beam_width, max_steps):

Constraints

• transition_logits is V × V, 1 ≤ V ≤ 6; entries are finite in [-20,20].
• 0 ≤ start_token,eos_id < V; 1 ≤ beam_width ≤ 8; 0 ≤ max_steps ≤ 6.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      0,
      2
    ],
    [
      0,
      0
    ]
  ],
  0,
  1,
  2,
  3
]
Expected result: {
  "tokens": [
    1
  ],
  "score": -0.1269280110429727,
  "finished": true
}

Arguments: [
  [
    [
      0,
      0
    ],
    [
      0,
      0
    ]
  ],
  0,
  1,
  1,
  1
]
Expected result: {
  "tokens": [
    0
  ],
  "score": -0.6931471805599453,
  "finished": false
}