predict_complexity_chunked_v2.py

#!/usr/bin/env python3
"""
Standalone PyTorch predictor for the chunked-history ModernBERT complexity router.

Designed for checkpoints produced by:
    train_complexity_router_chunked_v2.py

Runtime dependencies:
    torch
    transformers

No accelerate dependency is required.

Architecture
------------
- History is densely serialized into role-aware chunks.
- Final USER request is encoded separately.
- Shared ModernBERT encoder processes history chunks and final request.
- Tiny conversation Transformer aggregates embeddings.
- Output is a continuous complexity score in [0, 10].

The model directory should look like:

best/
├── encoder/
│   └── config.json
├── router_model.pt
├── router_config.json
└── tokenizer/
    ├── tokenizer.json
    └── tokenizer_config.json

Examples
--------
Single file:

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

Interactive mode, keeping the model resident in RAM:

python3 predict_complexity.py \
    --model-dir ./complexity-router/best

Then paste a complete pretty-printed JSON object. It is evaluated automatically
as soon as the JSON object is complete.

Optional CPU thread control:

python3 predict_complexity.py \
    --model-dir ./complexity-router/best \
    --threads 8

Force FP32 for comparison:

python3 predict_complexity.py \
    --model-dir ./complexity-router/best \
    --dtype fp32 \
    --input complex2.json
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Sequence, Tuple

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, AutoTokenizer


# =============================================================================
# CONSTANTS
# =============================================================================

MAX_USER_TURNS = 6

DEFAULT_SPECIAL_TOKENS = [
    "<|router_system|>",
    "<|router_user|>",
    "<|router_assistant|>",
    "<|router_system_cont|>",
    "<|router_user_cont|>",
    "<|router_assistant_cont|>",
]

DEFAULT_ROLE_START_TOKEN = {
    "system": "<|router_system|>",
    "user": "<|router_user|>",
    "assistant": "<|router_assistant|>",
}

DEFAULT_ROLE_CONT_TOKEN = {
    "system": "<|router_system_cont|>",
    "user": "<|router_user_cont|>",
    "assistant": "<|router_assistant_cont|>",
}


# =============================================================================
# VALIDATION
# =============================================================================

def expected_roles(
    user_turns: int,
    include_system: bool,
) -> List[str]:
    roles: List[str] = []

    if include_system:
        roles.append("system")

    for i in range(user_turns):
        roles.append("user")

        if i < user_turns - 1:
            roles.append("assistant")

    return roles


def validate_role_sequence(
    messages: List[Dict[str, Any]],
) -> None:
    if not isinstance(messages, list) or not messages:
        raise ValueError(
            "conversation contains no messages"
        )

    include_system = (
        messages[0].get("role") == "system"
    )

    conversational = (
        messages[1:]
        if include_system
        else messages
    )

    if (
        not conversational
        or len(conversational) % 2 == 0
    ):
        raise ValueError(
            "conversation after optional system message must have "
            "shape U (A U)* and end with an unanswered user message"
        )

    user_turns = (
        len(conversational) + 1
    ) // 2

    if not 1 <= user_turns <= MAX_USER_TURNS:
        raise ValueError(
            f"expected 1-{MAX_USER_TURNS} user turns; got {user_turns}"
        )

    expected = expected_roles(
        user_turns,
        include_system,
    )

    actual = [
        message.get("role")
        for message in messages
    ]

    if actual != expected:
        raise ValueError(
            f"invalid roles {actual}; expected {expected}"
        )

    for i, message in enumerate(messages):
        content = message.get("content")

        if (
            not isinstance(content, str)
            or not content.strip()
        ):
            raise ValueError(
                f"message {i} has empty/non-string content"
            )


# =============================================================================
# TOKENIZATION / CHUNKING
# =============================================================================

def tokenize_without_specials(
    tokenizer,
    text: str,
) -> List[int]:
    encoded = tokenizer(
        text,
        add_special_tokens=False,
        return_attention_mask=False,
        return_token_type_ids=False,
    )

    return list(
        encoded["input_ids"]
    )


def wrap_special_tokens(
    tokenizer,
    token_ids: Sequence[int],
) -> Dict[str, List[int]]:
    cls_token_id = tokenizer.cls_token_id
    sep_token_id = tokenizer.sep_token_id

    if (
        cls_token_id is None
        or sep_token_id is None
    ):
        raise ValueError(
            "Tokenizer must provide cls_token_id and sep_token_id"
        )

    input_ids = [
        int(cls_token_id),
        *map(int, token_ids),
        int(sep_token_id),
    ]

    return {
        "input_ids": input_ids,
        "attention_mask": [
            1
        ] * len(input_ids),
    }


def head_tail_content(
    token_ids: Sequence[int],
    content_budget: int,
) -> List[int]:
    token_ids = list(
        token_ids
    )

    if len(token_ids) <= content_budget:
        return token_ids

    head_tokens = (
        content_budget // 2
    )

    tail_tokens = (
        content_budget - head_tokens
    )

    return (
        token_ids[:head_tokens]
        + token_ids[-tail_tokens:]
    )


def encode_final_user(
    tokenizer,
    message: Dict[str, Any],
    max_final_user_tokens: int,
) -> Dict[str, List[int]]:
    content_budget = (
        max_final_user_tokens - 2
    )

    if content_budget < 8:
        raise ValueError(
            "max_final_user_tokens is too small"
        )

    user_ids = tokenize_without_specials(
        tokenizer,
        message["content"],
    )

    user_ids = head_tail_content(
        user_ids,
        content_budget,
    )

    return wrap_special_tokens(
        tokenizer,
        user_ids,
    )


def router_marker_ids(
    tokenizer,
    role_start_tokens: Dict[str, str],
    role_cont_tokens: Dict[str, str],
) -> Tuple[
    Dict[str, int],
    Dict[str, int],
]:
    start_ids: Dict[str, int] = {}
    cont_ids: Dict[str, int] = {}

    for role, token in role_start_tokens.items():
        token_id = tokenizer.convert_tokens_to_ids(
            token
        )

        if (
            token_id is None
            or token_id == tokenizer.unk_token_id
        ):
            raise ValueError(
                f"Missing router special token in tokenizer: {token}"
            )

        start_ids[
            role
        ] = int(
            token_id
        )

    for role, token in role_cont_tokens.items():
        token_id = tokenizer.convert_tokens_to_ids(
            token
        )

        if (
            token_id is None
            or token_id == tokenizer.unk_token_id
        ):
            raise ValueError(
                f"Missing router continuation token in tokenizer: {token}"
            )

        cont_ids[
            role
        ] = int(
            token_id
        )

    return (
        start_ids,
        cont_ids,
    )


def dense_history_payload_chunks(
    tokenizer,
    history_messages: Sequence[Dict[str, Any]],
    history_chunk_tokens: int,
    role_start_tokens: Dict[str, str],
    role_cont_tokens: Dict[str, str],
) -> List[List[int]]:
    """
    Must match training preprocessing exactly.

    Chunks are densely filled. A role marker begins each message. When a
    message crosses a chunk boundary, the continuation chunk begins with the
    role-specific continuation marker.
    """

    payload_budget = (
        history_chunk_tokens - 2
    )

    if payload_budget < 8:
        raise ValueError(
            "history_chunk_tokens is too small"
        )

    start_ids, cont_ids = router_marker_ids(
        tokenizer,
        role_start_tokens,
        role_cont_tokens,
    )

    chunks: List[List[int]] = []
    current: List[int] = []

    def flush_current() -> None:
        nonlocal current

        if current:
            chunks.append(
                current
            )

            current = []

    for message in history_messages:
        role = message["role"]

        content_ids = tokenize_without_specials(
            tokenizer,
            message["content"],
        )

        offset = 0
        first_segment = True

        while (
            offset < len(content_ids)
            or (
                first_segment
                and not content_ids
            )
        ):
            marker_id = (
                start_ids[role]
                if first_segment
                else cont_ids[role]
            )

            remaining = (
                payload_budget
                - len(current)
            )

            required_minimum = (
                2
                if offset < len(content_ids)
                else 1
            )

            if (
                current
                and remaining < required_minimum
            ):
                flush_current()

            current.append(
                marker_id
            )

            remaining = (
                payload_budget
                - len(current)
            )

            if offset < len(content_ids):
                take = min(
                    remaining,
                    len(content_ids) - offset,
                )

                current.extend(
                    content_ids[
                        offset:offset + take
                    ]
                )

                offset += take

            if offset < len(content_ids):
                flush_current()

                first_segment = False

            else:
                break

    flush_current()

    return chunks


@dataclass
class HistoryChunk:
    input_ids: List[int]
    attention_mask: List[int]
    gap_before: int = 0


def retain_history_with_budget(
    tokenizer,
    payload_chunks: Sequence[Sequence[int]],
    max_history_tokens: int,
) -> List[HistoryChunk]:
    """
    Match the training retention policy:
    first chunk + newest chunks when truncation is necessary.
    """

    if not payload_chunks:
        return []

    wrapped_lengths = [
        len(chunk) + 2
        for chunk in payload_chunks
    ]

    raw_total = sum(
        wrapped_lengths
    )

    if raw_total <= max_history_tokens:
        result: List[HistoryChunk] = []

        for payload in payload_chunks:
            wrapped = wrap_special_tokens(
                tokenizer,
                payload,
            )

            result.append(
                HistoryChunk(
                    input_ids=wrapped["input_ids"],
                    attention_mask=wrapped["attention_mask"],
                    gap_before=0,
                )
            )

        return result

    selected_indices: List[int] = []

    used = 0

    first_len = wrapped_lengths[
        0
    ]

    last_len = wrapped_lengths[
        -1
    ]

    if (
        len(payload_chunks) > 1
        and first_len + last_len
        <= max_history_tokens
    ):
        selected_indices.append(
            0
        )

        used += first_len

    newest_indices: List[int] = []

    for idx in range(
        len(payload_chunks) - 1,
        -1,
        -1,
    ):
        if (
            idx == 0
            and 0 in selected_indices
        ):
            continue

        length = wrapped_lengths[
            idx
        ]

        if (
            used + length
            <= max_history_tokens
        ):
            newest_indices.append(
                idx
            )

            used += length

    newest_indices.reverse()

    selected_indices.extend(
        newest_indices
    )

    selected_indices = sorted(
        set(
            selected_indices
        )
    )

    if not selected_indices:
        newest_payload = list(
            payload_chunks[-1]
        )

        content_budget = max(
            1,
            max_history_tokens - 2,
        )

        selected_payloads = [
            newest_payload[
                -content_budget:
            ]
        ]

        selected_indices = [
            len(payload_chunks) - 1
        ]

    else:
        selected_payloads = [
            list(
                payload_chunks[idx]
            )
            for idx in selected_indices
        ]

    result: List[HistoryChunk] = []

    for retained_pos, (
        original_idx,
        payload,
    ) in enumerate(
        zip(
            selected_indices,
            selected_payloads,
        )
    ):
        if retained_pos == 0:
            gap_before = (
                1
                if original_idx > 0
                else 0
            )

        else:
            previous_original_idx = (
                selected_indices[
                    retained_pos - 1
                ]
            )

            gap_before = (
                1
                if original_idx
                != previous_original_idx + 1
                else 0
            )

        wrapped = wrap_special_tokens(
            tokenizer,
            payload,
        )

        result.append(
            HistoryChunk(
                input_ids=wrapped["input_ids"],
                attention_mask=wrapped["attention_mask"],
                gap_before=gap_before,
            )
        )

    return result


def build_history_chunks(
    tokenizer,
    history_messages: Sequence[Dict[str, Any]],
    history_chunk_tokens: int,
    max_history_tokens: int,
    role_start_tokens: Dict[str, str],
    role_cont_tokens: Dict[str, str],
) -> List[HistoryChunk]:
    payload_chunks = dense_history_payload_chunks(
        tokenizer,
        history_messages,
        history_chunk_tokens,
        role_start_tokens,
        role_cont_tokens,
    )

    return retain_history_with_budget(
        tokenizer,
        payload_chunks,
        max_history_tokens,
    )


# =============================================================================
# MODEL
# =============================================================================

class ChunkedHistoryComplexityRouter(
    nn.Module
):
    def __init__(
        self,
        encoder_config,
        tokenizer_size: int,
        aggregator_layers: int,
        aggregator_heads: int,
        aggregator_dropout: float,
        max_history_chunks: int,
    ):
        super().__init__()

        self.max_history_chunks = (
            max_history_chunks
        )

        self.max_router_items = (
            max_history_chunks + 1
        )

        # CONFIG ONLY:
        # no encoder/model.safetensors is required.
        self.encoder = AutoModel.from_config(
            encoder_config
        )

        self.encoder.resize_token_embeddings(
            tokenizer_size
        )

        hidden = (
            self.encoder.config.hidden_size
        )

        self.item_type_embedding = nn.Embedding(
            2,
            hidden,
        )

        self.position_embedding = nn.Embedding(
            self.max_router_items,
            hidden,
        )

        self.final_user_marker = nn.Parameter(
            torch.zeros(
                hidden
            )
        )

        self.history_gap_marker = nn.Parameter(
            torch.zeros(
                hidden
            )
        )

        self.conversation_token = nn.Parameter(
            torch.zeros(
                1,
                1,
                hidden,
            )
        )

        layer = nn.TransformerEncoderLayer(
            d_model=hidden,
            nhead=aggregator_heads,
            dim_feedforward=hidden * 2,
            dropout=aggregator_dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )

        self.aggregator = nn.TransformerEncoder(
            layer,
            num_layers=aggregator_layers,
        )

        self.regressor = nn.Sequential(
            nn.LayerNorm(
                hidden * 2
            ),
            nn.Linear(
                hidden * 2,
                hidden // 2,
            ),
            nn.GELU(),
            nn.Dropout(
                aggregator_dropout
            ),
            nn.Linear(
                hidden // 2,
                1,
            ),
        )

    @staticmethod
    def masked_mean(
        last_hidden_state: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        mask = attention_mask.unsqueeze(
            -1
        ).to(
            last_hidden_state.dtype
        )

        summed = (
            last_hidden_state
            * mask
        ).sum(
            dim=1
        )

        denom = mask.sum(
            dim=1
        ).clamp_min(
            1e-6
        )

        return (
            summed / denom
        )

    def encode_sequences(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )

        return self.masked_mean(
            outputs.last_hidden_state,
            attention_mask,
        )

    def score_embeddings(
        self,
        history_vectors: torch.Tensor,
        final_vectors: torch.Tensor,
        history_positions: torch.Tensor,
        history_gap_flags: torch.Tensor,
        history_count: int,
    ) -> torch.Tensor:
        """
        Single-conversation inference head.
        """

        batch_size = 1

        hidden = (
            final_vectors.shape[-1]
        )

        device = (
            final_vectors.device
        )

        dtype = (
            final_vectors.dtype
        )

        router_items = torch.zeros(
            batch_size,
            self.max_router_items,
            hidden,
            device=device,
            dtype=dtype,
        )

        if history_count > 0:
            history_type_ids = torch.zeros(
                history_count,
                dtype=torch.long,
                device=device,
            )

            history_enriched = (
                history_vectors
                + self.item_type_embedding(
                    history_type_ids
                ).to(dtype)
                + self.position_embedding(
                    history_positions
                ).to(dtype)
            )

            if history_gap_flags.numel() > 0:
                history_enriched = (
                    history_enriched
                    + history_gap_flags.to(
                        dtype
                    ).unsqueeze(-1)
                    * self.history_gap_marker.to(
                        dtype
                    )
                )

            router_items[
                0,
                :history_count,
            ] = history_enriched

        final_position = (
            history_count
        )

        final_type_id = torch.ones(
            1,
            dtype=torch.long,
            device=device,
        )

        final_enriched = (
            final_vectors
            + self.item_type_embedding(
                final_type_id
            ).to(dtype)
            + self.position_embedding(
                torch.tensor(
                    [final_position],
                    dtype=torch.long,
                    device=device,
                )
            ).to(dtype)
            + self.final_user_marker.to(
                dtype
            )
        )

        router_items[
            0,
            final_position,
        ] = final_enriched[
            0
        ]

        router_length = (
            history_count + 1
        )

        valid_item_mask = (
            torch.arange(
                self.max_router_items,
                device=device,
            )
            < router_length
        ).unsqueeze(
            0
        )

        cls = self.conversation_token.expand(
            1,
            -1,
            -1,
        ).to(
            dtype
        )

        agg_input = torch.cat(
            [
                cls,
                router_items,
            ],
            dim=1,
        )

        valid_with_cls = torch.cat(
            [
                torch.ones(
                    1,
                    1,
                    device=device,
                    dtype=torch.bool,
                ),
                valid_item_mask,
            ],
            dim=1,
        )

        aggregated = self.aggregator(
            agg_input,
            src_key_padding_mask=(
                ~valid_with_cls
            ),
        )

        conversation_vector = (
            aggregated[:, 0, :]
        )

        final_user_vector = aggregated[
            :,
            final_position + 1,
            :,
        ]

        features = torch.cat(
            [
                conversation_vector,
                final_user_vector,
            ],
            dim=-1,
        )

        raw_score = self.regressor(
            features
        ).squeeze(
            -1
        )

        return (
            10.0
            * torch.sigmoid(
                raw_score
            )
        )


# =============================================================================
# INPUT PREPARATION
# =============================================================================

def pad_history_chunks(
    chunks: Sequence[HistoryChunk],
    pad_token_id: int,
    device: torch.device,
) -> Tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    if not chunks:
        return (
            torch.empty(
                0,
                1,
                dtype=torch.long,
                device=device,
            ),
            torch.empty(
                0,
                1,
                dtype=torch.long,
                device=device,
            ),
            torch.empty(
                0,
                dtype=torch.long,
                device=device,
            ),
            torch.empty(
                0,
                dtype=torch.long,
                device=device,
            ),
        )

    max_len = max(
        len(chunk.input_ids)
        for chunk in chunks
    )

    input_ids = torch.full(
        (
            len(chunks),
            max_len,
        ),
        pad_token_id,
        dtype=torch.long,
        device=device,
    )

    attention_mask = torch.zeros(
        (
            len(chunks),
            max_len,
        ),
        dtype=torch.long,
        device=device,
    )

    gap_flags = torch.empty(
        len(chunks),
        dtype=torch.long,
        device=device,
    )

    positions = torch.arange(
        len(chunks),
        dtype=torch.long,
        device=device,
    )

    for row, chunk in enumerate(
        chunks
    ):
        length = len(
            chunk.input_ids
        )

        input_ids[
            row,
            :length,
        ] = torch.tensor(
            chunk.input_ids,
            dtype=torch.long,
            device=device,
        )

        attention_mask[
            row,
            :length,
        ] = torch.tensor(
            chunk.attention_mask,
            dtype=torch.long,
            device=device,
        )

        gap_flags[
            row
        ] = chunk.gap_before

    return (
        input_ids,
        attention_mask,
        positions,
        gap_flags,
    )


def prepare_final_tensor(
    encoded: Dict[str, List[int]],
    device: torch.device,
) -> Tuple[
    torch.Tensor,
    torch.Tensor,
]:
    input_ids = torch.tensor(
        [
            encoded["input_ids"]
        ],
        dtype=torch.long,
        device=device,
    )

    attention_mask = torch.tensor(
        [
            encoded["attention_mask"]
        ],
        dtype=torch.long,
        device=device,
    )

    return (
        input_ids,
        attention_mask,
    )


# =============================================================================
# ROUTER WRAPPER
# =============================================================================

class ComplexityRouter:
    def __init__(
        self,
        model_dir: Path,
        device: str,
        dtype_name: str,
        threads: int,
    ):
        load_start = time.perf_counter()

        self.model_dir = Path(
            model_dir
        )

        if threads > 0:
            torch.set_num_threads(
                threads
            )

        self.device = torch.device(
            device
        )

        with open(
            self.model_dir
            / "router_config.json",
            "r",
            encoding="utf-8",
        ) as f:
            self.config = json.load(
                f
            )

        architecture = self.config.get(
            "architecture"
        )

        if architecture != "chunked_history_final_user_v2":
            raise ValueError(
                "This predictor requires a "
                "chunked_history_final_user_v2 checkpoint; "
                f"found architecture={architecture!r}"
            )

        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_dir
            / "tokenizer"
        )

        if self.tokenizer.pad_token_id is None:
            if self.tokenizer.eos_token_id is not None:
                self.tokenizer.pad_token = (
                    self.tokenizer.eos_token
                )
            else:
                raise ValueError(
                    "Tokenizer has no pad token or EOS token"
                )

        expected_vocab = self.config.get(
            "tokenizer_vocab_size"
        )

        if (
            expected_vocab is not None
            and len(self.tokenizer) != int(
                expected_vocab
            )
        ):
            raise ValueError(
                "Tokenizer vocabulary size does not match checkpoint: "
                f"tokenizer={len(self.tokenizer)}, "
                f"checkpoint={expected_vocab}"
            )

        encoder_config = AutoConfig.from_pretrained(
            self.model_dir
            / "encoder"
        )

        self.model = ChunkedHistoryComplexityRouter(
            encoder_config=encoder_config,
            tokenizer_size=len(
                self.tokenizer
            ),
            aggregator_layers=int(
                self.config[
                    "aggregator_layers"
                ]
            ),
            aggregator_heads=int(
                self.config[
                    "aggregator_heads"
                ]
            ),
            aggregator_dropout=float(
                self.config[
                    "aggregator_dropout"
                ]
            ),
            max_history_chunks=int(
                self.config[
                    "max_history_chunks"
                ]
            ),
        )

        checkpoint_dtype = self.config.get(
            "checkpoint_dtype",
            "bfloat16",
        )

        if dtype_name == "auto":
            if checkpoint_dtype == "bfloat16":
                dtype = torch.bfloat16
            else:
                dtype = torch.float32

        elif dtype_name == "bf16":
            dtype = torch.bfloat16

        elif dtype_name == "fp32":
            dtype = torch.float32

        else:
            raise ValueError(
                f"Unsupported dtype: {dtype_name}"
            )

        # IMPORTANT:
        # Cast the destination model BEFORE load_state_dict so a BF16
        # checkpoint remains BF16 rather than being silently copied into
        # FP32 parameters.
        self.model = self.model.to(
            device=self.device,
            dtype=dtype,
        )

        state = torch.load(
            self.model_dir
            / "router_model.pt",
            map_location="cpu",
            weights_only=True,
        )

        self.model.load_state_dict(
            state,
            strict=True,
        )

        self.model.eval()

        self.dtype = next(
            self.model.parameters()
        ).dtype

        self.role_start_tokens = (
            self.config.get(
                "role_start_tokens",
                DEFAULT_ROLE_START_TOKEN,
            )
        )

        self.role_cont_tokens = (
            self.config.get(
                "role_continuation_tokens",
                DEFAULT_ROLE_CONT_TOKEN,
            )
        )

        # Validate special-token availability now instead of during the first
        # request.
        router_marker_ids(
            self.tokenizer,
            self.role_start_tokens,
            self.role_cont_tokens,
        )

        load_ms = (
            time.perf_counter()
            - load_start
        ) * 1000

        print(
            json.dumps(
                {
                    "status": "model_loaded",
                    "device": str(
                        self.device
                    ),
                    "dtype": str(
                        self.dtype
                    ),
                    "threads": torch.get_num_threads(),
                    "load_time_ms": round(
                        load_ms,
                        2,
                    ),
                }
            ),
            file=sys.stderr,
            flush=True,
        )

    def predict(
        self,
        payload: Dict[str, Any],
    ) -> Dict[str, Any]:
        total_start = time.perf_counter()

        if not isinstance(
            payload,
            dict,
        ):
            raise ValueError(
                "Top-level JSON must be an object"
            )

        messages = payload.get(
            "messages"
        )

        validation_start = time.perf_counter()

        validate_role_sequence(
            messages
        )

        validation_ms = (
            time.perf_counter()
            - validation_start
        ) * 1000

        preprocessing_start = time.perf_counter()

        history_messages = (
            messages[:-1]
        )

        final_user_message = (
            messages[-1]
        )

        history_chunks = build_history_chunks(
            tokenizer=self.tokenizer,
            history_messages=history_messages,
            history_chunk_tokens=int(
                self.config[
                    "history_chunk_tokens"
                ]
            ),
            max_history_tokens=int(
                self.config[
                    "max_history_tokens"
                ]
            ),
            role_start_tokens=self.role_start_tokens,
            role_cont_tokens=self.role_cont_tokens,
        )

        final_encoded = encode_final_user(
            tokenizer=self.tokenizer,
            message=final_user_message,
            max_final_user_tokens=int(
                self.config[
                    "max_final_user_tokens"
                ]
            ),
        )

        preprocessing_ms = (
            time.perf_counter()
            - preprocessing_start
        ) * 1000

        tensor_start = time.perf_counter()

        (
            history_input_ids,
            history_attention_mask,
            history_positions,
            history_gap_flags,
        ) = pad_history_chunks(
            chunks=history_chunks,
            pad_token_id=int(
                self.tokenizer.pad_token_id
            ),
            device=self.device,
        )

        (
            final_input_ids,
            final_attention_mask,
        ) = prepare_final_tensor(
            final_encoded,
            self.device,
        )

        tensor_ms = (
            time.perf_counter()
            - tensor_start
        ) * 1000

        history_encoder_ms = 0.0
        final_encoder_ms = 0.0
        router_head_ms = 0.0

        with torch.inference_mode():
            if len(history_chunks) > 0:
                start = time.perf_counter()

                history_vectors = (
                    self.model.encode_sequences(
                        history_input_ids,
                        history_attention_mask,
                    )
                )

                history_encoder_ms = (
                    time.perf_counter()
                    - start
                ) * 1000

            else:
                hidden = (
                    self.model.encoder.config.hidden_size
                )

                history_vectors = torch.empty(
                    0,
                    hidden,
                    device=self.device,
                    dtype=self.dtype,
                )

            start = time.perf_counter()

            final_vectors = (
                self.model.encode_sequences(
                    final_input_ids,
                    final_attention_mask,
                )
            )

            final_encoder_ms = (
                time.perf_counter()
                - start
            ) * 1000

            start = time.perf_counter()

            score = (
                self.model.score_embeddings(
                    history_vectors=history_vectors,
                    final_vectors=final_vectors,
                    history_positions=history_positions,
                    history_gap_flags=history_gap_flags,
                    history_count=len(
                        history_chunks
                    ),
                )
                .float()
                .item()
            )

            router_head_ms = (
                time.perf_counter()
                - start
            ) * 1000

        inference_ms = (
            history_encoder_ms
            + final_encoder_ms
            + router_head_ms
        )

        total_ms = (
            time.perf_counter()
            - total_start
        ) * 1000

        history_lengths = [
            len(
                chunk.input_ids
            )
            for chunk in history_chunks
        ]

        final_tokens = len(
            final_encoded[
                "input_ids"
            ]
        )

        return {
            "complexity": round(
                score,
                4,
            ),
            "latency_ms": round(
                total_ms,
                2,
            ),
            "inference_ms": round(
                inference_ms,
                2,
            ),
            "history_encoder_ms": round(
                history_encoder_ms,
                2,
            ),
            "final_encoder_ms": round(
                final_encoder_ms,
                2,
            ),
            "router_head_ms": round(
                router_head_ms,
                2,
            ),
            "preprocessing_ms": round(
                preprocessing_ms,
                2,
            ),
            "tensor_build_ms": round(
                tensor_ms,
                2,
            ),
            "validation_ms": round(
                validation_ms,
                2,
            ),
            "device": str(
                self.device
            ),
            "dtype": str(
                self.dtype
            ),
            "messages": len(
                messages
            ),
            "history_chunks": len(
                history_chunks
            ),
            "history_chunk_lengths": (
                history_lengths
            ),
            "history_padded_tokens": (
                len(history_chunks)
                * max(
                    history_lengths,
                    default=0,
                )
            ),
            "history_actual_tokens": sum(
                history_lengths
            ),
            "final_tokens": (
                final_tokens
            ),
        }


# =============================================================================
# INPUT MODES
# =============================================================================

def run_file(
    router: ComplexityRouter,
    input_path: Path,
) -> None:
    with open(
        input_path,
        "r",
        encoding="utf-8",
    ) as f:
        payload = json.load(
            f
        )

    result = router.predict(
        payload
    )

    result = {
        "file": input_path.name,
        **result,
    }

    print(
        json.dumps(
            result,
            indent=2,
        ),
        flush=True,
    )


def run_interactive(
    router: ComplexityRouter,
) -> None:
    print(
        "",
        file=sys.stderr,
    )

    print(
        "Ready.",
        file=sys.stderr,
    )

    print(
        "Paste a complete JSON object.",
        file=sys.stderr,
    )

    print(
        "Pretty-printed / multi-line JSON is supported.",
        file=sys.stderr,
    )

    print(
        "The request is evaluated automatically as soon as "
        "the JSON object is complete.",
        file=sys.stderr,
    )

    print(
        "Press Ctrl+D to exit.",
        file=sys.stderr,
    )

    print(
        "",
        file=sys.stderr,
    )

    decoder = json.JSONDecoder()

    buffer = ""

    for line in sys.stdin:
        buffer += line

        while True:
            stripped = buffer.lstrip()

            if not stripped:
                buffer = ""
                break

            leading_whitespace = (
                len(buffer)
                - len(stripped)
            )

            try:
                payload, end_index = (
                    decoder.raw_decode(
                        stripped
                    )
                )

            except json.JSONDecodeError:
                # Most likely a valid-but-incomplete multi-line JSON object.
                break

            try:
                result = router.predict(
                    payload
                )

            except Exception as exc:
                result = {
                    "error": str(
                        exc
                    )
                }

            print(
                json.dumps(
                    result,
                    indent=2,
                ),
                flush=True,
            )

            consumed = (
                leading_whitespace
                + end_index
            )

            buffer = buffer[
                consumed:
            ]

    if buffer.strip():
        try:
            payload = json.loads(
                buffer
            )

            result = router.predict(
                payload
            )

        except Exception as exc:
            result = {
                "error": str(
                    exc
                )
            }

        print(
            json.dumps(
                result,
                indent=2,
            ),
            flush=True,
        )


# =============================================================================
# CLI
# =============================================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "Predict complexity with the chunked-history ModernBERT router"
        )
    )

    parser.add_argument(
        "--model-dir",
        required=True,
        help=(
            "Path to the best/ checkpoint directory"
        ),
    )

    parser.add_argument(
        "--input",
        default=None,
        help=(
            "Optional JSON file. If omitted, interactive mode keeps the "
            "model resident in RAM."
        ),
    )

    parser.add_argument(
        "--device",
        default=(
            "cuda"
            if torch.cuda.is_available()
            else "cpu"
        ),
    )

    parser.add_argument(
        "--dtype",
        choices=[
            "auto",
            "bf16",
            "fp32",
        ],
        default="auto",
        help=(
            "auto uses checkpoint_dtype from router_config.json. "
            "The v2 trainer saves BF16."
        ),
    )

    parser.add_argument(
        "--threads",
        type=int,
        default=0,
        help=(
            "PyTorch intra-op CPU threads. 0 leaves the PyTorch default."
        ),
    )

    return parser.parse_args()


def main():
    args = parse_args()

    model_dir = Path(
        args.model_dir
    )

    router = ComplexityRouter(
        model_dir=model_dir,
        device=args.device,
        dtype_name=args.dtype,
        threads=args.threads,
    )

    if args.input:
        run_file(
            router,
            Path(
                args.input
            ),
        )

    else:
        run_interactive(
            router
        )


if __name__ == "__main__":
    main()