"""Answer one System One request with a D1A Core ML bundle (this repository, downloaded), on the Neural Engine.

    pip install coremltools numpy torch transformers safetensors "d1a @ git+https://github.com/jonpol01/d1a"
    python reference/run.py <bundle dir> request.json

request.json is a /v1/systemone request: {"state": "...", "questions": {"id": {"type": "noul" | "choice" | "score", ...}}}.
The request must encode to 512 tokens or fewer. The first run compiles the chunks into coreml/.compiled/ (a few minutes);
later runs load them. Prints each question's options with their probabilities."""
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))   # workers.py and d1a_coreml.py sit next to this file


def main():
    from transformers import AutoTokenizer

    from d1a.core.api import SystemOneRequest, to_record
    from d1a.core.encoding import encode
    from workers import WorkerPipeline

    bundle, request = sys.argv[1], sys.argv[2]
    req = SystemOneRequest.model_validate({"model": "d1a-latest", **json.load(open(request))})
    rec, meta = to_record(req)
    enc = encode(AutoTokenizer.from_pretrained(bundle), rec)
    if len(enc["ids"]) > 512:
        raise SystemExit(f"the request encodes to {len(enc['ids'])} tokens; this build takes 512 or fewer")
    pipe = WorkerPipeline(bundle, os.path.join(bundle, "coreml"), 512)
    try:
        for m, p in zip(meta, pipe.probs(enc)):
            print(m["id"], {k: round(float(v), 4) for k, v in zip(m["keys"], p)})
    finally:
        pipe.close()


if __name__ == "__main__":
    main()