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

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 (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 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

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) provides a unified interface for text generation. It internally calls batched_generate from 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. This mode accepts an optional system prompt to set the model's behavior and automatically encodes the formatted prompt before generation.

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.

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:

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:

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:

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 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: Implements the core Transformer architecture (multi-head attention, feed-forward layers, positional embeddings) used by all checkpoints
• 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: Houses batched_generate, the low-level generation loop that implements temperature scaling and nucleus sampling
• 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 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 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 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 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 which accepts checkpoint paths, device specifications, and sampling parameters as command-line arguments, or launch the Streamlit UI (ui/pages/9_Chat.py) for an interactive browser-based interface that requires no coding.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}