Sampling strategies (greedy, top-k, top-p).

export type Candidate = { token: string; logit: number };
export type Strategy = 'greedy' | 'top-k' | 'top-p';
export type RankedCandidate = { token: string; probability: number; retained: boolean };

export const exampleCandidates: Candidate[] = [
	{ token: 'the', logit: 4 },
	{ token: 'a', logit: 3 },
	{ token: 'our', logit: 2 },
	{ token: 'this', logit: 1 },
	{ token: 'one', logit: 0 }
];

export function distribution(
	candidates: Candidate[],
	strategy: Strategy,
	options: { temperature: number; k: number; p: number }
): RankedCandidate[] {
	const { temperature, k, p } = options;
	if (!Number.isFinite(temperature) || temperature <= 0)
		throw new RangeError('temperature must be > 0');
	if (!Number.isInteger(k) || k < 1) throw new RangeError('k must be an integer >= 1');
	if (!Number.isFinite(p) || p <= 0 || p > 1) throw new RangeError('p must be in (0, 1]');
	if (candidates.length === 0) return [];
	if (candidates.some((candidate) => !Number.isFinite(candidate.logit))) {
		throw new TypeError('candidate logits must be finite');
	}

	const sorted = candidates
		.map((candidate, index) => ({ ...candidate, index }))
		.sort((a, b) => b.logit - a.logit || a.index - b.index);
	const maxLogit = sorted[0].logit;
	const weights = sorted.map((candidate) => Math.exp((candidate.logit - maxLogit) / temperature));
	const total = weights.reduce((sum, weight) => sum + weight, 0);
	const probabilities = weights.map((weight) => weight / total);
	let retainedCount = sorted.length;
	if (strategy === 'greedy') retainedCount = 1;
	if (strategy === 'top-k') retainedCount = Math.min(k, sorted.length);
	if (strategy === 'top-p') {
		let cumulative = 0;
		retainedCount = 0;
		for (const probability of probabilities) {
			cumulative += probability;
			retainedCount += 1;
			if (cumulative >= p) break;
		}
	}
	const retainedWeight = weights.slice(0, retainedCount).reduce((sum, weight) => sum + weight, 0);
	return sorted.map((candidate, index) => ({
		token: candidate.token,
		probability: index < retainedCount ? weights[index] / retainedWeight : 0,
		retained: index < retainedCount
	}));
}

// A seeded sampler is useful for a reproducible teaching trace, not for secrets or production sampling.
export function sampleWithSeed(candidates: RankedCandidate[], seed: number): string {
	if (!Number.isInteger(seed) || seed < 0 || seed > 0xffffffff)
		throw new RangeError('seed must be uint32');
	const total = candidates.reduce((sum, candidate) => sum + candidate.probability, 0);
	if (total <= 0) throw new RangeError('distribution must have positive mass');
	const random = ((Math.imul(seed >>> 0, 1664525) + 1013904223) >>> 0) / 0x100000000;
	let cumulative = 0;
	for (const candidate of candidates) {
		cumulative += candidate.probability / total;
		if (random < cumulative) return candidate.token;
	}
	return candidates.at(-1)!.token;
}

package main

import (
	"errors"
	"fmt"
	"math"
	"sort"
)

type Candidate struct {
	Token string
	Logit float64
}

type RankedCandidate struct {
	Token       string
	Probability float64
	Retained    bool
}

func Distribution(candidates []Candidate, strategy string, temperature float64, k int, p float64) ([]RankedCandidate, error) {
	if math.IsNaN(temperature) || math.IsInf(temperature, 0) || temperature <= 0 {
		return nil, errors.New("temperature must be finite and greater than zero")
	}
	if k < 1 {
		return nil, errors.New("k must be at least one")
	}
	if math.IsNaN(p) || p <= 0 || p > 1 {
		return nil, errors.New("p must be in (0, 1]")
	}
	if strategy != "greedy" && strategy != "top-k" && strategy != "top-p" {
		return nil, errors.New("strategy must be greedy, top-k, or top-p")
	}
	for _, candidate := range candidates {
		if math.IsNaN(candidate.Logit) || math.IsInf(candidate.Logit, 0) {
			return nil, errors.New("candidate logits must be finite")
		}
	}
	sorted := append([]Candidate(nil), candidates...)
	sort.SliceStable(sorted, func(i, j int) bool { return sorted[i].Logit > sorted[j].Logit })
	if len(sorted) == 0 {
		return []RankedCandidate{}, nil
	}
	maxLogit := sorted[0].Logit
	weights := make([]float64, len(sorted))
	var total float64
	for i, candidate := range sorted {
		weights[i] = math.Exp((candidate.Logit - maxLogit) / temperature)
		total += weights[i]
	}
	retained := len(sorted)
	switch strategy {
	case "greedy":
		retained = 1
	case "top-k":
		if k < retained {
			retained = k
		}
	case "top-p":
		var cumulative float64
		retained = 0
		for _, weight := range weights {
			cumulative += weight / total
			retained++
			if cumulative >= p {
				break
			}
		}
	}
	var keptWeight float64
	for _, weight := range weights[:retained] {
		keptWeight += weight
	}
	result := make([]RankedCandidate, len(sorted))
	for i, candidate := range sorted {
		result[i] = RankedCandidate{Token: candidate.Token, Retained: i < retained}
		if i < retained {
			result[i].Probability = weights[i] / keptWeight
		}
	}
	return result, nil
}

func main() {
	candidates := []Candidate{{"the", 4}, {"a", 3}, {"our", 2}, {"this", 1}, {"one", 0}}
	for _, strategy := range []string{"greedy", "top-k", "top-p"} {
		distribution, err := Distribution(candidates, strategy, 1, 3, 0.8)
		if err != nil {
			panic(err)
		}
		fmt.Println(strategy)
		for _, candidate := range distribution {
			fmt.Printf("  %-5s retained=%-5t p=%.3f\n", candidate.Token, candidate.Retained, candidate.Probability)
		}
	}
}