complexity-router-bf16

ModernBERT BF16 Complexity Router

A conversation-complexity classifier built on ModernBERT that predicts a continuous complexity score from 0.0 to 10.0.

This README covers the native PyTorch/BF16 model architecture and the primary inference script:

predict_complexity_chunked_v2.py

This document intentionally focuses on the BF16 ModernBERT model and does not cover OpenVINO or INT8 deployment.

---

Overview

The router is designed to estimate how difficult an unresolved user request is so that another system can select an appropriate downstream LLM.

Conceptually:

Conversation
    |
    +-- prior conversation history
    |       |
    |       +-- serialize with role-aware special tokens
    |       |
    |       +-- densely pack into bounded token chunks
    |       |
    |       +-- ModernBERT encoder
    |       |
    |       +-- history embeddings
    |
    +-- final unresolved user request
            |
            +-- encode separately
            |
            +-- ModernBERT encoder
            |
            +-- final-user embedding

history embeddings
        +
final-user embedding
        |
        v
small Transformer router head
        |
        v
complexity score: 0.0 - 10.0

The final user request is intentionally encoded separately from the previous conversation history.

This allows the model to distinguish between conversations such as:

User: Design a highly available Kubernetes authentication system...
Assistant: ...
User: Add multi-region failover...
Assistant: ...
User: What is the capital of France?

and:

User: Design a highly available Kubernetes authentication system...
Assistant: ...
User: Add multi-region failover...
Assistant: ...
User: Why is it still failing?

The first request should generally become simple despite the difficult history.

The second request may still depend heavily on the preceding technical context.

---

Architecture

The model uses a shared ModernBERT encoder for both conversation-history chunks and the final user request.

Unlike an architecture that independently pads every message to the length of the longest message, conversation history is serialized and densely packed into chunks.

Typical configuration:

history chunk size:        2048 tokens
maximum retained history:  8192 tokens
maximum final-user length: 2048 tokens

For example, instead of processing:

message 1 = 1700 tokens
message 2 = 700 tokens

as:

encoder call 1 = 1700 tokens
encoder call 2 = 700 tokens

the serialized history can be densely packed approximately as:

chunk 1 = 2048 tokens
chunk 2 = 352 tokens

This reduces wasted encoder work and makes the configured history-token budget more closely represent the actual amount of ModernBERT computation.

---

Role-aware serialization

Conversation history is serialized using tokenizer special tokens representing message roles.

Examples include:

<|router_system|>
<|router_user|>
<|router_assistant|>

If a message crosses a chunk boundary, continuation markers can be used:

<|router_system_cont|>
<|router_user_cont|>
<|router_assistant_cont|>

These are real tokenizer special tokens added during training.

Because of this:

Always use the tokenizer stored with the trained router checkpoint.

Do not replace it with a newly downloaded copy of the base ModernBERT tokenizer.

---

Long conversations

The router places an upper bound on history processing.

With:

--max-history-tokens 8192

the model will process at most the configured history budget.

When conversation history exceeds that budget, the architecture favors retaining:

opening context
+
most recent context

A learned gap signal informs the router that conversation content was omitted between the retained sections.

This keeps inference cost bounded even when the original conversation becomes very large.

---

Final-user encoding

The unresolved final user request is encoded separately from the conversation history.

The default maximum is typically:

2048 tokens

If the final user request exceeds the configured maximum, head/tail truncation is used so that information from both the beginning and end of the request is retained.

---

Model Precision

The trained router checkpoint is stored using BF16 weights.

Typical training configuration:

training precision:     BF16 mixed precision
saved checkpoint dtype: BF16

BF16 provides a useful reduction in model storage and memory usage while maintaining substantially more exponent range than FP16.

On modern NVIDIA GPUs with native BF16 support, BF16 is the recommended inference format.

Example:

--device cuda --dtype bf16

A BF16 checkpoint does not require the model to execute using BF16.

For example, the same checkpoint can be loaded into FP32 for CPU execution:

BF16 checkpoint
      |
      v
FP32 PyTorch model
      |
      v
FP32 inference

This is often preferable on CPUs without native BF16 acceleration.

---

Files

A typical BF16 deployment requires:

predict_complexity_chunked_v2.py

complexity-router-v2/
└── best/
    ├── router_model.pt
    ├── router_config.json
    ├── tokenizer/
    │   └── ...
    └── encoder/
        └── config.json

The complete trained ModernBERT weights are stored inside:

router_model.pt

The predictor reconstructs the ModernBERT architecture from:

encoder/config.json

and then loads the trained weights from:

router_model.pt

A duplicate:

encoder/model.safetensors

is therefore not required for this architecture.

---

Requirements

Python

A modern Python 3 environment is recommended.

For inference, the primary dependencies are:

torch
transformers

Accelerate is not required by predict_complexity_chunked_v2.py.

---

Installation

NVIDIA GPU / CUDA

Install PyTorch using the appropriate CUDA build for the system.

Then install Transformers:

pip install torch transformers

If PyTorch is already installed with the correct CUDA version:

pip install transformers

Verify CUDA is visible:

python3 -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'No CUDA GPU')"

Expected output on a GPU system resembles:

True
NVIDIA ...

---

CPU-only installation

For a CPU-only environment:

pip install transformers
pip install torch --index-url https://download.pytorch.org/whl/cpu

Verify PyTorch:

python3 -c "import torch; print(torch.__version__)"

---

Input Format

predict_complexity_chunked_v2.py expects JSON containing a messages array.

Example:

{
  "messages": [
    {
      "role": "user",
      "content": "Create a PHP page-view class."
    },
    {
      "role": "assistant",
      "content": "A reusable implementation could..."
    },
    {
      "role": "user",
      "content": "Refactor it into a reusable PHP 8.3 component with immutable value objects, tests, CSS, PDO support, and an injectable renderer."
    }
  ]
}

The supported conversation structure is:

[optional system] U (A U)*

Where:

U = user
A = assistant

The final message must always be an unresolved:

user

message.

Valid examples:

user

system
user

user
assistant
user

system
user
assistant
user
assistant
user

Invalid examples include:

user
assistant

because the conversation ends with an assistant response rather than an unresolved user request.

---

Running the Model

The main BF16 predictor is:

predict_complexity_chunked_v2.py

---

Basic GPU inference

Create an input file such as:

complex2.json

Then run:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json \
  --device cuda \
  --dtype bf16

This is the recommended configuration for a CUDA GPU with native BF16 support.

---

Automatic Device Selection

If --device is not specified, the predictor can automatically use CUDA when it is available.

For example:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json

For predictable production behavior, explicitly selecting the desired device is generally preferable.

---

Force CUDA

Use:

--device cuda

Example:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json \
  --device cuda \
  --dtype bf16

---

Force CPU

Use:

--device cpu

For example:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json \
  --device cpu \
  --dtype fp32 \
  --threads 8

---

Inference Data Types

The predictor supports:

--dtype auto
--dtype bf16
--dtype fp32

---

BF16 GPU

Recommended:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json \
  --device cuda \
  --dtype bf16

Modern NVIDIA GPUs with BF16 Tensor Core support are the intended environment for this mode.

---

FP32 CPU

For general CPU inference, start with:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json \
  --device cpu \
  --dtype fp32 \
  --threads 8

Do not assume that loading a BF16 checkpoint means BF16 CPU execution will be faster.

On CPUs without efficient native BF16 execution, FP32 may perform substantially better.

The BF16 model can safely execute in FP32.

---

BF16 CPU

BF16 CPU inference can also be tested:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json \
  --device cpu \
  --dtype bf16 \
  --threads 8

Whether this is beneficial depends heavily on the CPU architecture and available BF16 instructions.

Benchmark both:

BF16
FP32

before selecting a production configuration.

---

CPU Thread Configuration

CPU inference supports:

--threads N

Example:

--threads 8

Complete example:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input complex2.json \
  --device cpu \
  --dtype fp32 \
  --threads 8

More threads do not automatically mean lower latency.

For latency-oriented inference, test values appropriate for the host, for example:

4
6
8
10
12

and compare warmed request latency.

---

Interactive Mode

predict_complexity_chunked_v2.py can remain running and accept multiple requests.

Simply omit:

--input

Example:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cuda \
  --dtype bf16

The model is loaded once and remains resident.

Paste JSON:

{
  "messages": [
    {
      "role": "user",
      "content": "Explain DNS."
    }
  ]
}

The predictor evaluates the request once the complete JSON object has been entered.

Both compact JSON:

{"messages":[{"role":"user","content":"Explain DNS."}]}

and pretty-printed JSON are supported.

---

Why Interactive Mode Matters

Do not benchmark production inference by repeatedly launching:

python3 predict_complexity_chunked_v2.py ...

for every request.

Doing so includes:

Python startup
+
tokenizer loading
+
model construction
+
checkpoint loading
+
CUDA initialization
+
inference

in every measurement.

Production should instead resemble:

process starts
    |
    +-- load tokenizer
    |
    +-- reconstruct ModernBERT
    |
    +-- load BF16 checkpoint
    |
    +-- move model to GPU
    |
    v
ready

request 1
request 2
request 3
request 4
...

The interactive mode demonstrates this resident-process behavior.

An HTTP service should follow the same model.

---

Example Prediction Output

A prediction can resemble:

{
  "complexity": 6.8193,
  "latency_ms": 631.56,
  "inference_ms": 620.53,
  "history_encoder_ms": 548.41,
  "final_encoder_ms": 67.81,
  "router_head_ms": 4.31,
  "preprocessing_ms": 10.47,
  "tensor_build_ms": 0.47,
  "validation_ms": 0.03,
  "device": "cpu",
  "dtype": "torch.float32",
  "messages": 5,
  "history_chunks": 1,
  "history_chunk_lengths": [
    1090
  ],
  "history_padded_tokens": 1090,
  "history_actual_tokens": 1090,
  "final_tokens": 89
}

The primary result is:

"complexity": 6.8193

The score is continuous:

0.0 -------------------------------------- 10.0
easy                                         hard

---

Understanding Latency Metrics

Several timing values are exposed to make performance analysis easier.

The most important are:

history_encoder_ms
final_encoder_ms
router_head_ms

For most workloads, the ModernBERT encoder dominates inference time.

Conceptually:

total inference
    |
    +-- history ModernBERT encoding
    |
    +-- final-user ModernBERT encoding
    |
    +-- router Transformer/head

The small router head generally represents only a small portion of the overall compute cost.

---

Cold Start vs Warm Inference

The first CUDA inference request may be significantly slower than subsequent requests.

The initial request may trigger initialization of:

CUDA runtime
memory allocator
cuBLAS
attention kernels
PyTorch execution paths

Therefore, production benchmarks should use warmed inference.

A reasonable manual benchmark sequence is:

load model
    |
run 2-3 warmup requests
    |
run 10+ measured requests
    |
compare median latency

Do not use the first request as the representative production latency.

---

Expected GPU Performance Characteristics

For appropriately sized conversations on a suitable CUDA GPU, warmed BF16 inference can reach low-double-digit-millisecond latency.

A representative execution may resemble:

history encoder: ~5 ms
final encoder:   ~5 ms
router head:     <1 ms

Actual latency depends strongly on:

GPU architecture
input length
number of history chunks
PyTorch version
Transformers version
CUDA version
system load

Benchmark the model on the hardware where it will actually run.

---

Routing Scores

The model outputs a continuous value from:

0.0

through:

10.0

The downstream routing system is free to define whatever boundaries are appropriate.

The training configuration currently uses boundaries around:

0.0 - 4.5    tier 1
4.5 - 8.5    tier 2
8.5 - 10.0   tier 3

These thresholds are routing policy rather than a limitation of the model.

For example:

if complexity < 4.5:
    model = "simple"

elif complexity < 8.5:
    model = "advanced"

else:
    model = "maximum"

The continuous score also allows thresholds to be adjusted later without retraining ModernBERT.

---

Training the BF16 Architecture

The corresponding training script is:

train_complexity_router_chunked_v2.py

Training requires:

pip install torch transformers accelerate

A typical BF16 training run is:

accelerate launch \
  --mixed_precision bf16 \
  train_complexity_router_chunked_v2.py \
  --train-file train.jsonl \
  --validation-file validation.jsonl \
  --test-file test.jsonl \
  --output-dir ./complexity-router-v2 \
  --history-chunk-tokens 2048 \
  --max-history-tokens 8192 \
  --max-final-user-tokens 2048 \
  --epochs 3 \
  --batch-size 1 \
  --gradient-accumulation 8

Important configuration:

history chunk size:        2048
maximum history:           8192
maximum final user:        2048
training precision:        BF16
saved checkpoint:          BF16

The best validation checkpoint is written to:

./complexity-router-v2/best/

That directory can then immediately be used by:

predict_complexity_chunked_v2.py

---

Training Data Format

Training data uses JSONL.

Each line contains a complete conversation and its target complexity.

Example:

{"messages":[{"role":"user","content":"What is 2 + 2?"}],"complexity":0.5}

Another example:

{
  "messages": [
    {
      "role": "user",
      "content": "Design a production application using PostgreSQL and Kubernetes."
    },
    {
      "role": "assistant",
      "content": "A possible architecture is..."
    },
    {
      "role": "user",
      "content": "Now redesign the authentication layer for multi-region HA with SSO, MFA, RBAC, audit logging, and a zero-downtime migration plan."
    }
  ],
  "complexity": 8.4
}

Prediction input does not contain:

complexity

because that is the value being predicted.

---

Model Directory Validation

Before running inference, verify that the model directory resembles:

complexity-router-v2/best/
├── router_model.pt
├── router_config.json
├── tokenizer/
│   ├── tokenizer.json
│   ├── tokenizer_config.json
│   └── ...
└── encoder/
    └── config.json

The exact tokenizer files may vary depending on the Transformers/tokenizer version.

The critical pieces are:

router_model.pt
router_config.json
tokenizer/
encoder/config.json

---

Common Errors

router_model.pt not found

Verify:

ls ./complexity-router-v2/best/router_model.pt

and confirm that:

--model-dir

points to the best directory rather than its parent.

Correct:

--model-dir ./complexity-router-v2/best

Incorrect:

--model-dir ./complexity-router-v2

unless that directory itself contains the expected checkpoint files.

---

Missing encoder/model.safetensors

This model layout intentionally does not require:

encoder/model.safetensors

The predictor should reconstruct ModernBERT from:

encoder/config.json

and load its trained parameters from:

router_model.pt

If inference attempts to call:

AutoModel.from_pretrained(...)

and requires encoder/model.safetensors, make sure you are using:

predict_complexity_chunked_v2.py

and not an older predictor written for a different checkpoint format.

---

CUDA is unavailable

Check:

python3 -c "import torch; print(torch.cuda.is_available())"

If this returns:

False

the installed PyTorch build may not include CUDA support, the NVIDIA driver may not be visible, or the process may not have access to the GPU.

---

BF16 CPU inference is slow

Use:

--device cpu --dtype fp32

instead.

A BF16 checkpoint does not require BF16 execution.

---

Invalid conversation ordering

The expected structure is:

[optional system] U (A U)*

The final message must be:

user

For example, this is invalid:

{
  "messages": [
    {
      "role": "user",
      "content": "Explain DNS."
    },
    {
      "role": "assistant",
      "content": "DNS is..."
    }
  ]
}

because there is no unresolved user request to classify.

---

Production Deployment

For production use, keep the model permanently resident in memory.

Do not execute:

python3 predict_complexity_chunked_v2.py ...

as a new process for every incoming request.

Instead:

application starts
      |
      +-- initialize tokenizer
      |
      +-- construct ModernBERT architecture
      |
      +-- load router_model.pt
      |
      +-- move model to CUDA
      |
      +-- set inference/evaluation mode
      |
      v
   READY
      |
      +-- request
      +-- request
      +-- request
      +-- request

For a production routing service, predict_complexity_chunked_v2.py should be treated as the reference implementation for:

input validation
conversation serialization
history chunking
tokenization
ModernBERT execution
router scoring

and these components can be wrapped in a persistent HTTP API.

---

Recommended Production Configuration

NVIDIA GPU

For a GPU with native BF16 support:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cuda \
  --dtype bf16

This is the preferred native PyTorch deployment configuration.

---

CPU

Start with:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cpu \
  --dtype fp32 \
  --threads 8

Then benchmark different thread counts.

If the CPU has strong native BF16 acceleration, also test:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cpu \
  --dtype bf16 \
  --threads 8

Select the configuration based on warmed median latency rather than checkpoint dtype alone.

---

Quick Start

1. Install dependencies

GPU:

pip install torch transformers

CPU-only:

pip install transformers
pip install torch --index-url https://download.pytorch.org/whl/cpu

---

2. Create an input file

request.json:

{
  "messages": [
    {
      "role": "user",
      "content": "Explain why the sky is blue."
    }
  ]
}

---

3. Run with CUDA BF16

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input request.json \
  --device cuda \
  --dtype bf16

---

4. Or run with CPU FP32

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input request.json \
  --device cpu \
  --dtype fp32 \
  --threads 8

---

5. Run interactively

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cuda \
  --dtype bf16

Then submit JSON objects repeatedly without reloading the model.

---

Quick Reference

GPU BF16

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input request.json \
  --device cuda \
  --dtype bf16

CPU FP32

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --input request.json \
  --device cpu \
  --dtype fp32 \
  --threads 8

Interactive GPU

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cuda \
  --dtype bf16

Interactive CPU

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cpu \
  --dtype fp32 \
  --threads 8

Train BF16

accelerate launch \
  --mixed_precision bf16 \
  train_complexity_router_chunked_v2.py \
  --train-file train.jsonl \
  --validation-file validation.jsonl \
  --test-file test.jsonl \
  --output-dir ./complexity-router-v2 \
  --history-chunk-tokens 2048 \
  --max-history-tokens 8192 \
  --max-final-user-tokens 2048 \
  --epochs 3 \
  --batch-size 1 \
  --gradient-accumulation 8

---

Summary

The BF16 ModernBERT complexity router provides:

• a continuous 0-10 complexity score;
• a shared ModernBERT encoder;
• densely packed conversation-history chunks;
• separate encoding of the unresolved final user request;
• role-aware tokenizer special tokens;
• bounded processing of long histories;
• head/tail retention for long final-user requests;
• BF16 training and checkpoint storage;
• native CUDA BF16 inference;
• optional FP32 CPU inference from the same BF16 checkpoint;
• detailed latency instrumentation;
• resident interactive inference suitable as the basis for an HTTP service.

For normal GPU deployment, the primary command is:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./complexity-router-v2/best \
  --device cuda \
  --dtype bf16

For one-shot file prediction:

python3 predict_complexity_chunked_v2.py \
  --model-dir ./model \
  --input request.json \
  --device cuda \
  --dtype bf16

predict_complexity_chunked_v2.py should be considered the reference native PyTorch/BF16 inference implementation for this architecture.