Beam Search › Implementation with Priority Queues

from typing import List

import numpy as np

# A simple toy language model using random but reproducible probabilities
np.random.seed(42)

VOCAB = [
    "<EOS>",
    "the",
    "cat",
    "sat",
    "on",
    "mat",
    "dog",
    "ran",
    "away",
    "quickly",
]
VOCAB_SIZE = len(VOCAB)
EOS_ID = 0


def toy_lm_step(prev_tokens: List[int]) -> np.ndarray:
    """
    Toy language model: returns log-probabilities over the vocabulary.
    In a real system, this would be a transformer or RNN forward pass.
    """
    rng = np.random.RandomState(seed=sum(prev_tokens) % 2**31)
    logits = rng.randn(VOCAB_SIZE)
    # Suppress EOS early, encourage it later
    if len(prev_tokens) < 3:
        logits[EOS_ID] -= 5.0
    elif len(prev_tokens) >= 6:
        logits[EOS_ID] += 3.0
    # Softmax
    logits -= logits.max()
    probs = np.exp(logits)
    probs /= probs.sum()
    return np.log(probs)

class Hypothesis:
    """A beam search hypothesis: token sequence + cumulative log-prob."""

    def __init__(self, tokens: List[int], score: float):
        self.tokens = tokens
        self.score = score

    def __lt__(self, other: "Hypothesis") -> bool:
        # Sort descending by score: higher score is "less than" for min-heap inversion
        return self.score > other.score

    def __eq__(self, other: "Hypothesis") -> bool:
        return self.score == other.score

from typing import List, Tuple


def length_penalty(length: int, alpha: float = 0.6) -> float:
    """Google's length penalty from Wu et al. 2016."""
    return ((5 + length) ** alpha) / (6**alpha)


def beam_search(
    start_token: int,
    beam_width: int = 3,
    max_length: int = 8,
    alpha: float = 0.6,
) -> List[Tuple[List[int], float]]:
    """
    Run beam search and return the top completed hypotheses.
    Returns list of (token_sequence, normalized_score) tuples.
    """
    # Initialize beam with a single hypothesis containing the start token
    beam = [Hypothesis(tokens=[start_token], score=0.0)]
    completed = []

    for step in range(max_length):
        candidates = []
        for hyp in beam:
            # Get log-probs from the language model
            log_probs = toy_lm_step(hyp.tokens)
            # Expand this hypothesis by every vocabulary token
            for token_id, log_prob in enumerate(log_probs):
                new_tokens = hyp.tokens + [token_id]
                new_score = hyp.score + log_prob
                candidates.append(Hypothesis(new_tokens, new_score))

        # Sort all candidates, keep top beam_width
        candidates.sort()  # Sorts by neg_score ascending = score descending
        beam = candidates[:beam_width]

        # Move completed hypotheses (ending in EOS) to the completed list
        still_active = []
        for hyp in beam:
            if hyp.tokens[-1] == EOS_ID:
                # Apply length normalization (exclude start token from length)
                seq_len = len(hyp.tokens) - 1
                normalized = hyp.score / length_penalty(seq_len, alpha)
                completed.append((hyp.tokens, normalized))
            else:
                still_active.append(hyp)

        beam = still_active

        if not beam:
            break

    # Also collect remaining beam hypotheses with length normalization
    for hyp in beam:
        seq_len = len(hyp.tokens) - 1
        normalized = hyp.score / length_penalty(seq_len, alpha)
        completed.append((hyp.tokens, normalized))

    # Sort by normalized score (descending)
    completed.sort(key=lambda x: x[1], reverse=True)
    return completed