Implement greedy decoding and beam search

You are given access to an auto-regressive language model that can score the next token based on a prefix.

Assume:

• next_logprobs(prefix_tokens) -> Map[token, logp] returns log-probabilities (natural log) for the next token given the current prefix.
• EOS is a special end-of-sequence token.
• You are also given max_len and an initial prefix prompt_tokens.

Implement two decoding methods:

1. Greedy decoding (greedy sampling)

• Repeatedly select the single token with the highest log-probability at each step.
• Stop when you generate EOS or when the total length reaches max_len.
• Return the generated token sequence (including the prompt, or clearly specify whether you return only newly generated tokens).

2. Beam search (keep top-k paths)

• Given an integer k (the number of beams to keep), implement beam search that maintains the top-k candidate sequences by total sequence score.
• Use cumulative log-probability (sum of token log-probs) as the sequence score (no need for length normalization unless you choose to add it and explain).
• At each decoding step, expand each current beam by one token using next_logprobs, then keep only the top-k resulting sequences overall.
• Stop when all beams have generated EOS or when reaching max_len.
• Return the best sequence (and optionally the top-k sequences if requested).

Clarify any edge cases you handle (e.g., early-finished beams, ties, empty vocabulary, whether finished beams remain in the beam set without further expansion).

Part 1: Greedy Decoding from Prefix Log-Probabilities

You are given a simplified auto-regressive language model as a dictionary called scores. For any current token list prefix_tokens, the available next-token log-probabilities are stored in scores[' '.join(prefix_tokens)]. The empty prefix uses the key ''.

Implement greedy decoding. Starting from prompt_tokens, repeatedly append the token with the highest log-probability for the current prefix. If multiple tokens have the same log-probability, choose the lexicographically smallest token so the result is deterministic.

Stop decoding when one of the following happens:

• you append eos
• the current prefix is missing from scores or maps to an empty dictionary
• the total sequence length reaches max_len

Return the full decoded sequence, including the original prompt_tokens.

Assume all tokens are strings with no spaces.

Constraints

• 0 <= max_len <= 200
• 0 <= len(prompt_tokens) <= 200
• 0 <= total number of next-token scores across all prefixes <= 10^4
• All tokens, including eos, are non-empty strings without spaces
• Larger log-probability means a more likely next token

Examples

Example 1

Input: ({'': {'I': -0.1, 'You': -0.4}, 'I': {'like': -0.2, 'am': -0.3}, 'I like': {'pizza': -0.1, 'music': -0.5}, 'I like pizza': {'<EOS>': -0.05}}, [], 5, '<EOS>')
Output: ['I', 'like', 'pizza', '<EOS>']

Explanation: Greedy decoding always picks the highest-scoring next token until EOS is generated.

Example 2

Input: ({'Hello': {'world': -0.1}, 'Hello world': {'again': -0.2}}, ['Hello'], 2, '<EOS>')
Output: ['Hello', 'world']

Explanation: The sequence stops once total length reaches max_len, even though more transitions exist.

Hints

1. At each step, build the current prefix key with ' '.join(current_tokens).
2. To make ties deterministic, compare by highest log-probability first and token string second.

Part 2: Beam Search with Top-k Prefixes

You are given the same simplified language model as a dictionary scores, where scores[' '.join(prefix_tokens)] returns a map of next-token log-probabilities for that exact prefix. The empty prefix uses the key ''.

Implement beam search with beam width k.

Rules:

• Start with one beam: the initial prompt_tokens with cumulative score 0.0
• At each step, expand every unfinished beam by one token using the map for its current prefix
• The score of a beam is the sum of generated token log-probabilities along that beam
• After all expansions for the step, keep only the top k beams by cumulative score
• If scores tie, keep the lexicographically smaller full sequence first
• A beam is considered finished if it ends with eos, reaches max_len, or has no outgoing transitions; finished beams remain in the beam set but are not expanded further
• Stop when all remaining beams are finished

Return the single best full sequence, including the original prompt_tokens.

Do not use length normalization. Assume all tokens are strings with no spaces.

Constraints

• 1 <= k <= 50
• 0 <= max_len <= 200
• 0 <= len(prompt_tokens) <= 200
• 0 <= total number of next-token scores across all prefixes <= 10^4
• All tokens, including eos, are non-empty strings without spaces
• Larger cumulative log-probability means a better beam

Examples

Example 1

Input: ({'': {'A': -0.1, 'B': -0.2}, 'A': {'x': -2.0, '<EOS>': -10.0}, 'A x': {'<EOS>': -0.1}, 'B': {'y': -0.1}, 'B y': {'<EOS>': -0.1}}, [], 4, '<EOS>', 2)
Output: ['B', 'y', '<EOS>']

Explanation: Keeping two beams allows the search to recover from the locally best first token and choose the better total sequence.

Example 2

Input: ({'': {'a': -0.1, 'b': -0.2}, 'a': {'<EOS>': -0.1}, 'b': {'c': -0.05}, 'b c': {'<EOS>': -0.05}}, [], 3, '<EOS>', 2)
Output: ['a', '<EOS>']

Explanation: The finished beam ['a', '<EOS>'] remains in the beam set and still wins overall.

Hints

1. Each beam needs both its token list and its cumulative score.
2. Finished beams should stay in the candidate set unchanged; only unfinished beams are expanded.