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