File size: 669 Bytes
beb6a05 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | 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()
|