leo_ui7m/cli.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

from leo_ui7m import UIActionModel


def main() -> int:
    parser = argparse.ArgumentParser(description="LEO-UI7M local UI-action prediction CLI")
    parser.add_argument("--goal", required=True, help="Task goal")
    parser.add_argument("--elements", required=True, help="Path to JSON file containing a list of UI elements")
    parser.add_argument("--history", default=None, help="Optional path to JSON file containing action history")
    parser.add_argument("--checkpoint", default=None, help="Optional checkpoint path")
    parser.add_argument("--step-index", type=int, default=0)
    parser.add_argument("--pretty", action="store_true")
    args = parser.parse_args()

    elements = json.loads(Path(args.elements).read_text(encoding="utf-8"))
    history = json.loads(Path(args.history).read_text(encoding="utf-8")) if args.history else []

    model = UIActionModel(checkpoint_path=args.checkpoint)
    action = model.predict_action(
        goal=args.goal,
        elements=elements,
        history=history,
        step_index=args.step_index,
    )

    out = action.to_dict()
    print(json.dumps(out, indent=2 if args.pretty else None, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())