Handling Custom Chat Templates and Tokenizers

Custom chat templates and tokenizers are essential components for properly formatting conversational inputs and managing token generation in your imported models. Custom Model Import supports importing models with custom chat templates that define how multi-turn conversations are structured and tokenized.

Importing Models with Custom Chat Templates

When importing a model with custom chat templates, ensure that you follow the best practices outlined in the main Custom Model Import documentation, including bringing in model weights in safe tensors format and providing all necessary configuration files.

Chat Template Format Requirements

If your model uses a custom chat template and you want to use that template with Amazon Bedrock, you must include the chat template file in one of the following formats:

Custom Model Import will automatically process these files and include them in the correct directories during the import process.

Chat Template Precedence

If you provide multiple chat template sources, Amazon Bedrock applies the following precedence rules:

Chat Template Format Examples

Here are examples of the two supported chat template formats:

Jinja Format (chat_template.jinja)

A simplified example of a Jinja2-based chat template:

{% for message in messages %}
{% if loop.first and message['role'] != 'system' %}
<|im_start|>system
You are a helpful assistant.<|im_end|>
{% endif %}
<|im_start|>{{ message['role'] }}
{{ message['content'] }}<|im_end|>
{% endfor %}
{% if add_generation_prompt %}
<|im_start|>assistant
{% endif %}

For a complete example, see the GPT-OSS chat template on Hugging Face.

JSON Format (chat_template.json)

A simplified example of a JSON-based chat template with vision support:

{
    "chat_template": "{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' %}<|vision_start|><|image_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
}

For a complete example with multimodal support, see the Qwen2-VL chat template on Hugging Face.

Invoking Models with Custom Chat Templates

Once you have imported a model with a custom chat template, you have two options for invoking the model with properly formatted conversational inputs:

Using OpenAI ChatCompletion API with Messages

If you are providing inputs in the messages schema format, you should use the OpenAI ChatCompletion API. When you use this API format with a model that has an uploaded chat template (chat_template.jinja or chat_template.json), Amazon Bedrock will automatically convert your input messages into the correct format using the chat template.

This is the recommended approach as it provides the most seamless integration and allows Amazon Bedrock to handle the chat template application automatically.

Example: Using OpenAI ChatCompletion API with custom chat template

import json
import boto3

# Initialize Bedrock Runtime client
client = boto3.client('bedrock-runtime', region_name='us-east-1')

# Define the model ARN for your imported model with custom chat template
model_id = 'arn:aws:bedrock:us-east-1:123456789012:imported-model/your-model-id'

# Prepare the request payload using messages format
payload = {
    "messages": [
        {
            "role": "system",
            "content": "You are a helpful assistant."
        },
        {
            "role": "user",
            "content": "Hello, how are you?"
        }
    ],
    "max_tokens": 150,
    "temperature": 0.7
}

# Invoke the model
response = client.invoke_model(
    modelId=model_id,
    body=json.dumps(payload),
    accept='application/json',
    contentType='application/json'
)

# Parse and display the response
response_body = json.loads(response['body'].read())
print(json.dumps(response_body, indent=2))

In this example, Amazon Bedrock will automatically apply your custom chat template to the messages array, converting it into the appropriate format for your model.

Manual Tokenization Approach

Alternatively, if you prefer to have full control over the chat template application and tokenization process, you can manually apply the chat template to your conversation and then use the Completion API (not ChatCompletion) with the pre-formatted text.

This approach is useful when you need to customize the template application logic or when working with specialized tokenization requirements.

Step 1: Apply Chat Template Locally

Use the following code snippet to load your custom chat template and apply it to a conversation locally:

from transformers import AutoTokenizer

# Configuration paths - update these with your actual paths
TOKENIZER_PATH = ""  # Path to tokenizer directory

# Load tokenizer with updated chat template
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)

# Test chat template with sample conversation
chat_history = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Hello, how are you?"},
]

# Apply chat template and display formatted output to make sure chat template works as expected
formatted_chat = tokenizer.apply_chat_template(chat_history, tokenize=False)
print(formatted_chat)

This code demonstrates how to:

Step 2: Invoke with Completion API

After applying the chat template locally, use the formatted text with the Completion API:

import json
import boto3

# Initialize Bedrock Runtime client
client = boto3.client('bedrock-runtime', region_name='us-east-1')

# Define the model ARN for your imported model
model_id = 'arn:aws:bedrock:us-east-1:123456789012:imported-model/your-model-id'

# Use the formatted_chat output from Step 1 as the prompt
payload = {
    "prompt": formatted_chat,
    "max_tokens": 150,
    "temperature": 0.7
}

# Invoke the model using Completion format (not ChatCompletion)
response = client.invoke_model(
    modelId=model_id,
    body=json.dumps(payload),
    accept='application/json',
    contentType='application/json'
)

# Parse and display the response
response_body = json.loads(response['body'].read())
print(json.dumps(response_body, indent=2))

Best Practices