Test-Time Compute › Code Implementation: Multiple Sampling with Majority Vote

import random

# Simulate a model that answers math questions with a given per-question accuracy.
# In a real implementation, this would call an API or local inference endpoint.


class SimulatedModel:
    """Simulate a language model with known per-question accuracy."""

    def __init__(self, true_accuracy: float, seed: int = 42):
        self.true_accuracy = true_accuracy
        self.rng = random.Random(seed)

    def sample(self, question: str, n_samples: int = 1) -> list:
        """Generate n_samples independent answers, each correct with probability p."""
        answers = []
        for _ in range(n_samples):
            if self.rng.random() < self.true_accuracy:
                answers.append("correct")
            else:
                # Wrong answers are diverse; only one is correct.
                wrong_options = ["wrong_a", "wrong_b", "wrong_c"]
                answers.append(self.rng.choice(wrong_options))
        return answers

def majority_vote(answers: list) -> str:
    """Return the most common answer from a list of candidates."""
    counter = Counter(answers)
    return counter.most_common(1)[0][0]


def evaluate_majority_vote(
    model, n_questions: int, n_samples: int, seed: int = 0
) -> float:
    """Estimate accuracy using majority vote over n_samples per question."""
    rng = random.Random(seed)
    correct_count = 0
    for _ in range(n_questions):
        candidates = model.sample("question", n_samples)
        selected = majority_vote(candidates)
        if selected == "correct":
            correct_count += 1
    return correct_count / n_questions

def pass_at_k_exact(n: int, c: int, k: int) -> float:
    """
    Compute exact pass@k given n total samples and c correct ones.

    Formula: pass@k = 1 - C(n-c, k) / C(n, k)

    Args:
        n: Total number of generated samples.
        c: Number of correct samples among the n.
        k: Number of samples submitted (the budget).

    Returns:
        Probability that at least one of the k selected samples is correct.
    """
    if n - c < k:
        return 1.0
    return 1.0 - comb(n - c, k) / comb(n, k)


def expected_pass_at_k(
    p: float, n: int, k: int, n_questions: int = 10000
) -> float:
    """Estimate E[pass@k] over many questions, each with P(correct) = p."""
    rng = np.random.default_rng(42)
    total = 0.0
    for _ in range(n_questions):
        c = int(rng.binomial(n, p))
        total += pass_at_k_exact(n, c, k)
    return total / n_questions

def random_selection(answers: list) -> str:
    """Randomly select one answer (baseline)."""
    return random.choice(answers)


def oracle_selection(answers: list) -> str:
    """Select the correct answer if it exists (upper bound)."""
    if "correct" in answers:
        return "correct"
    return random.choice(answers)


def compare_strategies(
    model, n_questions: int = 2000, sample_sizes: list = None
) -> dict:
    """
    Compare three selection strategies across varying sample sizes.
    Returns dict mapping strategy name to list of accuracies.
    """
    if sample_sizes is None:
        sample_sizes = [1, 2, 4, 8, 16, 32, 64]

    results = {"random": [], "majority_vote": [], "oracle": []}

    for n in sample_sizes:
        counts = {"random": 0, "majority_vote": 0, "oracle": 0}
        for _ in range(n_questions):
            answers = model.sample("q", n)
            for strategy, fn in [
                ("random", random_selection),
                ("majority_vote", majority_vote),
                ("oracle", oracle_selection),
            ]:
                if fn(answers) == "correct":
                    counts[strategy] += 1
        for key in results:
            results[key].append(counts[key] / n_questions)

    return results


# Run comparison
model = SimulatedModel(true_accuracy=0.25, seed=99)
sample_sizes = [1, 2, 4, 8, 16, 32, 64]
strategy_results = compare_strategies(
    model, n_questions=2000, sample_sizes=sample_sizes
)