Your question is Beam Search for Sequence Generation. Start with the requirements on the right.
Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.
Implement beam search decoding for a simplified OpenAI model-style sequence generator. At each position, retain only the beam_width highest-scoring partial sequences, expand them with the model's next-token scores, and return the best completed sequence.
The model scores are supplied as model_scores, where model_scores[t] maps a prefix to a dictionary of possible next tokens and their log-probabilities. A prefix is represented by tokens joined with |. The initial prefix contains only start_token.
Write beam_search(start_token, end_token, beam_width, max_length, model_scores). Return a list of tokens containing the start token and, when generated, the end token. A sequence ending in end_token is complete and must not be expanded further. If no sequence reaches the end token, return the highest-scoring partial sequence after max_length total tokens.
Use cumulative log-probability as the sequence score. Larger scores are better. If fewer than beam_width next tokens are available, expand all available tokens. Every required prefix has an entry in model_scores.
def beam_search(start_token, end_token, beam_width, max_length, model_scores):