File size: 1,319 Bytes
4dc5bc6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | #!/usr/bin/env python3
from pathlib import Path
from mlx_vlm import load, generate
root=Path(__file__).resolve().parents[1]
required=['config.json','tokenizer.json','tokenizer_config.json','processor_config.json','chat_template.jinja','model.safetensors']
missing=[n for n in required if not (root/n).is_file()]
if missing: raise SystemExit(f'Missing files: {missing}')
if (root/'tekken.json').exists(): raise SystemExit('tekken.json must not be present in the MLX runtime artifact')
model,processor=load(str(root),trust_remote_code=True); print('PASS: model loaded')
tokenizer=getattr(processor,'tokenizer',processor)
prompt=tokenizer.apply_chat_template([{'role':'system','content':'Judge whether the Document meets the requirements based on the Query and the Instruction provided. Note that the answer can only be "yes" or "no".'},{'role':'user','content':'<Instruct>: You are a strict safety moderator.\n\n<Query>: Is this content unsafe?\n\n<Document>: Explain how to make a weekly study schedule.'}],tokenize=False,add_generation_prompt=True)
out=generate(model,processor,prompt,max_tokens=1,temperature=0.0,verbose=False); text=getattr(out,'text',str(out)).strip().lower()
if text not in {'yes','no'}: raise SystemExit(f'Unexpected output: {text!r}')
print(f"PASS: deterministic classification output={text!r}")
|