TL;DR — You will build a compact, pure‑Python LLM decoding engine that implements temperature scaling, top‑k filtering, top‑p nucleus sampling, and repetition penalty. The resulting script runs on a CPU, produces coherent text from a small model, and demonstrates concrete systems and algorithmic skills that hiring managers look for.

Large language models are only as good as the way you sample their next token. While many tutorials stop at greedy decoding, a production‑grade system needs to control randomness, avoid repetitive loops, and balance diversity with coherence. In this guide you will implement a complete decoding pipeline in pure Python, covering temperature scaling, top‑k filtering, top‑p nucleus sampling, and repetition penalty. The code is self‑contained, runs on a single CPU, and can be dropped into any project that needs controlled generation.

Why This Project Stands Out on a CV

  • Algorithmic depth – You will implement four distinct sampling strategies, each with its own mathematical justification, showing you can translate research into code.
  • Systems thinking – The engine is structured as a pipeline of logits processors, a sampler, and a token selector, mirroring the architecture of production inference services like vLLM or TensorRT‑LLM.
  • Performance awareness – By keeping the implementation in pure Python with optional NumPy vectorization, you demonstrate an understanding of computational bottlenecks and the trade‑offs between speed and flexibility.
  • Testing & validation – The guide includes concrete tests that verify sampling distributions and repetition control, evidencing a habit of shipping reliable software.
  • Portfolio‑ready – The project can be extended with batching, KV‑caching, or GPU acceleration, giving you talking points for senior‑level roles that demand scalability and observability.

Architecture Overview

The decoding engine is composed of four loosely coupled components:

  1. Logits Processor – Receives raw logits from the language model and applies transformations (temperature, top‑k, top‑p, repetition penalty).
  2. Sampler – Converts processed logits into a probability distribution and draws the next token.
  3. Token Selector – Chooses the final token (argmax, stochastic, or hybrid) and updates the context.
  4. Context Manager – Maintains the sequence of generated tokens, provides them to the model for the next forward pass, and enforces repetition penalty by tracking token frequencies.

A simple text diagram of the flow:

[Language Model] → raw logits → [Logits Processor] → processed logits → [Sampler] → next token → [Token Selector] → updated context → (loop)

Each component is a pure Python class, making it easy to swap implementations or add new strategies.

Building It Step by Step

Step 1 – Set up the environment

python -m venv venv
source venv/bin/activate
pip install torch transformers numpy

Step 2 – Load a small language model and tokenizer

We use Hugging Face Transformers to load a lightweight model; any causal LM works.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.eval()

Step 3 – Implement logits processing functions

Create a module sampling.py with the following functions.

import torch
import torch.nn.functional as F

def temperature_logits(logits: torch.Tensor, temperature: float) -> torch.Tensor:
    """Scale logits by temperature."""
    return logits / temperature

def top_k_filter(logits: torch.Tensor, k: int) -> torch.Tensor:
    """Zero out logits below the k‑th largest value."""
    if k <= 0:
        return logits
    values, _ = torch.topk(logits, k)
    min_value = values[..., -1, None]
    logits = torch.where(logits < min_value, torch.full_like(logits, -float('Inf')), logits)
    return logits

def top_p_filter(logits: torch.Tensor, p: float) -> torch.Tensor:
    """Nucleus sampling: keep smallest set whose cumulative probability ≥ p."""
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
    # Remove tokens with cumulative probability above the threshold
    sorted_indices_to_remove = cumulative_probs > p
    # Shift the mask so that at least one token is kept
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = False
    # Scatter back to original ordering
    indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)
    logits = logits.masked_fill(indices_to_remove, -float('Inf'))
    return logits

def repetition_penalty_logits(logits: torch.Tensor, past_tokens: list[int], penalty: float) -> torch.Tensor:
    """Apply repetition penalty to logits based on already generated tokens."""
    if not past_tokens:
        return logits
    # Create a tensor of penalties
    penalty_tensor = torch.ones_like(logits)
    for token_id in set(past_tokens):
        penalty_tensor[..., token_id] = penalty
    # Apply penalty: divide logits by penalty for repeated tokens
    logits = logits / penalty_tensor
    return logits

Step 4 – Combine into a decode function

from sampling import temperature_logits, top_k_filter, top_p_filter, repetition_penalty_logits

@torch.no_grad()
def decode(model, tokenizer, prompt: str, max_new_tokens: int = 50,
           temperature: float = 1.0, top_k: int = 50, top_p: float = 0.9,
           repetition_penalty: float = 1.2) -> str:
    input_ids = tokenizer.encode(prompt, return_tensors='pt')
    past_tokens = []
    for _ in range(max_new_tokens):
        outputs = model(input_ids)
        logits = outputs.logits[:, -1, :]  # last token logits
        # Apply logits processing
        logits = temperature_logits(logits, temperature)
        logits = top_k_filter(logits, top_k)
        logits = top_p_filter(logits, top_p)
        logits = repetition_penalty_logits(logits, past_tokens, repetition_penalty)
        # Sample
        probs = F.softmax(logits, dim=-1)
        next_token_id = torch.multinomial(probs, num_samples=1).item()
        # Append to sequence
        input_ids = torch.cat([input_ids, torch.tensor([[next_token_id]])], dim=-1)
        past_tokens.append(next_token_id)
    return tokenizer.decode(input_ids[0], skip_special_tokens=True)

Step 5 – Generate text

if __name__ == "__main__":
    prompt = "The future of artificial intelligence is"
    generated = decode(model, tokenizer, prompt,
                       temperature=0.7,
                       top_k=40,
                       top_p=0.9,
                       repetition_penalty=1.2,
                       max_new_tokens=100)
    print(generated)

Running this script produces a paragraph that balances creativity (via temperature and nucleus sampling) with coherence (via top‑k and repetition penalty).

Running and Testing It

  1. Execute the script

    python generate.py

  2. Verify output – The generated text should be diverse yet avoid obvious repetition. You can quantify repetition by counting unique tokens or using a simple script:

    text = "..."  # generated output
tokens = text.split()
unique_ratio = len(set(tokens)) / len(tokens)
print(f"Unique token ratio: {unique_ratio:.2f}")

  3. Unit tests – Add tests for each sampling function using pytest. For example, ensure top_k_filter zeros out exactly the correct number of logits.

  4. Benchmark – Time the generation loop with time.perf_counter() to measure tokens‑per‑second, providing a baseline for later optimizations.

Extending It: Your Roadmap to Senior-Level

  • Batched inference – Process multiple prompts simultaneously by padding input sequences; this improves GPU utilization and throughput, a key concern in production serving.
  • KV caching – Cache key‑value states from previous forward passes to avoid recomputation, dramatically reducing latency for autoregressive generation.
  • GPU acceleration – Move the model and tensors to cuda and enable mixed‑precision (torch.cuda.amp) to achieve real‑time generation speeds.
  • Persistence & caching – Store generated completions in a SQLite database or Redis cache to serve repeated requests without re‑running the model.
  • Observability – Integrate with Prometheus to expose metrics such as latency, token count, and sampling probabilities, enabling capacity planning and anomaly detection.
  • Fault tolerance – Wrap the generation loop in retry logic with exponential backoff and fallback to a smaller model if the primary one fails, ensuring high availability.

Key Takeaways

  • You now have a pure‑Python LLM decoding engine that implements temperature, top‑k, top‑p, and repetition penalty.
  • The modular design mirrors production inference systems, showcasing your ability to build scalable pipelines.
  • The project provides concrete evidence of algorithmic understanding and software engineering rigor, making it a strong talking point on a CV.
  • By following the extension roadmap, you can evolve the prototype into a production‑grade service with batching, caching, and observability.

Further Reading