| """Hugging Face Inference Endpoints custom handler for Banyan Sapling. |
| |
| Loads the model exactly like the reference `inference.py` shipped with the |
| weights package (verified against the H200 reference runtime), with two |
| serving-only deviations: |
| - skips the strict package-version assertion loop (requirements.txt pins them) |
| - skips the full 46 GB sha256 pre-hash (the Hub delivers checksummed storage) |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| import sys |
| import time |
| from typing import Any, Dict |
|
|
| CURRENT_DIR = Path(__file__).resolve().parent |
| if str(CURRENT_DIR) not in sys.path: |
| sys.path.insert(0, str(CURRENT_DIR)) |
|
|
| import os |
| os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8') |
|
|
|
|
| def load_model_for_endpoint(root: Path, device: str = "cuda"): |
| root = Path(root).resolve() |
| 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) |
| 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())) |
|
|
| manifest_path = root / 'distribution_manifest.json' |
| if manifest_path.exists(): |
| manifest = json.loads(manifest_path.read_text()) |
| else: |
| manifest = {'source_release_sha256': '', 'files': {'model.safetensors': ''}} |
| return model, tokenizer, calibrator, manifest |
|
|
|
|
| def predict_for_endpoint(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.get('source_release_sha256', ''), |
| 'weights_sha256': manifest.get('files', {}).get('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}} |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = "") -> None: |
| model_dir = Path(path) if path else CURRENT_DIR |
| self.model, self.tokenizer, self.calibrator, self.manifest = load_model_for_endpoint( |
| model_dir, device="cuda" |
| ) |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
| payload = data.get("inputs", data) |
| if isinstance(payload, str): |
| try: |
| payload = json.loads(payload) |
| except Exception: |
| return {"error": "Invalid payload: expected JSON object with policy, evidence, and claim strings."} |
|
|
| if not isinstance(payload, dict): |
| return {"error": "Payload must be a dictionary with policy, evidence, and claim fields."} |
|
|
| for field in ("policy", "evidence", "claim"): |
| val = payload.get(field) |
| if not val or not str(val).strip(): |
| return {"error": f"Missing required non-empty string field: '{field}'"} |
|
|
| return predict_for_endpoint(self.model, self.tokenizer, self.calibrator, self.manifest, payload) |
|
|
