TorchedUp
LearnBetaProblemsSystem DesignSoonPremium
TorchedUp
LearnBetaProblemsSystem DesignSoonPremium
←

177. Beam Search Step

Hard

Implement a single expansion step of beam search for autoregressive decoding.

Signature: def beam_search_step(beams: list, log_probs: np.ndarray, beam_size: int) -> list

  • beams: a list of dicts {'tokens': list[int], 'score': float}
  • log_probs: shape (len(beams), vocab_size) — next-token log probabilities for each beam
  • beam_size: how many beams to keep

For each existing beam i and each vocab token v:

  • new_score = beams[i]['score'] + log_probs[i, v]
  • new_tokens = beams[i]['tokens'] + [v]

Return the top beam_size candidates ranked by new_score (descending). Break ties by lower beam index, then lower token id.

Math

si,v′​=si​+logp(v∣beami​),keep top-k by s′

Asked at

NumPy

import numpy as np

 

def beam_search_step(...):

    pass

🔒

Premium problem

Free accounts include problems #1–20. Upgrade to unlock the editor, hidden test cases, and reference solutions for every problem.

Upgrade to PremiumBack to problems

Already premium?