| """Standalone Nima inference. Imported MLX and weights only inside load/predict.""" |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import time |
| from encoding import RULES, choices, encode |
| from head import create_head, pack |
|
|
| HERE = Path(__file__).resolve().parent |
|
|
|
|
| def digest(path): |
| h = hashlib.sha256() |
| with Path(path).open('rb') as handle: |
| for block in iter(lambda: handle.read(1024 * 1024), b''): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def load(release_dir=HERE, local_files_only=False): |
| import mlx.core as mx |
| import mlx.nn as nn |
| from mlx.utils import tree_flatten |
| from mlx_lm import load as load_qwen |
| from huggingface_hub import snapshot_download |
|
|
| folder = Path(release_dir) |
| config = json.loads((folder / 'config.json').read_text()) |
| architecture = json.loads((folder / 'architecture.json').read_text()) |
| head_path = folder / config['head_file'] |
| if digest(head_path) != config['head_sha256']: |
| raise ValueError('Nima head checksum mismatch') |
| snapshot = Path(snapshot_download( |
| config['base_model'], revision=config['base_revision'], |
| allow_patterns=list(config['base_file_sha256']), |
| local_files_only=local_files_only, |
| )) |
| for name, expected in config['base_file_sha256'].items(): |
| if digest(snapshot / name) != expected: |
| raise ValueError(f'Frozen backbone file checksum mismatch: {name}') |
| model, tokenizer, base_config = load_qwen(str(snapshot), return_config=True) |
| if model.model_type != 'qwen3' or base_config.get('hidden_size') != architecture['hidden_size']: |
| raise ValueError('Expected the pinned Qwen3 architecture') |
| if base_config.get('quantization', {}).get('bits') != 4: |
| raise ValueError('Expected a 4-bit backbone') |
| if not any(isinstance(module, nn.QuantizedLinear) for _, module in model.named_modules()): |
| raise ValueError('Quantized backbone layers missing') |
| model.freeze() |
| model.eval() |
| if tree_flatten(model.trainable_parameters()): |
| raise ValueError('Backbone must be completely frozen') |
| head = create_head(architecture) |
| head.load_weights(str(head_path)) |
| head.eval() |
| if sum(value.size for _, value in tree_flatten(head.parameters())) != config['head_parameters']: |
| raise ValueError('Head parameter count mismatch') |
| mx.eval(model.parameters(), head.parameters()) |
| return Nima(model, tokenizer, head, architecture, config) |
|
|
|
|
| class Nima: |
| def __init__(self, model, tokenizer, head, architecture, config): |
| self.model, self.tokenizer, self.head = model, tokenizer, head |
| self.architecture, self.config = architecture, config |
|
|
| def predict(self, request): |
| import mlx.core as mx |
| if not isinstance(request.get('prompt'), str) or not request['prompt'].strip(): |
| raise ValueError('A nonempty prompt is required') |
| if not isinstance(request.get('context'), dict) or not {'current_date', 'timezone'} <= request['context'].keys(): |
| raise ValueError('context requires current_date and timezone') |
| names = choices(request) |
| if 'no_tool' not in names: |
| raise ValueError('Include an explicit no_tool choice to allow abstention') |
| case = {**request, 'id': 'custom', 'target': names[0]} |
| started = time.perf_counter() |
| encoded = encode(case, names, self.tokenizer, RULES, self.architecture['max_tokens']) |
| |
| hidden = self.model.model(mx.array([encoded['token_ids']], dtype=mx.int32), cache=None) |
| evidence = mx.stop_gradient(hidden[0, mx.array(encoded['evidence_positions']), :].astype(mx.float32)) |
| candidates = mx.stack([mx.mean(hidden[0, span['start']:span['end'], :], axis=0) for span in encoded['candidate_spans']]) |
| candidates = mx.stop_gradient(candidates.astype(mx.float32)) |
| mx.eval(evidence, candidates) |
| q, e, qv, ev, _ = pack([{'candidates': candidates, 'evidence': evidence, 'label': 0}]) |
| logits = self.head(q, e, qv, ev)[0, :len(names)] |
| probabilities = mx.softmax(logits / self.config['temperature']) |
| mx.eval(logits, probabilities) |
| elapsed = time.perf_counter() - started |
| scores = dict(zip(names, probabilities.tolist())) |
| winner = min(scores, key=lambda name: (-scores[name], name)) |
| return {'prediction': winner, 'probabilities': scores, |
| 'logits': dict(zip(names, logits.tolist())), 'latency_seconds': elapsed, |
| 'input_tokens': len(encoded['token_ids']), 'qwen_forward_passes': 1, |
| 'note': 'Probabilities are uncalibrated and relative to the supplied choices; no tool executes.'} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--input', type=Path, required=True) |
| parser.add_argument('--local-files-only', action='store_true') |
| args = parser.parse_args() |
| request = json.loads(args.input.read_text()) |
| router = load(local_files_only=args.local_files_only) |
| print(json.dumps(router.predict(request), indent=2, ensure_ascii=False)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
