inference.py

"""Standalone OpenSML-150M inference on Apple Silicon using native MLX."""
import argparse,json
from pathlib import Path
import mlx.core as mx
from native_model import TransformerConfig,TransformerLM,count_parameters
from native_tokenizer import Tokenizer
from native_utils import file_sha256

def load_model(directory):
 root=Path(directory).resolve();cfg=json.loads((root/'config.json').read_text())
 if cfg['format']!='opensml-native-mlx-v1':raise ValueError('Unsupported model format')
 if file_sha256(root/'model.safetensors')!=cfg['weights_sha256']:raise ValueError('Weights checksum mismatch')
 tok=Tokenizer(root)
 if tok.vocab_size!=cfg['model']['vocab_size']:raise ValueError('Tokenizer/model vocabulary mismatch')
 model=TransformerLM(TransformerConfig(**cfg['model']));model.load_weights(str(root/'model.safetensors'),strict=True);model.eval();mx.eval(model.parameters())
 return model,tok

def generate(model,tokenizer,prompt,max_new_tokens=128,raw_completion=False):
 if max_new_tokens<1:raise ValueError('max_new_tokens must be positive')
 text=prompt if raw_completion else f'User: {prompt}\nAssistant:'
 ids=tokenizer.encode(text)
 if not ids or len(ids)+max_new_tokens>model.cfg.max_seq_len:raise ValueError('Prompt plus output budget must fit the 2048-token context')
 generated=[];caches=None;inputs=mx.array([ids]);stop='max_new_tokens'
 for _ in range(max_new_tokens):
  logits,caches=model.step(inputs,caches=caches)
  token=int(mx.argmax(logits,axis=-1).item());generated.append(token)
  if token==tokenizer.eos:stop='eos';break
  inputs=mx.array([[token]])
 answer=generated[:-1] if generated and generated[-1]==tokenizer.eos else generated
 return {'text':tokenizer.decode(answer),'token_ids':generated,'stop_reason':stop,'prompt_tokens':len(ids),'generated_tokens':len(generated)}

def main():
 p=argparse.ArgumentParser(description=__doc__);p.add_argument('--model-directory',type=Path,default=Path(__file__).resolve().parent);p.add_argument('--prompt',required=True);p.add_argument('--max-new-tokens',type=int,default=128);p.add_argument('--raw-completion',action='store_true',help='Skip the User/Assistant wrapper for base-style completion prompts');a=p.parse_args()
 model,tok=load_model(a.model_directory)
 print(json.dumps(generate(model,tok,a.prompt,a.max_new_tokens,a.raw_completion),ensure_ascii=False,indent=2))
if __name__=='__main__':main()