# How to Format Instructions and Use Chat Templates for LLM Input

> Learn to format instructions and use chat templates for LLM input with our tokenizer-agnostic implementation. Convert messages to token sequences and masks efficiently.

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

---

**The train-llm-from-scratch repository provides a tokenizer-agnostic chat template implementation that converts role-based message lists into token sequences and binary loss masks, using plain-text role headers and a single end-of-text token for turn separation.**

This minimal implementation, found in [`src/post_training/chat_template.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py), handles the complete pipeline from conversation formatting to training-ready tokenization. It works with any tokenizer that supports byte-pair encoding (BPE) by treating role markers as ordinary text rather than special tokens, allowing the model to learn turn-taking through supervised fine-tuning (SFT).

## Understanding the Chat Template Design

The chat template system centers on a single special token—the end-of-text marker—while keeping role identifiers as regular text tokens.

### Special Tokens and Role Markers

The only true special token in this implementation is the end-of-text token (`<|endoftext|>`) with a fixed ID of **50256**, defined as `EOT_ID` in [`src/post_training/chat_template.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py) (lines 38-40). This token serves dual purposes: it terminates each turn in the conversation and acts as the generation stop token during inference.

Role headers (`<|user|>\n`, `<|assistant|>\n`, `<|system|>\n`) are implemented as plain strings via `USER_HEADER`, `ASSISTANT_HEADER`, and `SYSTEM_HEADER` (lines 42-45). Because these are ordinary token strings rather than registered special tokens, they are encoded using standard BPE tokenization without requiring custom tokenizer configuration.

### Conversation Format Structure

A single complete turn follows this strict pattern:

```

<|user|>
{user_content}<|endoftext|><|assistant|>
{assistant_content}<|endoftext|>

```

For inference scenarios where you want the model to generate a response, you append only the assistant header to cue generation. The system prompt, if present, appears at the beginning of the sequence using the `<|system|>\n` header format.

## Encoding Conversations for Training and Inference

The module provides three primary utilities for converting between message dictionaries and model-ready inputs.

### The encode_chat Function

The `encode_chat` function (lines 99-136) processes a list of message dictionaries and returns two aligned lists: token IDs and a binary loss mask. The encoding pipeline executes these steps:

1. Encode the role header using `tiktoken.encode_ordinary`
2. Encode the message content
3. Append the EOT token (ID 50256)
4. Build the mask concurrently, assigning **1** to assistant response tokens (including the final EOT) and **0** to all role headers, user content, and system prompts

```python
from src.post_training.chat_template import encode_chat

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Explain the difference between SFT and RLHF."},
    {"role": "assistant", "content": "Supervised fine-tuning uses labeled data..."}
]

token_ids, loss_mask = encode_chat(messages)

# loss_mask contains 1s only for assistant tokens and final EOT

```

### The encode_prompt Function

For inference workflows, `encode_prompt` (line 95) returns only token IDs for a prompt that ends with an assistant header, preparing the sequence for autoregressive generation. This function sets `add_generation_prompt=True` internally, ensuring the model receives the cue to begin generating without any target tokens to mask.

```python
from src.post_training.chat_template import encode_prompt

prompt_ids = encode_prompt(messages)

# Sequence ends with <|assistant|>\n, ready for generation

```

### Rendering and Decoding Utilities

The `render_chat` function (line 73) produces a human-readable string representation for debugging, while `decode` (lines 139, 145) safely converts token IDs back to text, stripping any IDs greater than or equal to `EOT_ID` to prevent junk output.

```python
from src.post_training.chat_template import render_chat

print(render_chat(messages, add_generation_prompt=True))

```

## Preparing Instruction Data

Structure your conversation data as a list of dictionaries with specific role keys before encoding. The system message is optional but recommended for setting behavioral context.

```python
conversation = [
    {"role": "system", "content": "You are a concise coding assistant."},
    {"role": "user", "content": "Write a Python function to reverse a string."},
    {"role": "assistant", "content": "def reverse_string(s):\n    return s[::-1]"}
]

```

When `add_generation_prompt=True` is passed to encoding functions, the sequence ends immediately after the `<|assistant|>\n` header, resulting in a loss mask of all zeros since no target tokens exist for prediction.

## Applying Loss Masks During Training

The binary loss mask enables selective gradient computation on only the assistant's responses. Integrate this mask with your loss function to ignore predictions on user instructions and role headers.

```python
import torch
import torch.nn.functional as F

# token_ids and loss_mask from encode_chat()

input_ids = torch.tensor([token_ids], dtype=torch.long)
mask = torch.tensor([loss_mask], dtype=torch.float)

logits, _ = model(input_ids)  # shape [B, T, vocab_size]

logits = logits[:, :-1, :].reshape(-1, logits.size(-1))
targets = input_ids[:, 1:].reshape(-1)

loss = F.cross_entropy(logits, targets, reduction='none')
masked_loss = (loss * mask[:, 1:].reshape(-1)).mean()

```

This approach ensures the model learns to generate assistant responses while ignoring the prompt tokens during backpropagation, which is essential for both SFT and reinforcement learning from human feedback (RLHF) workflows.

## Summary

- **Tokenizer-agnostic design**: Role headers (`<|user|>`, `<|assistant|>`, `<|system|>`) are plain text strings encoded via standard BPE, requiring no special tokenizer configuration.
- **Single special token**: Only `<|endoftext|>` (ID 50256) is treated as special, serving as both turn delimiter and generation stop token.
- **Dual encoding modes**: Use `encode_chat` for training (returns IDs + mask) and `encode_prompt` for inference (returns IDs ending with assistant header).
- **Binary masking**: Loss masks contain 1s only for assistant tokens and the final EOT, enabling selective training on completions while ignoring prompts.

## Frequently Asked Questions

### What is the difference between encode_chat and encode_prompt?

The `encode_chat` function processes complete conversations including assistant responses, returning both token IDs and a binary loss mask where assistant tokens are marked with 1s. The `encode_prompt` function is designed for inference, returning only token IDs for the prompt ending with an assistant header, effectively setting the loss mask to all zeros since no target response exists yet.

### Why are role headers implemented as plain text instead of special tokens?

Implementing role headers (`<|user|>`, `<|assistant|>`, etc.) as ordinary text strings makes the implementation tokenizer-agnostic. According to the source code 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 design avoids dependencies on custom tokenizer configurations and allows the model to learn turn-taking patterns through supervised fine-tuning rather than relying on hardcoded special token IDs.

### How does the loss mask work in this implementation?

The loss mask is a binary array aligned with the token sequence, created in `encode_chat` (lines 99-136). It assigns **1** to tokens belonging to the assistant's content (including the final EOT token) and **0** to all system prompts, user messages, and role headers. When computing cross-entropy loss, this mask zeroes out gradients for non-assistant tokens, ensuring the model only learns to predict the assistant's responses.

### Can I use this chat template with any tokenizer?

Yes, the implementation is specifically designed to work with any BPE-based tokenizer because it relies only on standard text encoding for role headers and the fixed end-of-text token ID 50256 (compatible with GPT-2 and similar tokenizers). The code uses `tiktoken.encode_ordinary` for the headers, meaning no custom special token registration is required in your tokenizer configuration.