# How to Perform Inference and Generate Text with a Trained LLM: A Complete Guide

> Easily perform inference and generate text with a trained LLM. Learn how to load your model and use generate_reply for seamless text generation with custom controls.

- Repository: [Fareed Khan/train-llm-from-scratch](https://github.com/FareedKhan-dev/train-llm-from-scratch)
- Tags: tutorial
- Published: 2026-06-11

---

**To perform inference and generate text with a trained LLM, load your checkpoint using `load_model_from_ckpt` and call `generate_reply` with your prompt, which handles tokenization, sampling with temperature and top-k/p controls, and decoding back to text.**

The train-llm-from-scratch repository provides a streamlined inference API that works with any checkpoint produced by its training pipelines, including base pre-training, supervised fine-tuning (SFT), and RLHF stages like DPO, PPO, or GRPO. Whether you need programmatic access or a quick command-line test, the library abstracts away the complexity of checkpoint loading, chat template formatting, and token sampling strategies.

## Loading a Model Checkpoint

The first step to perform inference is restoring the model from a saved checkpoint. The `load_model_from_ckpt` function in [`src/post_training/inference.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py) (lines 19-33) handles this automatically:

- Reads the configuration (`cfg`) stored in the checkpoint to reconstruct the model architecture
- Instantiates the `Transformer` class from [`src/models/transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py) with the correct dimensions (n_embed, n_head, n_layer, etc.)
- Strips common DDP prefixes (`module.` or `transformer.`) from state dict keys when models are saved with Distributed Data Parallel
- Places the model on the specified device and sets it to evaluation mode

```python
from src.post_training.inference import load_model_from_ckpt
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_model_from_ckpt("checkpoints/sft_stage.pt", device=device)

```

## Generating Text: Two Inference Modes

Once loaded, the `generate_reply` function (lines 36-71 in [`src/post_training/inference.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py)) provides a unified interface for text generation. It internally calls `batched_generate` from [`src/post_training/evaluation.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/evaluation.py) to handle the token-level sampling loop with support for temperature, top-k, top-p, and greedy decoding.

### Chat Mode

By default, `generate_reply` operates in chat mode, which formats your input using the instruction-tuned chat template defined in [`src/post_training/chat_template.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py). This mode accepts an optional `system` prompt to set the model's behavior and automatically encodes the formatted prompt before generation.

```python
reply = generate_reply(
    model,
    user_text="Explain quantum computing in simple terms.",
    device=device,
    system="You are a helpful physics tutor.",
    max_new_tokens=200,
    temperature=0.7,
    top_k=50,
    top_p=0.95,
)

```

### Raw Mode

For base models or when you need the model to continue text without conversational formatting, set `raw=True`. This bypasses the chat template and treats the input as a plain continuation prompt.

```python
reply = generate_reply(
    model,
    user_text="Once upon a time,",
    device=device,
    raw=True,
    max_new_tokens=100,
    temperature=1.0,
)

```

## Sampling Parameters: Temperature, Top-k, and Top-p

The inference API exposes fine-grained control over generation randomness through parameters passed to `batched_generate`:

- **`temperature`**: Scales logits before softmax (values < 1 produce more deterministic output, > 1 increases diversity)
- **`top_k`**: Limits sampling to the k highest probability tokens (set to `None` to disable)
- **`top_p`**: Implements nucleus sampling by restricting to tokens whose cumulative probability exceeds p (set to `None` to disable)
- **`greedy`**: When `True`, selects the highest probability token at each step instead of sampling

## Practical Code Examples

### Python API: Chat Inference

This example demonstrates loading a checkpoint and generating a response with system instructions:

```python
from src.post_training.inference import load_model_from_ckpt, generate_reply
import torch

ckpt_path = "checkpoints/dpo_stage.pt"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model with automatic config restoration

model = load_model_from_ckpt(ckpt_path, device)

# Generate chat response

reply = generate_reply(
    model,
    user_text="What are the health benefits of meditation?",
    device=device,
    system="You are a knowledgeable wellness coach.",
    max_new_tokens=256,
    temperature=0.8,
    top_k=40,
    top_p=0.9,
)

print("Assistant:", reply)

```

### Python API: Raw Continuation

Use this approach for base models or when you need the model to complete a partial text:

```python
reply = generate_reply(
    model,
    user_text="The key to effective machine learning is",
    device=device,
    raw=True,  # Bypass chat template

    max_new_tokens=150,
    temperature=0.9,
    top_k=None,
    top_p=None,
    greedy=False,
)
print(reply)

```

### Command-Line Interface

For quick testing without writing a Python script, use the CLI wrapper in [`scripts/chat.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/chat.py):

```bash
python -m scripts.chat \
    --ckpt checkpoints/ppo_stage.pt \
    --device cuda \
    --text "Write a Python function to sort a list" \
    --system "You are an expert programmer." \
    --max_new_tokens 200 \
    --temperature 0.7

```

The CLI supports all generation parameters including `--raw`, `--top_k`, and `--top_p`.

### Streamlit UI

For an interactive browser-based experience, launch the Streamlit application. The chat interface in [`ui/pages/9_Chat.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/ui/pages/9_Chat.py) imports the same `load_model_from_ckpt` and `generate_reply` utilities, allowing you to select checkpoints from the sidebar and conduct multi-turn conversations with the model.

## Core Implementation Details

The inference stack relies on several key components:

- **[`src/models/transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py)**: Implements the core Transformer architecture (multi-head attention, feed-forward layers, positional embeddings) used by all checkpoints
- **[`src/post_training/chat_template.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py)**: Contains `encode_prompt` and `decode` functions that construct the instruction-following prompt format and convert token IDs back to strings
- **[`src/post_training/evaluation.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/evaluation.py)**: Houses `batched_generate`, the low-level generation loop that implements temperature scaling and nucleus sampling
- **[`src/post_training/inference.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py)**: The public API entry point with `load_model_from_ckpt` and `generate_reply`

## Summary

- **Use `load_model_from_ckpt`** to automatically handle DDP prefixes, restore model configuration from the checkpoint, and place the model on your target device.
- **Select the appropriate inference mode**: Chat mode (default) for instruction-tuned models with system prompts, or raw mode (`raw=True`) for base model text continuation.
- **Control generation quality** via `temperature`, `top_k`, and `top_p` parameters passed directly to `generate_reply` to balance creativity against determinism.
- **Choose your interface**: Import the functions directly for Python scripts, use [`scripts/chat.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/chat.py) for command-line testing, or launch the Streamlit UI for interactive sessions.

## Frequently Asked Questions

### How do I load a checkpoint saved with Distributed Data Parallel (DDP)?

The `load_model_from_ckpt` function in [`src/post_training/inference.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py) automatically strips `module.` or `transformer.` prefixes from the state dictionary keys that DDP adds during multi-GPU training, allowing seamless loading without manual key renaming.

### What is the difference between chat mode and raw mode?

Chat mode (default) wraps your input in an instruction-tuned chat template using `encode_prompt` from [`src/post_training/chat_template.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py) and supports system-level instructions, while raw mode (`raw=True`) treats the input as a plain text continuation prompt, bypassing the template entirely for base model inference.

### Which sampling parameters control the randomness of generated text?

The `generate_reply` function accepts `temperature` (scaling logits before softmax), `top_k` (limiting vocabulary to the k highest probability tokens), and `top_p` (nucleus sampling threshold), which are passed to `batched_generate` in [`src/post_training/evaluation.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/evaluation.py) to control the diversity of the generated output.

### Can I perform inference without writing Python code?

Yes, you can use the CLI wrapper in [`scripts/chat.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/chat.py) which accepts checkpoint paths, device specifications, and sampling parameters as command-line arguments, or launch the Streamlit UI ([`ui/pages/9_Chat.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/ui/pages/9_Chat.py)) for an interactive browser-based interface that requires no coding.