"""Download Reflex-1 and exercise single-request and batch inference."""
import argparse
import json
from dataclasses import asdict

import torch
from transformers import AutoModel


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="gai-labs/reflex-1")
    parser.add_argument("--revision", help="Explicit Hub checkpoint branch, tag or commit")
    parser.add_argument("--local-files-only", action="store_true")
    args = parser.parse_args()
    torch.set_num_threads(1)
    model = AutoModel.from_pretrained(args.model, trust_remote_code=True,
                                     revision=args.revision, local_files_only=args.local_files_only)
    requests = [
        {"state": "The lamp is red.", "question": "Which color is the lamp?",
         "options": ["red", "blue", "green"]},
        {"state": "The customer was charged twice for one card payment.",
         "question": "Choose the matching issue.",
         "options": ["duplicate charge", "lost card", "unknown fee", "cash withdrawal"]},
    ]
    single = model.predict(requests[0])[0]
    batch = model.predict_batch(requests)
    assert single.choice == "red"
    assert [decisions[0].choice for decisions in batch] == ["red", "duplicate charge"]
    torch.testing.assert_close(torch.tensor(single.probabilities),
                               torch.tensor(batch[0][0].probabilities), atol=2e-5, rtol=1e-4)
    print(json.dumps({"passed": True, "revision": getattr(model.config, "_commit_hash", None),
                      "decisions": [asdict(row[0]) for row in batch]}, indent=2))


if __name__ == "__main__":
    main()