"""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)