"""Banyan Sapling inference for a downloaded, checksummed weights package.""" from __future__ import annotations import argparse import hashlib import importlib.metadata import json import os from pathlib import Path import sys import time os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8') os.environ.setdefault('OMP_NUM_THREADS', '8') os.environ.setdefault('MKL_NUM_THREADS', '8') sys.dont_write_bytecode = True def sha256(path): digest = hashlib.sha256() with Path(path).open('rb') as stream: for block in iter(lambda: stream.read(8 * 1024 * 1024), b''): digest.update(block) return digest.hexdigest() def verify_files(root): """Verify listed assets; Hub download metadata may coexist with the package.""" root = Path(root).resolve() manifest = json.loads((root / 'distribution_manifest.json').read_text()) if manifest.get('format_version') != 1 or manifest.get('model_id') != 'banyan-sapling': raise ValueError('Unsupported distribution manifest') for name, expected in manifest['files'].items(): path = (root / name).resolve() if not path.is_relative_to(root) or not path.is_file(): raise ValueError(f'Missing or invalid asset: {name}') if sha256(path) != expected: raise ValueError(f'Checksum mismatch: {name}') required = {'model.safetensors', 'config.json', 'architecture.json', 'calibration.json', 'runtime_versions.json', 'runtime/model.py', 'runtime/backbone.py', 'runtime/conformal_calibrator.py', 'runtime/interaction_head.py', 'runtime/verifier_batching.py', 'runtime/__init__.py', 'tokenizer/tokenizer.json', 'tokenizer/tokenizer_config.json', 'tokenizer/chat_template.jinja', 'inference.py'} if not required.issubset(manifest['files']): raise ValueError('Incomplete distribution manifest') return manifest def load_model(root, device='cuda'): root = Path(root).resolve() manifest = verify_files(root) versions = json.loads((root / 'runtime_versions.json').read_text()) for package, version in versions['packages'].items(): if importlib.metadata.version(package) != version: raise ValueError(f'Runtime requires {package}=={version}; use requirements.txt') if str(root) not in sys.path: sys.path.insert(0, str(root)) import torch from safetensors.torch import load_file # The text-only runtime does not require torchvision or torchaudio. sys.modules.setdefault('torchvision', None) sys.modules.setdefault('torchaudio', None) from transformers import AutoConfig, AutoModel, AutoTokenizer from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeRotaryEmbedding from runtime.backbone import attach_lora, set_shared from runtime.model import FullModel, identity, seed_all, tensor_hash from runtime.conformal_calibrator import ConformalCalibrator seed_all(1707) if device != 'cuda' or not torch.cuda.is_available(): raise ValueError('This reference runtime requires an NVIDIA CUDA device') config_data = json.loads((root / 'config.json').read_text()) architecture = json.loads((root / 'architecture.json').read_text()) expert_counts = [config_data[k] for k in ('num_local_experts', 'num_experts') if k in config_data] if (not expert_counts or any(n != 96 for n in expert_counts) or config_data.get('model_type') != 'qwen3_moe' or identity(config_data) != architecture['hf_config_sha256'] or architecture.get('modify_rope') is not False or architecture.get('factorized_lora') is not True): raise ValueError('Unsupported model architecture') config = AutoConfig.from_pretrained(root, local_files_only=True) with torch.device('meta'): backbone = AutoModel.from_config(config, dtype=torch.bfloat16, attn_implementation='sdpa', experts_implementation='grouped_mm') attach_lora(backbone, architecture['config']) if architecture['shared_ids'] is not None: set_shared(backbone, architecture['shared_ids']) model = FullModel(backbone) state = load_file(str(root / 'model.safetensors'), device=device) model.load_state_dict(state, strict=True, assign=True) del state model.backbone.rotary_emb = Qwen3MoeRotaryEmbedding(config, device=device) rope = model.backbone.rotary_emb scaling = rope.attention_scaling fingerprint = {'inv_freq': tensor_hash({'inv_freq': rope.inv_freq}), 'attention_scaling': scaling.item() if hasattr(scaling, 'item') else scaling} if fingerprint != architecture['rotary_fingerprint']: raise ValueError('Native positional encoding does not match the checkpoint') if any(t.is_meta for t in list(model.parameters()) + list(model.buffers())): raise ValueError('Incomplete model materialization') for name, parameter in model.named_parameters(): expected = torch.float32 if name.startswith('head.') or name.endswith(('.A', '.B')) or '.mlp.gate.' in name else torch.bfloat16 if parameter.dtype != expected: raise ValueError(f'Unexpected precision: {name}') model.requires_grad_(False) model.eval() tokenizer = AutoTokenizer.from_pretrained(root / 'tokenizer', local_files_only=True) calibrator = ConformalCalibrator.from_dict(json.loads((root / 'calibration.json').read_text())) return model, tokenizer, calibrator, manifest def predict(model, tokenizer, calibrator, manifest, request): import torch from runtime.backbone import encode from runtime.model import infer row = encode(tokenizer, request, 16384) pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id start = time.perf_counter() logits = infer(model, [row], pad)[0] detail = calibrator.predict_detailed(logits.reshape(1, 4), empty_set_policy='rounded_argmax')[0] probabilities = torch.softmax(logits.double() / calibrator.temperature, -1).tolist() return {'model': 'banyan-sapling', 'release_stage': 'research-preview', 'source_release_sha256': manifest['source_release_sha256'], 'weights_sha256': manifest['files']['model.safetensors'], 'logits': logits.tolist(), **detail, 'probabilities': dict(zip(calibrator.LABELS, probabilities)), 'usage': {'input_tokens': len(row['input_ids']), 'generated_tokens': 0}, 'timing_ms': {'model': (time.perf_counter() - start) * 1000}} def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--model-dir', type=Path, default=Path(__file__).resolve().parent) parser.add_argument('--input', type=Path, help='JSON file with policy, evidence, and claim') parser.add_argument('--check-files', action='store_true', help='Verify package checksums without loading the model') args = parser.parse_args() if args.check_files: manifest = verify_files(args.model_dir) print(json.dumps({'model': manifest['model_id'], 'verified_files': len(manifest['files'])})) return if args.input is None: parser.error('--input is required for inference') request = json.loads(args.input.read_text()) model, tokenizer, calibrator, manifest = load_model(args.model_dir) print(json.dumps(predict(model, tokenizer, calibrator, manifest, request), indent=2, allow_nan=False)) if __name__ == '__main__': main()