#!/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()