"""FP32 inference and candidate scoring using the included VimeML V2 source."""

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt", default="今日は雨が降っているので、")
    parser.add_argument("--candidate", action="append", help="Candidate string; repeat to rerank.")
    parser.add_argument("--max-new-tokens", type=int, default=8)
    parser.add_argument("--threads", type=int, default=4)
    args = parser.parse_args()
    if args.threads < 1 or args.max_new_tokens < 1:
        parser.error("threads and max-new-tokens must be positive.")
    sys.path.insert(0, str(ROOT / "source/src"))
    import torch
    from vimeml.deployment.v2 import BundleLM
    from vimeml.training.evaluate_ime import score_candidates

    torch.set_num_threads(args.threads)
    lm = BundleLM(ROOT / "inference")
    if args.candidate:
        result = score_candidates(lm, args.prompt, args.candidate)
        result["ranked_candidates"] = sorted(
            result["candidates"], key=lambda item: item["log_probability_sum"], reverse=True
        )
    else:
        result = lm.generate(args.prompt, max_new_tokens=args.max_new_tokens, temperature=0)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()