import argparse
import json
from pathlib import Path

from mlx_lm import generate, load
from mlx_lm.sample_utils import make_sampler

MODEL = "mlx-community/Qwen2.5-0.5B-Instruct-4bit"

parser = argparse.ArgumentParser(description="Score a model on held-out library requests.")
parser.add_argument("--adapter", type=Path)
args = parser.parse_args()

if args.adapter and not (args.adapter / "adapters.safetensors").is_file():
    parser.error(f"No adapters.safetensors in {args.adapter}")

examples = [
    json.loads(line)["messages"]
    for line in Path("library_data/test.jsonl").read_text(encoding="utf-8").splitlines()
    if line.strip()
]
if not examples:
    raise SystemExit("The held-out test set is empty.")
model, tokenizer = load(MODEL, adapter_path=str(args.adapter) if args.adapter else None)
greedy = make_sampler(temp=0)

correct = 0
for messages in examples:
    expected = messages[-1]["content"]
    prompt = tokenizer.apply_chat_template(
        messages[:-1], tokenize=False, add_generation_prompt=True
    )
    answer = generate(
        model, tokenizer, prompt=prompt, max_tokens=12, sampler=greedy, verbose=False
    ).strip()
    passed = answer == expected
    correct += passed
    print(f"{'PASS' if passed else 'MISS'} | {expected} | {answer!r}")

print(f"{correct}/{len(examples)} exact matches")