4B-5949/examples.py

"""Run against a local merged HF snapshot; no project-local dependencies."""

import argparse
import json

from openjet_runtime import OpenJet


def decision_examples():
    binary = {
        "id": "example-binary",
        "group_id": "example-binary",
        "primitive": "choice",
        "state": "Order 731 has been delivered. The customer's message says thank you.",
        "instructions": "Select the appropriate next workflow action.",
        "criteria": [
            {"id": "close", "description": "Close the resolved support ticket."},
            {"id": "refund", "description": "Refund an undelivered order."},
        ],
    }
    browser = {
        "id": "example-browser",
        "group_id": "example-browser",
        "primitive": "choice",
        "state": "A settings page has 16 visible buttons labeled Page 1 through Page 16.",
        "instructions": "Navigate to Page 12 by choosing its matching button.",
        "criteria": [
            {"id": f"click-{i}", "description": f"Click the Page {i} button."}
            for i in range(1, 17)
        ],
    }
    return binary, browser


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("model", help="Local merged snapshot directory")
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--dtype", choices=["float32", "bfloat16"], default="bfloat16")
    parser.add_argument("--effort", choices=["low", "high", "both"], default="both")
    parser.add_argument(
        "--text", action="store_true", help="Also run slow TYPE reference"
    )
    args = parser.parse_args()
    runtime = OpenJet.from_pretrained(args.model, args.device, args.dtype)
    efforts = ("low", "high") if args.effort == "both" else (args.effort,)
    for effort in efforts:
        for request in decision_examples():
            result = runtime.decide(request, effort)
            print(json.dumps({"example": request["id"], **result}, ensure_ascii=False))
        if args.text:
            result = runtime.generate_text(
                "Return only the literal text to type into a search box for 'red shoes'.",
                effort=effort,
                max_new_tokens=32,
            )
            print(json.dumps({"example": "browser-type", **result}, ensure_ascii=False))


if __name__ == "__main__":
    main()