File size: 1,102 Bytes
bfd78b5 | 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 | """Run an independent synthetic structured-choice example."""
import argparse
import json
from pathlib import Path
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--backend', choices=('torch', 'mlx'), default='torch')
parser.add_argument('--model', type=Path, default=Path(__file__).resolve().parent)
parser.add_argument('--request', type=Path)
parser.add_argument('--device')
parser.add_argument('--output', type=Path)
args = parser.parse_args()
request = json.loads((args.request or args.model / 'example_request.json').read_text(encoding='utf-8'))
if args.backend == 'mlx':
from laya_mlx.agent import Agent
agent = Agent(str(args.model), dtype='float32', device=args.device)
else:
import torch
from laya.agent import Agent
torch.set_num_threads(4)
agent = Agent(str(args.model), device=args.device or 'cuda', compile=False)
result = agent.predict(request['state'], request['questions'])
serialized = json.dumps(result, ensure_ascii=False, indent=2) + '\n'
if args.output:
args.output.write_text(serialized, encoding='utf-8')
print(serialized)
|