File size: 1,324 Bytes
03223d7 | 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 26 27 28 29 30 31 32 | """Run OpenJEV E4B 1.0 on a closed-choice question."""
import argparse
from pathlib import Path
from openjev import OpenJEV
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--state", default="Checkout fails after a deployment; customers cannot place orders.")
parser.add_argument("--instruction", default="Which team should own this incident?")
parser.add_argument("--option", dest="options", action="append", help="Repeat for each choice (at least two).")
parser.add_argument("--load-mode", choices=("nf4", "bf16"), default="nf4")
parser.add_argument("--device", default="cuda")
parser.add_argument("--repo-dir", type=Path, default=Path(__file__).resolve().parents[1])
args = parser.parse_args()
options = args.options or ["Billing", "Technical outages", "Sales"]
if len(options) < 2:
parser.error("pass at least two --option values")
model = OpenJEV.from_pretrained(args.repo_dir, device=args.device, load_mode=args.load_mode)
result = model.choice(state=args.state, instruction=args.instruction, options=options)
for option, probability in zip(options, result["probabilities"]):
print(f"{probability:.2%} {option}")
print(f"Selected: {result['selected_option']}")
if __name__ == "__main__":
main()
|