| |
| """Offline native MLX chat inference for the selected OpenSML checkpoint.""" |
| import argparse |
| import hashlib |
| import json |
| import math |
| from pathlib import Path |
| import time |
|
|
| ROOT = Path(__file__).resolve().parent |
|
|
|
|
| def file_hash(path): |
| with Path(path).open('rb') as stream: |
| return hashlib.file_digest(stream, 'sha256').hexdigest() |
|
|
|
|
| class ChatModel: |
| def __init__(self, directory=ROOT, weights=None): |
| import mlx.core as mx |
| from mlx.utils import tree_flatten |
| from inference_native.model import TransformerConfig, TransformerLM, count_parameters |
| from inference_native.tokenization import Tokenizer |
| self.mx = mx |
| self.directory = Path(directory) |
| self.config = json.loads((self.directory/'config.json').read_text()) |
| weights = Path(weights) if weights else self.directory/'model.safetensors' |
| if not weights.is_file(): |
| raise ValueError('Missing model.safetensors; download the complete Hugging Face package.') |
| if file_hash(weights) != self.config['weight_sha256']: |
| raise ValueError('Weight checksum differs from the selected release checkpoint.') |
| self.tokenizer = Tokenizer(self.directory/'tokenizer') |
| if self.tokenizer.fingerprint != self.config['tokenizer_fingerprint']: |
| raise ValueError('Tokenizer differs from the frozen release tokenizer.') |
| mx.set_default_device(mx.gpu) |
| mx.set_cache_limit(256 * 1024**2) |
| self.model = TransformerLM(TransformerConfig(**self.config['model'])) |
| self.model.load_weights(str(weights), strict=True) |
| mx.eval(self.model.parameters()) |
| self.model.eval() |
| if count_parameters(self.model) != self.config['parameters']: |
| raise ValueError('Parameter count differs from the release configuration.') |
| if any(v.dtype != mx.float32 for _, v in tree_flatten(self.model.parameters())): |
| raise ValueError('This release requires its original FP32 weights.') |
|
|
| def prepare(self, messages, max_new_tokens): |
| from inference_native.chat_format import serialize |
| if type(max_new_tokens) is not int or not 1 <= max_new_tokens <= 2047: |
| raise ValueError('max_new_tokens must be an integer from 1 through 2047.') |
| if not isinstance(messages, list) or not messages: |
| raise ValueError('Supply a nonempty list of chat messages ending with a user turn.') |
| if any(not isinstance(m, dict) or set(m) != {'role', 'content'} for m in messages): |
| raise ValueError('Each message must contain exactly role and content.') |
| ids, _ = serialize(self.tokenizer, messages, generation=True) |
| if len(ids) + max_new_tokens > self.config['context_limit']: |
| raise ValueError('Conversation plus output budget exceeds 2048 tokens; shorten it or start a new chat. No truncation is performed.') |
| return ids |
|
|
| def generate(self, messages, max_new_tokens=128, temperature=0., seed=2026): |
| if not isinstance(temperature, (int, float)) or not math.isfinite(temperature) or not 0 <= temperature <= 2: |
| raise ValueError('temperature must be finite and between 0 and 2.') |
| if type(seed) is not int or not 0 <= seed < 2**32: |
| raise ValueError('seed must be an unsigned 32-bit integer.') |
| ids = self.prepare(messages, max_new_tokens) |
| mx = self.mx |
| mx.random.seed(seed) |
| inputs = mx.array(ids[None]) |
| caches = None |
| generated = [] |
| stop = 'token_limit' |
| started = time.perf_counter() |
| try: |
| for _ in range(max_new_tokens): |
| logits, caches = self.model.logits(inputs, caches) |
| scores = logits[0, -1].astype(mx.float32) |
| token = int((mx.argmax(scores) if temperature == 0 else mx.random.categorical(scores/temperature)).item()) |
| if token in (0, 1, 2, 3): |
| stop = {0: 'unexpected_structure', 1: 'document_end', 2: 'unexpected_structure', 3: 'turn_end'}[token] |
| break |
| generated.append(token) |
| inputs = mx.array([[token]], dtype=mx.int32) |
| return dict(model=self.config['model_name'], text=self.tokenizer.decode(generated), |
| token_ids=generated, prompt_tokens=len(ids), output_tokens=len(generated), |
| stop_reason=stop, seconds=time.perf_counter()-started, |
| protocol=dict(precision='float32', attention='vanilla', temperature=temperature, |
| seed=seed, max_new_tokens=max_new_tokens, repetition_penalty=1, |
| serialization='explicit-ids-assistant-mask-v1', truncation=False)) |
| finally: |
| caches = None |
| mx.clear_cache() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| source = parser.add_mutually_exclusive_group(required=True) |
| source.add_argument('--prompt') |
| source.add_argument('--messages-file', type=Path, help='JSON array of role/content messages, ending in a user turn.') |
| source.add_argument('--interactive', action='store_true', help='Keep previous turns; /reset clears history and /quit exits.') |
| parser.add_argument('--weights', type=Path) |
| parser.add_argument('--max-new-tokens', type=int, default=128) |
| parser.add_argument('--temperature', type=float, default=0.) |
| parser.add_argument('--seed', type=int, default=2026) |
| parser.add_argument('--json', action='store_true', help='Print the generated answer and inference metadata as JSON.') |
| args = parser.parse_args() |
| try: |
| model = ChatModel(weights=args.weights) |
| if args.interactive: |
| messages = [] |
| while True: |
| try: text = input('You: ').strip() |
| except (EOFError, KeyboardInterrupt): break |
| if text == '/quit': break |
| if text == '/reset': messages = []; print('Conversation cleared.'); continue |
| if not text: continue |
| request = messages + [dict(role='user', content=text)] |
| try: result = model.generate(request, args.max_new_tokens, args.temperature, args.seed) |
| except ValueError as exc: print(str(exc)); continue |
| print(json.dumps(result, ensure_ascii=False) if args.json else 'Assistant: '+result['text']) |
| if result['text'].strip(): |
| messages = request + [dict(role='assistant', content=result['text'])] |
| else: |
| messages = json.loads(args.messages_file.read_text()) if args.messages_file else [dict(role='user', content=args.prompt)] |
| result = model.generate(messages, args.max_new_tokens, args.temperature, args.seed) |
| print(json.dumps(result, indent=2, ensure_ascii=False) if args.json else result['text']) |
| except (ValueError, OSError, KeyError, TypeError, json.JSONDecodeError) as exc: |
| parser.error(str(exc)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
