runtime.py

"""Inference compatible with Untitled54 (2).ipynb's existing checkpoint.

No training, model upload or generation API calls occur in this module.
"""
import json
import math
import os
from pathlib import Path
import torch
from torch import nn
from transformers import AutoConfig, AutoModel, AutoTokenizer
from huggingface_hub import HfApi, hf_hub_download

BACKBONE = 'jhu-clsp/mmBERT-base'
MAX_OPTIONS = 4
MAX_LENGTH = 256
VALID_SCORE_LEVELS = 5  # The notebook's labels are 0..4, although its head has 10 outputs.


class DynamicDecisionBaseline(nn.Module):
    """Exact parameter names/shapes of the user's saved state_dict."""
    def __init__(self, config, max_score_levels=10):
        super().__init__()
        # ModernBERT can compile internally even without explicit torch.compile.
        # Disable that path for ZeroGPU/CPU and use the portable attention path.
        config.reference_compile = False
        self.encoder = AutoModel.from_config(config, attn_implementation='eager',
                                             trust_remote_code=False)
        hidden = config.hidden_size
        self.choice_scorer = nn.Linear(hidden, 1)
        self.noul_head = nn.Linear(hidden, 1)
        self.score_head = nn.Linear(hidden, max_score_levels)

    def forward(self, input_ids, attention_mask):
        # One representation per candidate; no padding candidates needed at inference.
        cls = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
        # In eval mode these are identical to the notebook's three independent
        # encoder calls: the auxiliary calls both reused candidate zero.
        return {'choice_logits': self.choice_scorer(cls).squeeze(-1),
                'noul_logit': self.noul_head(cls[0:1]).reshape(()),
                'score_logits': self.score_head(cls[0:1]).squeeze(0)}


def parse_criteria(text):
    def unique_pairs(pairs):
        result = {}
        for key, value in pairs:
            if key in result:
                raise ValueError(f'Duplicate option key: {key}')
            result[key] = value
        return result
    try:
        criteria = json.loads(text, object_pairs_hook=unique_pairs)
    except json.JSONDecodeError as exc:
        raise ValueError(f'Criteria must be valid JSON: {exc.msg}') from exc
    if not isinstance(criteria, dict) or not 2 <= len(criteria) <= MAX_OPTIONS:
        raise ValueError('Criteria must be a JSON object with 2 to 4 options.')
    for key, description in criteria.items():
        if not isinstance(key, str) or not key.strip():
            raise ValueError('Every option needs a nonempty string key.')
        if not isinstance(description, str) or not description.strip():
            raise ValueError('Every option description must be a nonempty string.')
    return criteria


def prepare_inputs(tokenizer, state, question, criteria):
    if not isinstance(state, str) or not state.strip():
        raise ValueError('Enter a nonempty state.')
    if not isinstance(question, str) or not question.strip():
        raise ValueError('Enter a nonempty routing question.')
    # Preserve EXACT training serialization. [SEP] is literal text here;
    # changing to a new prompt/paired-tokenizer format requires retraining.
    texts = [f'{state} [SEP] {question} [SEP] {key}: {description}'
             for key, description in criteria.items()]
    tokenized = tokenizer(texts, truncation=False, padding=False,
                          return_attention_mask=True, return_token_type_ids=False)
    lengths = [len(ids) for ids in tokenized['input_ids']]
    if max(lengths) > MAX_LENGTH:
        raise ValueError(f'Input reaches {max(lengths)} tokens; this checkpoint was trained with '
                         f'{MAX_LENGTH}. Shorten the state/question/options. Nothing was silently truncated.')
    # Fixed padding matches the notebook. Only real candidates are encoded.
    encoded = tokenizer.pad(tokenized, padding='max_length', max_length=MAX_LENGTH,
                            return_tensors='pt')
    return {k: encoded[k] for k in ('input_ids', 'attention_mask')}, lengths


def format_result(outputs, criteria, temperature, metadata):
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError('Temperature must be positive and finite.')
    keys = list(criteria)
    choice_logits = outputs['choice_logits'].detach().float().cpu()
    noul_logit = outputs['noul_logit'].detach().float().cpu()
    raw_score = outputs['score_logits'].detach().float().cpu()
    if not all(torch.isfinite(x).all() for x in (choice_logits, noul_logit, raw_score)):
        raise RuntimeError('Model produced nonfinite logits; inspect checkpoint and dependencies.')
    p = torch.softmax(choice_logits/temperature, dim=-1)
    best = int(p.argmax())
    raw_p = torch.softmax(raw_score, dim=-1)
    # Renormalize over the only label levels that actually existed in training.
    score_p = torch.softmax(raw_score[:VALID_SCORE_LEVELS], dim=-1)
    warnings = [
        'Escalation and severity are experimental fixed tasks; they are not general Noul/Score questions.',
        'Both auxiliary heads use the first routing option, as in training; changing that option can change their outputs.',
        'Temperature scaling applies to routing only; it does not establish correctness or calibrate the auxiliary heads.'
    ]
    if not metadata.get('calibration_loaded'):
        warnings.append('No calibration_config.json was available; routing uses temperature 1.0.')
    return {
        'routing_decision': {
            'choice': keys[best], 'description': criteria[keys[best]],
            'max_probability': float(p[best]),
            'probabilities': {k: float(v) for k, v in zip(keys, p)},
            'temperature': temperature},
        'human_escalation': {'probability': float(torch.sigmoid(noul_logit)),
                             'calibrated': False, 'experimental': True},
        'severity_score': {'score_mode': int(score_p.argmax()),
                           'expected_score': float((torch.arange(VALID_SCORE_LEVELS)*score_p).sum()),
                           'probabilities': {str(i):float(v) for i,v in enumerate(score_p)},
                           'raw_probability_mass_outside_0_to_4': float(raw_p[VALID_SCORE_LEVELS:].sum()),
                           'calibrated': False, 'experimental': True},
        'metadata': {**metadata, 'auxiliary_input_option': keys[0]},
        'limitations': warnings}


class Runtime:
    def __init__(self, model, tokenizer, temperature, metadata, device):
        self.device = torch.device(device)
        self.model = model.to(self.device).eval()
        self.tokenizer = tokenizer
        self.temperature = temperature
        self.metadata = metadata

    @torch.inference_mode()
    def predict(self, encoded, criteria, lengths):
        inputs = {k:v.to(self.device) for k,v in encoded.items()}
        outputs = self.model(**inputs)
        return format_result(outputs, criteria, self.temperature,
                             {**self.metadata, 'device':str(self.device), 'input_token_lengths':lengths})


def load_runtime(device='cpu'):
    repo_id = os.environ.get('MODEL_REPO_ID', 'TD-jayadeera/sroute')
    requested_revision = os.environ.get('MODEL_REVISION', 'main')
    local_dir = os.environ.get('SROUTE_LOCAL_MODEL_DIR')
    token = os.environ.get('HF_TOKEN') or False
    if local_dir:
        root = Path(local_dir)
        def get_file(name): return str(root/name)
        files = {p.relative_to(root).as_posix() for p in root.rglob('*') if p.is_file()}
        revision = 'local'
    else:
        try:
            info = HfApi().model_info(repo_id, revision=requested_revision, token=token)
        except Exception as exc:
            raise RuntimeError('Cannot read the model repository. Set the Space HF_TOKEN secret to a '
                               'read token with access to TD-jayadeera/sroute, and verify MODEL_REPO_ID.') from exc
        revision = info.sha
        files = {f.rfilename for f in info.siblings}
        def get_file(name):
            return hf_hub_download(repo_id=repo_id,filename=name,revision=revision,token=token)
    if 'model_weights.pt' not in files:
        raise RuntimeError('Repository root is missing model_weights.pt. Upload the actual saved model folder contents.')
    tokenizer_subfolder = '' if 'tokenizer_config.json' in files else 'tokenizer'
    tokenizer_source = str(Path(local_dir)/tokenizer_subfolder) if local_dir else repo_id
    tok_kwargs = {'local_files_only':True} if local_dir else {
        'revision':revision,'token':token,'subfolder':tokenizer_subfolder}
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_source,trust_remote_code=False,**tok_kwargs)
    # Prefer a saved training backbone config if present. The uploaded notebook
    # did not save it; use the public mmBERT config as the compatibility fallback.
    if 'encoder_config.json' in files:
        from transformers import ModernBertConfig
        config = ModernBertConfig.from_json_file(get_file('encoder_config.json'))
        config_source = 'saved_encoder_config'
    else:
        config = AutoConfig.from_pretrained(BACKBONE,trust_remote_code=False,
                  revision=os.environ.get('BACKBONE_REVISION','main'),token=False)
        config_source = 'public_mmbert_config_fallback'
    weights_path = get_file('model_weights.pt')
    # This checkpoint is a plain tensor state_dict. Never use weights_only=False.
    weights = torch.load(weights_path, map_location='cpu', weights_only=True, mmap=True)
    if not isinstance(weights, dict) or 'score_head.weight' not in weights:
        raise RuntimeError('Expected a plain DynamicDecisionBaseline state_dict, not a pickled model or wrapped checkpoint.')
    levels = weights['score_head.weight'].shape[0]
    if levels != 10:
        raise RuntimeError(f'Expected the notebook 10-output score head; found {levels}. Verify checkpoint architecture.')
    model = DynamicDecisionBaseline(config, max_score_levels=levels)
    model.load_state_dict(weights, strict=True)
    del weights
    embedding_count = model.encoder.get_input_embeddings().num_embeddings
    if max(tokenizer.get_vocab().values()) >= embedding_count:
        raise RuntimeError('Tokenizer IDs exceed embedding vocabulary. Use the tokenizer saved with these weights.')
    temperature = 1.0
    calibration_loaded = 'calibration_config.json' in files
    if calibration_loaded:
        with open(get_file('calibration_config.json'),encoding='utf-8') as f:
            temperature = float(json.load(f)['optimal_temperature'])
        if not math.isfinite(temperature) or temperature <= 0:
            raise RuntimeError('Invalid optimal_temperature in calibration_config.json.')
    metadata = {'model_repo':repo_id,'model_revision':revision,'backbone':BACKBONE,
                'config_source':config_source,'calibration_loaded':calibration_loaded,
                'input_format':'notebook_v2_compatible','max_tokens':MAX_LENGTH,
                'score_labels':[0,1,2,3,4],'attention':'eager','reference_compile':False}
    return Runtime(model,tokenizer,temperature,metadata,device)