| """Native MLX inference from a checkpoint or self-contained xDecision GGUF. |
| |
| Weights remain resident in the requested precision. Q8_0 is decoded once on |
| load; inference uses dense MLX operations, not quantized matrix multiplication. |
| Prompt construction and result decoding use the pinned Laya contract. |
| """ |
| import json |
| import math |
| import os |
| import threading |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import numpy as np |
| from laya.agent import Agent |
| from laya.common import collapsed_options, _reuse_question_tokens |
| from tokenizers import Tokenizer |
|
|
| from .mlx_model import LayaMLX |
|
|
|
|
| class LocalTokenizer: |
| """Read-only Rust tokenizer adapter for Laya's exact sequence builder.""" |
| def __init__(self, text, config): |
| self.backend = Tokenizer.from_str(text) |
| self.backend.no_padding() |
| self.backend.no_truncation() |
| for name in ('cls', 'sep', 'mask', 'pad'): |
| token = config.get(name + '_token') |
| if isinstance(token, dict): |
| token = token['content'] |
| token_id = self.backend.token_to_id(token) if isinstance(token, str) else None |
| if token_id is None: |
| raise ValueError(f'Missing {name} token') |
| setattr(self, name + '_token', token) |
| setattr(self, name + '_token_id', token_id) |
|
|
| def __call__(self, text, *, add_special_tokens=False, truncation=False, max_length=None): |
| ids = self.backend.encode(text, add_special_tokens=add_special_tokens).ids |
| if truncation: |
| ids = ids[:max_length] |
| return {'input_ids': ids} |
|
|
|
|
| def validate_temperatures(cfg): |
| """Use fitted values unchanged; reject invalid values instead of clamping.""" |
| temperatures = cfg.get('temperature', [1.0, 1.0, 1.0]) |
| buckets = cfg.get('temperature_by_options', {}) |
| if len(temperatures) != 3 or any( |
| not math.isfinite(float(t)) or float(t) <= 0 |
| for t in [*temperatures, *buckets.values()]): |
| raise ValueError('Temperatures must be finite and positive, with three question types') |
| return [float(t) for t in temperatures], {k: float(t) for k, t in buckets.items()} |
|
|
|
|
| def read_checkpoint(path, dtype): |
| """Return metadata and resident arrays without writing restored weights to disk.""" |
| path = Path(path) |
| if path.is_dir(): |
| cfg = json.loads((path/'rl_agent_config.json').read_text()) |
| enc = json.loads((path/'encoder/config.json').read_text()) |
| tok_text = (path/'tokenizer/tokenizer.json').read_text() |
| tok_cfg = json.loads((path/'tokenizer/tokenizer_config.json').read_text()) |
| weights = mx.load(str(path/'model.safetensors')) |
| weights = {k: v.astype(dtype) for k, v in weights.items()} |
| else: |
| import gguf |
| from .gguf_io import _decoded_tensor |
| reader = gguf.GGUFReader(path) |
| arch = reader.get_field('general.architecture') |
| if arch is None or arch.contents() != 'xdecision': |
| raise ValueError('Expected a complete xDecision GGUF') |
| def field(name): |
| value = reader.get_field('xdecision.' + name) |
| if value is None: |
| raise ValueError(f'Missing embedded {name}') |
| return value.contents() |
| cfg = json.loads(field('agent_config')) |
| enc = json.loads(field('encoder_config')) |
| tok_text = field('tokenizer_json') |
| tok_cfg = json.loads(field('tokenizer_config')) |
| weights = {} |
| for tensor in reader.tensors: |
| if tensor.name in weights: |
| raise ValueError(f'Duplicate tensor {tensor.name}') |
| |
| value = _decoded_tensor(tensor).astype( |
| np.float32 if tensor.name == 'temperature' else np.float16) |
| weights[tensor.name] = mx.array(value).astype(dtype) |
| mx.eval(weights) |
| if 'temperature' not in weights or weights['temperature'].shape != (3,): |
| raise ValueError('Missing or invalid temperature tensor') |
| weights.pop('temperature') |
| return cfg, enc, tok_text, tok_cfg, weights |
|
|
|
|
| def collate(items, pad_id): |
| n = len(items) |
| length = max(len(row['ids']) for row in items) |
| count = max(2, max(len(row['markers']) for row in items)) |
| ids = np.full((n, length), pad_id, dtype=np.int32) |
| att = np.zeros((n, length), dtype=np.bool_) |
| mpos = np.zeros((n, count), dtype=np.int32) |
| mmask = np.zeros((n, count), dtype=np.bool_) |
| for i, row in enumerate(items): |
| ids[i, :len(row['ids'])] = row['ids'] |
| att[i, :len(row['ids'])] = True |
| mpos[i, :len(row['markers'])] = row['markers'] |
| mmask[i, :len(row['markers'])] = True |
| return ids, att, mpos, mmask, np.array([r['qtype'] for r in items], dtype=np.int32) |
|
|
|
|
| class MLXModel: |
| def __init__(self, path, device=None, *, dtype='float16', batch_size=16): |
| if dtype not in ('float16', 'float32'): |
| raise ValueError('MLX inference dtype must be float16 or float32') |
| if dtype == 'float32' and os.environ.get('MLX_ENABLE_TF32') != '0': |
| raise ValueError('For strict float32, launch Python with MLX_ENABLE_TF32=0 before importing MLX') |
| if device not in (None, 'auto', 'gpu', 'metal', 'mps', 'cpu'): |
| raise ValueError('MLX requires an Apple GPU or CPU; use backend="torch" for CUDA') |
| if type(batch_size) is not int or batch_size < 1: |
| raise ValueError('batch_size must be a positive integer') |
| self.device = mx.cpu if device == 'cpu' else mx.gpu |
| if self.device == mx.gpu and not mx.metal.is_available(): |
| raise RuntimeError('MLX Metal device is not available') |
| self.dtype = {'float16': mx.float16, 'float32': mx.float32}[dtype] |
| self.batch_size = batch_size |
| self._lock = threading.RLock() |
| with mx.stream(self.device): |
| self.cfg, enc, text, tc, weights = read_checkpoint(path, self.dtype) |
| self.temperature, self.temperature_by_options = validate_temperatures(self.cfg) |
| self.lang_temperatures = {} |
| if self.cfg.get('lang_temperatures'): |
| raise ValueError('Language-specific calibration is not supported by this MLX entry point') |
| if (enc.get('model_type') != 'modernbert' or enc.get('hidden_activation', 'gelu') != 'gelu' |
| or any(enc.get(k, False) for k in ('norm_bias', 'mlp_bias', 'attention_bias'))): |
| raise ValueError('Unsupported encoder configuration') |
| if any(v.get('rope_type', 'default') != 'default' |
| for v in enc.get('rope_parameters', {}).values()): |
| raise ValueError('Scaled RoPE is not supported') |
| self.tok = LocalTokenizer(text, tc) |
| self.model = LayaMLX(enc, self.cfg['head_layers'], len(self.cfg.get('act_costs', {}))+1, |
| hidden_dtype=self.dtype) |
| self.model.load_weights(list(weights.items()), strict=True) |
| self.model.eval() |
| mx.eval(self.model.parameters()) |
|
|
| @_reuse_question_tokens |
| def prepare(self, state, questions): |
| if not isinstance(questions, dict): |
| raise ValueError('questions must be a dictionary') |
| for qid, question in questions.items(): |
| Agent._check_question(qid, question) |
| internal = {qid: Agent._to_internal(q) for qid, q in questions.items()} |
| items = Agent._encode_state(self, state, list(questions), internal) if questions else [] |
| return items, internal |
|
|
| def forward(self, items): |
| with mx.stream(self.device): |
| batch = [mx.array(v) for v in collate(items, self.tok.pad_token_id)] |
| logits, act = self.model(*batch, cdt=self.dtype, train=False) |
| probabilities = mx.softmax(act.astype(mx.float32), axis=-1) |
| mx.eval(logits, probabilities) |
| logits, probabilities = np.asarray(logits), np.asarray(probabilities) |
| if not np.isfinite(logits).all() or not np.isfinite(probabilities).all(): |
| raise FloatingPointError('Non-finite MLX result; restart with MLX_ENABLE_TF32=0 and dtype="float32"') |
| return logits, probabilities |
|
|
| def predict(self, state, questions): |
| with self._lock: |
| if self.model is None: |
| raise RuntimeError('Model is closed') |
| items, internal = self.prepare(state, questions) |
| qids, answers = list(questions), {} |
| for start in range(0, len(items), self.batch_size): |
| chunk = items[start:start+self.batch_size] |
| logits, act = self.forward(chunk) |
| answers.update(Agent._decode_answers(self, logits, act, chunk, |
| qids[start:start+self.batch_size], internal, 0)) |
| usage = {'input_tokens': sum(len(r['ids']) for r in items), 'output_tokens': 0} |
| collapsed = collapsed_options(qids, items) |
| if collapsed: |
| usage['options'] = collapsed |
| return {'model': 'xDecision', 'answers': answers, 'usage': usage} |
|
|
| def close(self): |
| with self._lock: |
| self.model = None |
|
|
