from __future__ import annotations

import argparse
from pathlib import Path

from model import Add2ModelError, TinyAdd2LLM


def main() -> None:
    parser = argparse.ArgumentParser(description="Tiny Add2 LLM inference")
    parser.add_argument("prompt", help="Example: '12 + 30'")
    parser.add_argument("--model", default=str(Path("model") / "add2_model.bin"), help="Path to model file")
    args = parser.parse_args()

    try:
        model = TinyAdd2LLM.load(args.model)
        print(model.answer(args.prompt))
    except (Add2ModelError, FileNotFoundError) as e:
        print(f"ERROR: {e}")
        raise SystemExit(1)


if __name__ == "__main__":
    main()