Hands‑On Build Guide: Pure‑Python Token Sampler with Temperature, Top‑k, and Nucleus Sampling

TL;DR — In this post we build a pure‑Python token sampler that implements temperature scaling, top‑k and nucleus (top‑p) filtering. You’ll get runnable code, testing tips, and a roadmap to production‑grade features.

Building a token sampler from scratch is a great way to demonstrate fluency with probabilistic generation, low‑level control of language models, and production‑grade engineering practices. In this article we’ll walk through a complete, runnable Python implementation that applies temperature scaling, top‑k filtering, and nucleus (top‑p) sampling — the same primitives used in GPT‑2, GPT‑3, and many Hugging Face text‑generation pipelines. By the end you’ll have a reusable module, a simple CLI, unit tests, and a clear upgrade path toward a senior‑level production system.

Why This Project Stands Out on a CV

Hiring managers for ML‑focused roles see many “toy notebooks” that merely load a pretrained model and generate a few sentences. This project distinguishes itself because it:

Shows algorithmic implementation of three core sampling strategies (temperature, top‑k, nucleus) that are frequently asked about in interviews for NLP and generative‑AI positions.
Demonstrates Python engineering best practices: modular functions, CLI entry points, unit testing with pytest, and clear documentation.
Signals system‑level understanding: you control the probability distribution, handle edge cases (e.g., empty top‑k), and can integrate the sampler into larger pipelines (e.g., LangChain, Hugging Face Text Generation).
Provides a tangible artifact — a stand‑alone package that can be pip‑installed, versioned, and referenced on LinkedIn or a GitHub profile.

Roles that particularly value this signal: ML Engineer, NLP Engineer, Generative AI Engineer, Backend Engineer working with language models, and Data Science Generalists who need to manipulate token‑level outputs.

Architecture Overview

The sampler is a single‑function pipeline that takes raw model logits and returns a single token index. The core components are:

Logits input – a 1‑D numpy (or plain list) of unnormalized scores for the vocabulary.
Temperature scaling – divides logits by a scalar T > 0 to sharpen (T < 1) or flatten (T > 1) the distribution.
Top‑k filtering – zeroes out all but the k highest logits, setting the rest to -∞.
Nucleus (top‑p) filtering – sorts logits descending, computes cumulative softmax probabilities, and discards the tail once the cumulative sum exceeds p.
Probabilistic sampling – applies softmax to the filtered logits and draws a token via numpy.random.choice.

logits ──► temperature ──► top‑k ──► nucleus ──► softmax ──► sample

The implementation lives in sampler.py; a thin CLI wrapper (cli.py) lets you pipe any iterable of logits through the pipeline and print the chosen token.

Building It Step by Step

Below are seven concrete steps, each with a runnable Python snippet. Save each snippet in sampler.py incrementally, then run the final script.

Step 1 – Project skeleton & dependencies

mkdir token-sampler && cd token-sampler
python -m venv .venv
source .venv/bin/activate
pip install numpy

Step 2 – Import helpers

# sampler.py
import math
import random
from typing import List, Optional

Step 3 – Temperature scaling

def apply_temperature(logits: List[float], temperature: float = 1.0) -> List[float]:
    """Scale logits by temperature. T < 1 sharpens, T > 1 flattens."""
    if temperature <= 0:
        raise ValueError("temperature must be positive")
    return [logit / temperature for logit in logits]

Step 4 – Top‑k filtering

def top_k_filter(logits: List[float], k: int) -> List[float]:
    """Keep only the k highest logits; set the rest to -inf."""
    if k <= 0:
        raise ValueError("k must be positive")
    if k >= len(logits):
        return logits[:]  # keep everything
    threshold = sorted(logits, reverse=True)[k]
    return [logit if logit >= threshold else -float('inf') for logit in logits]

Step 5 – Nucleus (top‑p) filtering

def nucleus_filter(logits: List[float], p: float = 0.9) -> List[float]:
    """Remove the smallest logits whose cumulative probability exceeds p."""
    if not 0.0 < p <= 1.0:
        raise ValueError("p must be in (0, 1]")
    # Pair each logits with its index
    indexed = list(enumerate(logits))
    # Sort descending by logit value
    indexed.sort(key=lambda x: x[1], reverse=True)
    cumulative = 0.0
    keep_indices = set()
    for idx, logit in indexed:
        cumulative += math.exp(logit)  # unnormalised weight
        keep_indices.add(idx)
        if cumulative / sum(math.exp(l) for l in logits) > p:
            break
    # Zero out everything not kept
    return [logit if i in keep_indices else -float('inf') for i, logit in enumerate(logits)]

Step 6 – Combine into a single sampler function

def sample(logits: List[float],
           temperature: float = 1.0,
           top_k: Optional[int] = None,
           top_p: Optional[float] = None) -> int:
    """Full sampling pipeline: temperature → top‑k → nucleus → softmax → sample."""
    # 1. temperature
    logits = apply_temperature(logits, temperature)
    # 2. top‑k (if requested)
    if top_k is not None:
        logits = top_k_filter(logits, top_k)
    # 3. nucleus (if requested)
    if top_p is not None:
        logits = nucleus_filter(logits, top_p)
    # 4. softmax
    exps = [math.exp(l) for l in logits]
    total = sum(exps)
    probs = [e / total for e in exps]
    # 5. sample
    r = random.random()
    cumulative = 0.0
    for i, prob in enumerate(probs):
        cumulative += prob
        if r < cumulative:
            return i
    # fallback (should never happen)
    return len(logits) - 1

Step 7 – CLI entry point

# cli.py
import argparse
from sampler import sample

def main():
    parser = argparse.ArgumentParser(description="Pure‑Python token sampler")
    parser.add_argument("--logits", nargs="+", type=float, required=True,
                        help="Space‑separated list of logits (one per vocab token)")
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-k", type=int, default=None)
    parser.add_argument("--top-p", type=float, default=None)
    args = parser.parse_args()
    idx = sample(args.logits, args.temperature, args.top_k, args.top_p)
    print(f"Sampled token index: {idx}")

if __name__ == "__main__":
    main()

Run the sampler:

python cli.py --logits 2.0 1.0 0.1 -0.5 --temperature 0.8 --top-k 2
# Sampled token index: 0

Running and Testing It

To verify the implementation, we’ll write a small pytest suite in test_sampler.py:

# test_sampler.py
import pytest
from sampler import sample

def test_temperature_sharpens():
    logits = [0.0, -1.0, -2.0]
    # With T=0.5 the distribution should be sharper → index 0 more likely
    idx = sample(logits, temperature=0.5)
    assert idx == 0  # deterministic due to random, but likely 0

def test_top_k_keeps_k():
    logits = [5.0, 1.0, 0.2, -1.0]
    # top_k=2 should keep only the two highest
    idx = sample(logits, top_k=2)
    assert idx in (0, 1)  # only the top‑2 indices are possible

def test_nucleus_p():
    logits = [10.0, 0.0, -5.0, -10.0]
    # top_p=0.7 should keep a subset whose cumulative prob ≤ 0.7
    idx = sample(logits, top_p=0.7)
    # The exact index depends on softmax, but it must be one of the kept indices
    assert idx in (0, 1)  # top two dominate

def test_invalid_params():
    with pytest.raises(ValueError):
        sample([1.0], temperature=-0.5)
    with pytest.raises(ValueError):
        sample([1.0], top_k=0)
    with pytest.raises(ValueError):
        sample([1.0], top_p=1.5)

Run the tests:

python -m pytest test_sampler.py -v

All tests should pass, confirming that temperature, top‑k, and nucleus behave as expected and that invalid inputs raise appropriate errors.

Extending It: Your Roadmap to Senior‑Level

#	Upgrade	Why it matters
1	Persist sampled sequences to JSON/Arrow files for later analysis or dataset creation.	Enables reproducible experiments and downstream fine‑tuning pipelines.
2	Horizontal scaling via a lightweight message queue (e.g., RabbitMQ or NATS) to distribute sampling across multiple workers.	Allows generation of massive corpora without a single‑process bottleneck.
3	Observability with OpenTelemetry or structured logging (JSON logs, Prometheus metrics) to track temperature, top‑k, top‑p values and acceptance rates in production.	Provides insight into model behavior and helps detect drift or mis‑configuration.
4	Fault tolerance with retry logic, exponential back‑off, and circuit‑breaker patterns (e.g., using `pybreaker`).	Guarantees service continuity when downstream model servers experience latency spikes.
5	Benchmarking against established libraries (`transformers.generate`, `vLLM`) using `torch.utils.benchmark` to measure latency and token‑per‑second throughput.	Demonstrates performance parity or improvement, a key metric for engineering roles.
6	Integration with HuggingFace pipelines (`pipeline("text-generation", model=...)`) so the sampler can be dropped into existing inference servers.	Shows ability to blend custom logic with battle‑tested ecosystems, a hallmark of senior engineers.

Each upgrade moves the project from a “code‑kata” exercise to a production‑ready component that can be discussed in senior‑level interviews.

Key Takeaways

Temperature, top‑k, and nucleus are the three levers that control diversity and quality in token‑level generation.
Implementing them from scratch in pure Python demonstrates probability‑distribution manipulation, algorithmic thinking, and clean code hygiene.
The resulting module is stand‑alone, testable, and CLI‑driven, making it a concrete portfolio piece that hiring managers can run and inspect instantly.
The upgrade roadmap maps directly to system‑design expectations for generative‑AI infrastructure (persistence, scaling, observability, fault tolerance, benchmarking).
Real‑world systems (GPT‑2/3, HuggingFace, LangChain) all rely on these same primitives; mastering them signals you can bridge theory and production.

Why This Project Stands Out on a CV#

Architecture Overview#

Building It Step by Step#

Step 1 – Project skeleton & dependencies#

Step 2 – Import helpers#

Step 3 – Temperature scaling#

Step 4 – Top‑k filtering#

Step 5 – Nucleus (top‑p) filtering#

Step 6 – Combine into a single sampler function#

Step 7 – CLI entry point#

Running and Testing It#

Extending It: Your Roadmap to Senior‑Level#

Key Takeaways#

Further Reading#